scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix: tests.test_litellm.interactions.test_gemini_interactions.TestGeminiInteractions

Browse source
This commit is contained in:
Sameer Kankute 2026-01-05 16:41:24 +05:30
parent d590c820c2
commit 840617de73
1 changed files with 23 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/llms/gemini/interactions/transformation.py
View file

@ -139,7 +139,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
) -> InteractionsAPIResponse:
"""Parse response - it already matches our response type."""
"""Parse response and transform to OpenAI-compatible format."""
try:
logging_obj.post_call(
original_response=raw_response.text,
@ -155,6 +155,18 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
verbose_logger.debug("Google AI Interactions response: %s", raw_json)
# Transform usage to OpenAI format
if "usage" in raw_json and raw_json["usage"]:
usage = raw_json["usage"]
# Map Google's field names to OpenAI's field names
if "total_input_tokens" in usage:
usage["input_tokens"] = usage["total_input_tokens"]
if "total_output_tokens" in usage:
usage["output_tokens"] = usage["total_output_tokens"]
if "total_tokens" in usage:
# Keep total_tokens as is - it's the same in both formats
pass
response = InteractionsAPIResponse(**raw_json)
response._hidden_params["headers"] = dict(raw_response.headers)
response._hidden_params["additional_headers"] = process_response_headers(dict(raw_response.headers))
@ -200,6 +212,16 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig):
status_code=raw_response.status_code,
headers=dict(raw_response.headers),
)
# Transform usage to OpenAI format
if "usage" in raw_json and raw_json["usage"]:
usage = raw_json["usage"]
# Map Google's field names to OpenAI's field names
if "total_input_tokens" in usage:
usage["input_tokens"] = usage["total_input_tokens"]
if "total_output_tokens" in usage:
usage["output_tokens"] = usage["total_output_tokens"]
response = InteractionsAPIResponse(**raw_json)
response._hidden_params["headers"] = dict(raw_response.headers)
return response