scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

_optional_combine_thinking_block_in_choices

Browse source
This commit is contained in:
Ishaan Jaff 2025-03-06 00:18:38 -08:00
parent ea2c805d93
commit bbee3b7ecd
1 changed files with 51 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/litellm_core_utils/streaming_handler.py
View file

@ -70,6 +70,11 @@ class CustomStreamWrapper:
self.completion_stream = completion_stream
self.sent_first_chunk = False
self.sent_last_chunk = False
self.sent_first_thinking_block = False
self.sent_last_thinking_block = False
self.thinking_content = ""
self.system_fingerprint: Optional[str] = None
self.received_finish_reason: Optional[str] = None
self.intermittent_finish_reason: Optional[str] = (
@ -873,6 +878,10 @@ class CustomStreamWrapper:
_index: Optional[int] = completion_obj.get("index")
if _index is not None:
model_response.choices[0].index = _index
model_response = self._optional_combine_thinking_block_in_choices( # type: ignore
model_response=model_response
)
print_verbose(f"returning model_response: {model_response}")
return model_response
else:
@ -929,6 +938,48 @@ class CustomStreamWrapper:
self.chunks.append(model_response)
return
def _optional_combine_thinking_block_in_choices(
self, model_response: ModelResponseStream
) -> Optional[ModelResponseStream]:
"""
UI's Like OpenWebUI expect to get 1 chunk with <thinking>...</thinking> tags in the chunk content
This helper handles that logic
Collects all `<thinking>...</thinking>` blocks and combines them into 1 chunk
"""
if litellm.merge_reasoning_content_in_choices:
reasoning_content = getattr(
model_response.choices[0].delta, "reasoning_content", None
)
if reasoning_content:
if self.sent_first_thinking_block is False:
self.thinking_content += (
(model_response.choices[0].delta.content or "")
+ f"<thinking>{model_response.choices[0].delta.reasoning_content}"
)
self.sent_first_thinking_block = True
return None
elif (
self.sent_first_thinking_block is True
and hasattr(model_response.choices[0].delta, "reasoning_content")
and model_response.choices[0].delta.reasoning_content
):
self.thinking_content += (
model_response.choices[0].delta.content or ""
) + f"{model_response.choices[0].delta.reasoning_content}"
return None
elif (
self.sent_first_thinking_block is True
and not self.sent_last_thinking_block
and model_response.choices[0].delta.content
):
self.thinking_content += "</thinking>"
model_response.choices[0].delta.content = self.thinking_content
self.sent_last_thinking_block = True
return model_response
def chunk_creator(self, chunk: Any): # type: ignore # noqa: PLR0915
model_response = self.model_response_creator()
response_obj: Dict[str, Any] = {}