From 3aea951e6cce99787cd8e7d1bdf73b4a5aa4d46c Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:36:43 -0700 Subject: [PATCH] refactor(xai): keep Responses tool usage pass-through in llms/xai Revert shared responses/utils.py extras forwarding. Attach server_side_tool_usage_details on chat Usage inside XAIResponsesAPIConfig so cost calc keeps web_search_calls without provider logic in shared utils. --- litellm/llms/xai/responses/transformation.py | 50 ++++++++++++++++++- litellm/responses/utils.py | 47 ----------------- .../responses/test_responses_utils.py | 17 ------- 3 files changed, 48 insertions(+), 66 deletions(-) diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 48fb95d9411..fbff9876d49 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,16 +1,22 @@ from typing import TYPE_CHECKING, Any, Final +import httpx + import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ( + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LlmProviders +from litellm.types.utils import LlmProviders, Usage if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -51,6 +57,46 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return supported_params + def transform_response_api_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> ResponsesAPIResponse: + """ + Attach xAI tool usage details onto a chat Usage object. + + Cost calculation normalizes Responses usage via + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage, which + drops non-standard fields unless usage is already a chat Usage instance. + """ + response = super().transform_response_api_response( + model=model, raw_response=raw_response, logging_obj=logging_obj + ) + self._attach_server_side_tool_usage_details_to_usage(response) + return response + + @staticmethod + def _attach_server_side_tool_usage_details_to_usage( + response: ResponsesAPIResponse, + ) -> None: + if response.usage is None: + return + + details = getattr(response.usage, "server_side_tool_usage_details", None) + if details is None and isinstance(response.usage, dict): + details = response.usage.get("server_side_tool_usage_details") + if details is None: + return + + if isinstance(response.usage, Usage): + setattr(response.usage, "server_side_tool_usage_details", details) + return + + chat_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) + setattr(chat_usage, "server_side_tool_usage_details", details) + response.usage = chat_usage # type: ignore[assignment] + def _transform_web_search_tool(self, tool: dict[str, Any]) -> XAIWebSearchTool | dict[str, Any]: """ Transform web_search tool to XAI format. diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 5e1412ee86b..db2e515609c 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1022,20 +1022,6 @@ class ResponsesAPIRequestUtils: class ResponseAPILoggingUtils: - # Standard Responses usage keys mapped explicitly below; extras pass through to Usage. - _RESPONSE_API_USAGE_MAPPED_KEYS = frozenset( - { - "input_tokens", - "output_tokens", - "total_tokens", - "input_tokens_details", - "output_tokens_details", - "input_token_details", - "output_token_details", - "cost", # handled separately after Usage construction - } - ) - @staticmethod def _is_response_api_usage(usage: dict | ResponseAPIUsage) -> bool: """returns True if usage is from OpenAI Response API""" @@ -1045,33 +1031,6 @@ class ResponseAPILoggingUtils: return True return False - @staticmethod - def _extra_fields_from_response_api_usage( - usage_input: dict | ResponseAPIUsage, - response_api_usage: ResponseAPIUsage, - ) -> dict[str, Any]: - """ - Preserve provider/extension usage fields not part of the standard token mapping. - - ResponseAPIUsage allows extra attributes; without forwarding them, the rebuilt - chat Usage would drop fields needed for provider-specific cost tracking. - """ - extras: dict[str, Any] = {} - if isinstance(usage_input, dict): - for key, value in usage_input.items(): - if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: - extras[key] = value - return extras - - model_extra = getattr(response_api_usage, "model_extra", None) or getattr( - response_api_usage, "__pydantic_extra__", None - ) - if isinstance(model_extra, dict): - for key, value in model_extra.items(): - if key not in ResponseAPILoggingUtils._RESPONSE_API_USAGE_MAPPED_KEYS and value is not None: - extras[key] = value - return extras - @staticmethod def _transform_response_api_usage_to_chat_usage( usage_input: dict | ResponseAPIUsage | None, @@ -1130,18 +1089,12 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(output_tokens_details, "audio_tokens", None), ) - usage_kwargs: Final = ResponseAPILoggingUtils._extra_fields_from_response_api_usage( - usage_input=usage_input, - response_api_usage=response_api_usage, - ) - chat_usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, prompt_tokens_details=prompt_tokens_details, completion_tokens_details=completion_tokens_details, - **usage_kwargs, ) # Preserve cost attribute if it exists on ResponseAPIUsage diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 8785ff1da43..0141cf5d96a 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -278,23 +278,6 @@ class TestResponseAPILoggingUtils: and result.prompt_tokens_details.cached_tokens == 2 ) - def test_transform_response_api_usage_preserves_extra_usage_fields(self): - """Non-standard usage keys pass through for provider cost tracking.""" - usage = { - "input_tokens": 10, - "output_tokens": 5, - "total_tokens": 15, - "server_side_tool_usage_details": {"web_search_calls": 2}, - "num_server_side_tools_used": 2, - } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) - assert getattr(result, "server_side_tool_usage_details", None) == { - "web_search_calls": 2 - } - assert getattr(result, "num_server_side_tools_used", None) == 2 - def test_transform_response_api_usage_with_none_values(self): """Test transformation handles None values properly""" # Setup