scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Make gemini stream thinking as reasoning_content (#11290)

Browse source 
When "Thought": True, return text as reasoning_content instead of
content.

fixes #10563
fixes #11000
This commit is contained in:
Adam Holmberg 2025-05-31 11:13:00 -05:00 • committed by GitHub
parent 51f716c762
commit e0daa3da68
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 88 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

76
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -43,7 +43,6 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParamFunctionChunk,
ChatCompletionUsageBlock,
OpenAIChatCompletionFinishReason,
)
from litellm.types.llms.vertex_ai import (
@ -64,8 +63,10 @@ from litellm.types.utils import (
ChatCompletionTokenLogprob,
ChoiceLogprobs,
CompletionTokensDetailsWrapper,
GenericStreamingChunk,
Delta,
ModelResponseStream,
PromptTokensDetailsWrapper,
StreamingChoices,
TopLogprob,
Usage,
)
@ -1650,14 +1651,15 @@ class ModelResponseIterator:
self.accumulated_json = ""
self.sent_first_chunk = False
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
def chunk_parser(self, chunk: dict) -> ModelResponseStream:
try:
processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore
text = ""
reasoning_content = None
tool_use: Optional[ChatCompletionToolCallChunk] = None
finish_reason = ""
usage: Optional[ChatCompletionUsageBlock] = None
usage: Optional[Usage] = None
_candidates: Optional[List[Candidates]] = processed_chunk.get("candidates")
gemini_chunk: Optional[Candidates] = None
if _candidates and len(_candidates) > 0:
@ -1669,7 +1671,11 @@ class ModelResponseIterator:
and "parts" in gemini_chunk["content"]
):
if "text" in gemini_chunk["content"]["parts"][0]:
text = gemini_chunk["content"]["parts"][0]["text"]
if gemini_chunk["content"]["parts"][0].get("thought"):
reasoning_content = gemini_chunk["content"]["parts"][0]["text"]
else:
text = gemini_chunk["content"]["parts"][0]["text"]
elif "functionCall" in gemini_chunk["content"]["parts"][0]:
function_call = ChatCompletionToolCallFunctionChunk(
name=gemini_chunk["content"]["parts"][0]["functionCall"][
@ -1695,7 +1701,7 @@ class ModelResponseIterator:
## GEMINI SETS FINISHREASON ON EVERY CHUNK!
if "usageMetadata" in processed_chunk:
usage = ChatCompletionUsageBlock(
usage = Usage(
prompt_tokens=processed_chunk["usageMetadata"].get(
"promptTokenCount", 0
),
@ -1705,20 +1711,26 @@ class ModelResponseIterator:
total_tokens=processed_chunk["usageMetadata"].get(
"totalTokenCount", 0
),
completion_tokens_details={
"reasoning_tokens": processed_chunk["usageMetadata"].get(
completion_tokens_details=CompletionTokensDetailsWrapper(
reasoning_tokens=processed_chunk["usageMetadata"].get(
"thoughtsTokenCount", 0
)
},
),
)
returned_chunk = GenericStreamingChunk(
text=text,
tool_use=tool_use,
is_finished=False,
finish_reason=finish_reason,
returned_chunk = ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(
content=text,
tool_calls=[tool_use] if tool_use is not None else None,
reasoning_content=reasoning_content,
),
finish_reason=finish_reason,
)
],
usage=usage,
index=0,
)
return returned_chunk
except json.JSONDecodeError:
@ -1729,7 +1741,7 @@ class ModelResponseIterator:
self.response_iterator = self.streaming_response
return self
def handle_valid_json_chunk(self, chunk: str) -> GenericStreamingChunk:
def handle_valid_json_chunk(self, chunk: str) -> ModelResponseStream:
chunk = chunk.strip()
try:
json_chunk = json.loads(chunk)
@ -1747,7 +1759,7 @@ class ModelResponseIterator:
return self.chunk_parser(chunk=json_chunk)
def handle_accumulated_json_chunk(self, chunk: str) -> GenericStreamingChunk:
def handle_accumulated_json_chunk(self, chunk: str) -> ModelResponseStream:
chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
message = chunk.replace("\n\n", "")
@ -1761,16 +1773,18 @@ class ModelResponseIterator:
return self.chunk_parser(chunk=_data)
except json.JSONDecodeError:
# If it's not valid JSON yet, continue to the next event
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason="",
)
],
usage=None,
index=0,
tool_use=None,
)
def _common_chunk_parsing_logic(self, chunk: str) -> GenericStreamingChunk:
def _common_chunk_parsing_logic(self, chunk: str) -> ModelResponseStream:
try:
chunk = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
if len(chunk) > 0:
@ -1784,13 +1798,15 @@ class ModelResponseIterator:
elif self.chunk_type == "accumulated_json":
return self.handle_accumulated_json_chunk(chunk=chunk)
return GenericStreamingChunk(
text="",
is_finished=False,
finish_reason="",
return ModelResponseStream(
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason="",
)
],
usage=None,
index=0,
tool_use=None,
)
except Exception:
raise

48
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -334,15 +334,52 @@ def test_streaming_chunk_includes_reasoning_tokens():
}
iterator = ModelResponseIterator(streaming_response=[], sync_stream=True)
streaming_chunk = iterator.chunk_parser(chunk)
assert streaming_chunk["usage"] is not None
assert streaming_chunk["usage"]["prompt_tokens"] == 5
assert streaming_chunk["usage"]["completion_tokens"] == 7
assert streaming_chunk["usage"]["total_tokens"] == 12
assert streaming_chunk.usage is not None
assert streaming_chunk.usage.prompt_tokens == 5
assert streaming_chunk.usage.completion_tokens == 7
assert streaming_chunk.usage.total_tokens == 12
assert (
streaming_chunk["usage"]["completion_tokens_details"]["reasoning_tokens"] == 3
streaming_chunk.usage.completion_tokens_details.reasoning_tokens == 3
)
def test_streaming_chunk_includes_reasoning_content():
"""
Ensure that when Gemini returns a chunk with `thought=True`, the parser maps it to `reasoning_content`.
"""
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator,
)
# Simulate a streaming chunk from Gemini which contains reasoning (thought) content
chunk = {
"candidates": [
{
"content": {
"parts": [
{
"text": "I'm thinking through the problem...",
"thought": True,
}
]
}
}
],
"usageMetadata": {},
}
iterator = ModelResponseIterator(streaming_response=[], sync_stream=True)
streaming_chunk = iterator.chunk_parser(chunk)
# The text content should be empty and reasoning_content should be populated
assert streaming_chunk.choices[0].delta.content == ""
assert (
streaming_chunk.choices[0].delta.reasoning_content
== "I'm thinking through the problem..."
)
def test_check_finish_reason():
config = VertexGeminiConfig()
finish_reason_mappings = config.get_finish_reason_mapping()
@ -446,4 +483,3 @@ def test_vertex_ai_map_tool_with_anyof():
] == {
"anyOf": [{"type": "string", "nullable": True, "title": "Base Branch"}]
}, f"Expected only anyOf field and its contents to be kept, but got {tools[0]['function_declarations'][0]['parameters']['properties']['base_branch']}"