From 840617de73b95a9ab5d22a54f92d7cb6de8640fe Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 5 Jan 2026 16:41:24 +0530 Subject: [PATCH] fix: tests.test_litellm.interactions.test_gemini_interactions.TestGeminiInteractions --- .../gemini/interactions/transformation.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index d21775eb236..769d70473b9 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -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