tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-14 22:21:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(vertex_and_google_ai_studio_gemini.py): bubble up thoughtsignature back to client

Browse Source
This commit is contained in:
Krrish Dholakia
2025-08-30 17:26:18 -07:00
parent 599071321d
commit 51c73dc60b
4 changed files with 137 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+55 -22
View File
@@ -43,6 +43,7 @@ from litellm.types.llms.gemini import BidiGenerateContentServerMessage
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionResponseMessage,
ChatCompletionThinkingBlock,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParamFunctionChunk,
@@ -792,7 +793,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
content_str += _content_str
return content_str, reasoning_content_str
def _extract_thinking_blocks_from_parts(
self, parts: List[HttpxPartType]
) -> List[ChatCompletionThinkingBlock]:
"""Extract thinking blocks from parts if present"""
thinking_blocks: List[ChatCompletionThinkingBlock] = []
for part in parts:
if "thoughtSignature" in part:
part_copy = part.copy()
part_copy.pop("thoughtSignature")
thinking_blocks.append(
ChatCompletionThinkingBlock(
type="thinking",
thinking=json.dumps(part_copy),
signature=part["thoughtSignature"],
)
)
return thinking_blocks
def _extract_image_response_from_parts(
self, parts: List[HttpxPartType]
) -> Optional[ImageURLObject]:
@@ -804,10 +823,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if mime_type.startswith("image/"):
# Convert base64 data to data URI format
data_uri = f"data:{mime_type};base64,{data}"
return ImageURLObject(
url=data_uri,
detail="auto"
)
return ImageURLObject(url=data_uri, detail="auto")
return None
def _extract_audio_response_from_parts(
@@ -1127,7 +1143,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
elif web_search_queries:
web_search_requests = len(grounding_metadata)
return web_search_requests
@staticmethod
def _create_streaming_choice(
chat_completion_message: ChatCompletionResponseMessage,
@@ -1151,9 +1167,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
index=candidate.get("index", idx),
delta=Delta(
content=chat_completion_message.get("content"),
reasoning_content=chat_completion_message.get(
"reasoning_content"
),
reasoning_content=chat_completion_message.get("reasoning_content"),
tool_calls=tools,
image=image_response,
function_call=functions,
@@ -1164,13 +1178,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return choice
@staticmethod
def _extract_candidate_metadata(candidate: Candidates) -> Tuple[List[dict], List[dict], List, List]:
def _extract_candidate_metadata(
candidate: Candidates,
) -> Tuple[List[dict], List[dict], List, List]:
"""
Extract metadata from a single candidate response.
Returns:
grounding_metadata: List[dict]
url_context_metadata: List[dict]
url_context_metadata: List[dict]
safety_ratings: List
citation_metadata: List
"""
@@ -1178,7 +1194,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
url_context_metadata: List[dict] = []
safety_ratings: List = []
citation_metadata: List = []
if "groundingMetadata" in candidate:
if isinstance(candidate["groundingMetadata"], list):
grounding_metadata.extend(candidate["groundingMetadata"]) # type: ignore
@@ -1194,8 +1210,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if "urlContextMetadata" in candidate:
# Add URL context metadata to grounding metadata
url_context_metadata.append(cast(dict, candidate["urlContextMetadata"]))
return grounding_metadata, url_context_metadata, safety_ratings, citation_metadata
return (
grounding_metadata,
url_context_metadata,
safety_ratings,
citation_metadata,
)
@staticmethod
def _process_candidates(
@@ -1227,6 +1248,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
tools: Optional[List[ChatCompletionToolCallChunk]] = []
functions: Optional[ChatCompletionToolCallFunctionChunk] = None
cumulative_tool_call_index: int = 0
thinking_blocks: Optional[List[ChatCompletionThinkingBlock]] = None
for idx, candidate in enumerate(_candidates):
if "content" not in candidate:
@@ -1239,7 +1261,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
candidate_safety_ratings,
candidate_citation_metadata,
) = VertexGeminiConfig._extract_candidate_metadata(candidate)
grounding_metadata.extend(candidate_grounding_metadata)
url_context_metadata.extend(candidate_url_context_metadata)
safety_ratings.extend(candidate_safety_ratings)
@@ -1264,6 +1286,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
)
)
thinking_blocks = (
VertexGeminiConfig()._extract_thinking_blocks_from_parts(
parts=candidate["content"]["parts"]
)
)
if audio_response is not None:
cast(Dict[str, Any], chat_completion_message)[
"audio"
@@ -1271,7 +1299,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
chat_completion_message["content"] = None # OpenAI spec
if image_response is not None:
# Handle image response - combine with text content into structured format
cast(Dict[str, Any], chat_completion_message)["image"] = image_response
cast(Dict[str, Any], chat_completion_message)[
"image"
] = image_response
if content is not None:
chat_completion_message["content"] = content
@@ -1298,15 +1328,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if functions is not None:
chat_completion_message["function_call"] = functions
if thinking_blocks is not None:
chat_completion_message["thinking_blocks"] = thinking_blocks # type: ignore
if isinstance(model_response, ModelResponseStream):
choice = VertexGeminiConfig._create_streaming_choice(
chat_completion_message=chat_completion_message,
candidate=candidate,
idx=idx,
tools=tools,
functions=functions,
candidate=candidate,
idx=idx,
tools=tools,
functions=functions,
chat_completion_logprobs=chat_completion_logprobs,
image_response=image_response
image_response=image_response,
)
model_response.choices.append(choice)
elif isinstance(model_response, ModelResponse):
litellm/types/llms/openai.py
+6 -2
View File
@@ -43,10 +43,14 @@ from openai.types.responses.response import (
# Handle OpenAI SDK version compatibility for Text type
try:
from openai.types.responses.response_create_params import Text as ResponseText
from openai.types.responses.response_create_params import (
Text as ResponseText, # type: ignore
)
except (ImportError, AttributeError):
# Fall back to the concrete config type available in all SDK versions
from openai.types.responses.response_text_config_param import ResponseTextConfigParam as ResponseText
from openai.types.responses.response_text_config_param import (
ResponseTextConfigParam as ResponseText,
)
from openai.types.responses.response_create_params import (
Reasoning,
litellm/types/llms/vertex_ai.py
+3 -1
View File
@@ -72,6 +72,7 @@ class HttpxPartType(TypedDict, total=False):
executableCode: HttpxExecutableCode
codeExecutionResult: HttpxCodeExecutionResult
thought: bool
thoughtSignature: str
class HttpxContentType(TypedDict, total=False):
@@ -245,10 +246,11 @@ class UsageMetadata(TypedDict, total=False):
class TokenCountDetailsResponse(TypedDict):
"""
Response structure for token count details with modality breakdown.
Example:
{'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
"""
totalTokens: int
promptTokensDetails: List[PromptTokensDetails]
tests/llm_translation/test_gemini.py
+73 -23
View File
@@ -436,7 +436,10 @@ def test_gemini_with_empty_function_call_arguments():
async def test_claude_tool_use_with_gemini():
response = await litellm.anthropic.messages.acreate(
messages=[
{"role": "user", "content": "Hello, can you tell me the weather in Boston. Please respond with a tool call?"}
{
"role": "user",
"content": "Hello, can you tell me the weather in Boston. Please respond with a tool call?",
}
],
model="gemini/gemini-2.5-flash",
stream=True,
@@ -578,11 +581,17 @@ def test_gemini_tool_use():
assert stop_reason is not None
assert stop_reason == "tool_calls"
@pytest.mark.asyncio
async def test_gemini_image_generation_async():
litellm._turn_on_debug()
response = await litellm.acompletion(
messages=[{"role": "user", "content": "Generate an image of a banana wearing a costume that says LiteLLM"}],
messages=[
{
"role": "user",
"content": "Generate an image of a banana wearing a costume that says LiteLLM",
}
],
model="gemini/gemini-2.5-flash-image-preview",
)
@@ -597,12 +606,16 @@ async def test_gemini_image_generation_async():
assert IMAGE_URL["url"].startswith("data:image/png;base64,")
@pytest.mark.asyncio
async def test_gemini_image_generation_async_stream():
#litellm._turn_on_debug()
# litellm._turn_on_debug()
response = await litellm.acompletion(
messages=[{"role": "user", "content": "Generate an image of a banana wearing a costume that says LiteLLM"}],
messages=[
{
"role": "user",
"content": "Generate an image of a banana wearing a costume that says LiteLLM",
}
],
model="gemini/gemini-2.5-flash-image-preview",
stream=True,
)
@@ -611,35 +624,72 @@ async def test_gemini_image_generation_async_stream():
model_response_image = None
async for chunk in response:
print("CHUNK: ", chunk)
if hasattr(chunk.choices[0].delta, "image") and chunk.choices[0].delta.image is not None:
if (
hasattr(chunk.choices[0].delta, "image")
and chunk.choices[0].delta.image is not None
):
model_response_image = chunk.choices[0].delta.image
print("MODEL_RESPONSE_IMAGE: ", model_response_image)
assert model_response_image is not None
assert model_response_image["url"].startswith("data:image/png;base64,")
break
#########################################################
# Important: Validate we did get an image in the response
#########################################################
assert model_response_image is not None
assert model_response_image["url"].startswith("data:image/png;base64,")
def test_system_message_with_no_user_message():
"""
Test that the system message is translated correctly for non-OpenAI providers.
"""
messages = [
{
"role": "system",
"content": "Be a good bot!",
"""
Test that the system message is translated correctly for non-OpenAI providers.
"""
messages = [
{
"role": "system",
"content": "Be a good bot!",
},
]
response = litellm.completion(
model="gemini/gemini-2.5-flash",
messages=messages,
)
assert response is not None
assert response.choices[0].message.content is not None
def test_gemini_with_thinking():
from litellm import completion
litellm._turn_on_debug()
tools = [
{
"type": "function",
"function": {
"name": "get_current_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
]
}
]
messages = [{"role": "user", "content": "What's the weather like in Boston today?"}]
response = litellm.completion(
model="gemini/gemini-2.5-flash",
messages=messages,
)
assert response is not None
assert response.choices[0].message.content is not None
result = completion(
model="gemini/gemini-2.5-flash",
messages=messages,
tools=tools,
)
print(f"result: {result}")