diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 3744be5bc79..2863c9c15cb 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -2,6 +2,7 @@ Helper utilities for tracking the cost of built-in tools. """ +from collections.abc import Mapping from typing import Any, Final, Literal import litellm @@ -23,6 +24,14 @@ from litellm.types.utils import ( ) +def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: + details: Final = getattr(usage, "server_side_tool_usage_details", None) + if not isinstance(details, Mapping): + return False + calls: Final = details.get("web_search_calls") + return isinstance(calls, int) and calls > 0 + + class StandardBuiltInToolCostTracking: """ Helper class for tracking the cost of built-in tools @@ -351,6 +360,10 @@ class StandardBuiltInToolCostTracking: # and _handle_web_search_cost() is never called. if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None: return True + # xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched + # answer with no url_citation annotations has no other chat-path signal + if _usage_reports_server_side_web_search_calls(usage): + return True return False elif isinstance(response_object, ResponsesAPIResponse): # response api explicitly includes web_search_call in the output @@ -370,6 +383,8 @@ class StandardBuiltInToolCostTracking: ) ): return True + if _usage_reports_server_side_web_search_calls(usage): + return True return False @@ -432,7 +447,9 @@ class StandardBuiltInToolCostTracking: """ output: Final = response_object.output for output_item in output: - _output_type: str | None = getattr(output_item, "type", None) + _output_type: str | None = ( + output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None) + ) if _output_type == output_type: return True return False diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 9d06b609752..ae5849812bf 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from typing import Any, Final import httpx @@ -12,13 +12,15 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_messages, ) from litellm.llms.xai.common_utils import XAIModelInfo +from litellm.llms.xai.cost_calculator import ( + apply_server_side_tool_usage_details_to_usage, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ( Choices, ModelResponse, ModelResponseStream, - PromptTokensDetailsWrapper, Usage, ) @@ -248,7 +250,7 @@ class XAIChatConfig(OpenAIGPTConfig): XAI API returns empty string for finish_reason when using tools, so we need to fix this after the standard OpenAI transformation. - Also handles X.AI web search usage tracking by extracting num_sources_used. + Also handles X.AI web search usage tracking. """ # First, let the parent class handle the standard transformation @@ -351,25 +353,20 @@ class XAIChatConfig(OpenAIGPTConfig): def _enhance_usage_with_xai_web_search_fields(self, model_response: ModelResponse, raw_response_json: dict) -> None: """ - Extract num_sources_used from X.AI response and map it to web_search_requests. + Copy usage.server_side_tool_usage_details from the provider usage block + onto model_response.usage for tool cost calculation. """ if not hasattr(model_response, "usage") or model_response.usage is None: return usage: Final[Usage] = model_response.usage - num_sources_used = None - response_usage: Final = raw_response_json.get("usage", {}) - if isinstance(response_usage, dict) and "num_sources_used" in response_usage: - num_sources_used = response_usage.get("num_sources_used") - - # Map num_sources_used to web_search_requests for cost detection - if num_sources_used is not None and num_sources_used > 0: - if usage.prompt_tokens_details is None: - usage.prompt_tokens_details = PromptTokensDetailsWrapper() - - usage.prompt_tokens_details.web_search_requests = int(num_sources_used) - setattr(usage, "num_sources_used", int(num_sources_used)) - verbose_logger.debug("X.AI web search sources used: %s", num_sources_used) + response_usage: Final = raw_response_json.get("usage") + if not isinstance(response_usage, dict): + return + details: Final = response_usage.get("server_side_tool_usage_details") + if isinstance(details, Mapping): + apply_server_side_tool_usage_details_to_usage(usage, details) + verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details) @staticmethod def _normalize_openai_compatible_usage_totals( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 384388f3300..dd77b8d5d09 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -4,14 +4,37 @@ Helper util for handling XAI-specific cost calculation - Handles XAI-specific reasoning token billing (billed as part of completion tokens) """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Final from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import Usage +from litellm.types.utils import PromptTokensDetailsWrapper, Usage if TYPE_CHECKING: from litellm.types.utils import ModelInfo +# https://docs.x.ai/developers/pricing#tools-pricing — default when unset in model map +_DEFAULT_WEB_SEARCH_COST_PER_CALL: Final = 5.0 / 1000.0 + + +def apply_server_side_tool_usage_details_to_usage(usage: Usage, details: Mapping[str, object] | None) -> None: + """ + Attach server_side_tool_usage_details and mirror web_search_calls onto + prompt_tokens_details.web_search_requests for built-in tool cost gating. + """ + if details is None: + return + usage.server_side_tool_usage_details = details # pyright: ignore[reportAttributeAccessIssue] # extra # rebind-ok: extras + try: + web_search_calls: Final = int(details.get("web_search_calls") or 0) + except (TypeError, ValueError): + return + if web_search_calls <= 0: + return + prompt_tokens_details: Final = usage.prompt_tokens_details or PromptTokensDetailsWrapper() + prompt_tokens_details.web_search_requests = web_search_calls + usage.prompt_tokens_details = prompt_tokens_details # rebind-ok: write details onto caller usage + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ @@ -32,9 +55,11 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: prompt_tokens: Final = int(getattr(usage, "prompt_tokens", 0) or 0) completion_tokens: Final = int(getattr(usage, "completion_tokens", 0) or 0) total_tokens: Final = int(getattr(usage, "total_tokens", 0) or 0) - reasoning_tokens = 0 - if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details: - reasoning_tokens = int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0) + reasoning_tokens: Final = ( + int(getattr(usage.completion_tokens_details, "reasoning_tokens", 0) or 0) + if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details + else 0 + ) already_normalised: Final = total_tokens == prompt_tokens + completion_tokens total_completion_tokens: Final = completion_tokens if already_normalised else completion_tokens + reasoning_tokens @@ -52,33 +77,48 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: return prompt_cost, completion_cost +def _web_search_cost_per_call_from_model_info(model_info: "ModelInfo") -> float: + """ + Per-invocation web_search price from model_info when configured. + + Prefer ``search_context_cost_per_query`` (same shape as Gemini/Anthropic web + search pricing in the model cost map). Fall back to current xAI list pricing. + """ + search_costs: Final = model_info.get("search_context_cost_per_query") + if not isinstance(search_costs, Mapping): + return _DEFAULT_WEB_SEARCH_COST_PER_CALL + for key in ( + "search_context_size_medium", + "search_context_size_low", + "search_context_size_high", + ): + value = search_costs.get(key) + if value is None: + continue + try: + cost = float(value) + except (TypeError, ValueError): + continue + if cost > 0: + return cost + return _DEFAULT_WEB_SEARCH_COST_PER_CALL + + def cost_per_web_search_request(usage: "Usage", model_info: "ModelInfo") -> float: """ Calculate the cost of web search requests for X.AI models. - X.AI Live Search costs $25 per 1,000 sources used. - Each source costs $0.025. - - The number of sources is stored in prompt_tokens_details.web_search_requests - by the transformation layer to be compatible with the existing detection system. + Counts invocations from usage.server_side_tool_usage_details.web_search_calls. + Per-call rate comes from model_info.search_context_cost_per_query when set, + otherwise the default xAI tools rate ($5 / 1k calls). """ - # Cost per source used: $25 per 1,000 sources = $0.025 per source - cost_per_source: Final = 25.0 / 1000.0 # $0.025 - - num_sources_used = 0 - - if ( - hasattr(usage, "prompt_tokens_details") - and usage.prompt_tokens_details is not None - and hasattr(usage.prompt_tokens_details, "web_search_requests") - and usage.prompt_tokens_details.web_search_requests is not None - ): - num_sources_used = int(usage.prompt_tokens_details.web_search_requests) - - # Fallback: try to get from num_sources_used if set directly - elif hasattr(usage, "num_sources_used") and usage.num_sources_used is not None: - num_sources_used = int(usage.num_sources_used) - - total_cost: Final = cost_per_source * num_sources_used - - return total_cost + details: Final = getattr(usage, "server_side_tool_usage_details", None) + if not isinstance(details, Mapping): + return 0.0 + try: + web_search_calls: Final = int(details.get("web_search_calls") or 0) + except (TypeError, ValueError): + return 0.0 + if web_search_calls <= 0: + return 0.0 + return _web_search_cost_per_call_from_model_info(model_info) * web_search_calls diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 48fb95d9411..d79e7d4c146 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import Any, Final import litellm from litellm._logging import verbose_logger @@ -12,13 +12,6 @@ from litellm.types.llms.xai import XAIWebSearchTool, XAIXSearchTool from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - - LiteLLMLoggingObj = _LiteLLMLoggingObj -else: - LiteLLMLoggingObj = Any - class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): """ diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index db2e515609c..1907b5aa447 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1033,13 +1033,18 @@ class ResponseAPILoggingUtils: @staticmethod def _transform_response_api_usage_to_chat_usage( - usage_input: dict | ResponseAPIUsage | None, + usage_input: Mapping[str, object] | ResponseAPIUsage | Usage | None, ) -> Usage: """ Transforms ResponseAPIUsage or ImageUsage to a Usage object. Both have the same spec with input_tokens, output_tokens, and input_tokens_details (text_tokens, image_tokens). + + Usage inputs are returned as-is so re-running this helper never drops + fields. Non-standard provider fields (e.g. xAI's + server_side_tool_usage_details) are carried onto the returned Usage so + provider cost calculators can read them after normalization. """ if usage_input is None: return Usage( @@ -1047,6 +1052,10 @@ class ResponseAPILoggingUtils: completion_tokens=0, total_tokens=0, ) + if isinstance(usage_input, Usage): + return usage_input + if isinstance(usage_input, dict) and not ResponseAPILoggingUtils._is_response_api_usage(usage_input): + return Usage(**usage_input) response_api_usage: ResponseAPIUsage if isinstance(usage_input, dict): usage_input = dict(usage_input) # shallow copy; avoid mutating caller @@ -1055,13 +1064,11 @@ class ResponseAPILoggingUtils: usage_input["input_tokens_details"] = usage_input["input_token_details"] if usage_input.get("output_tokens_details") is None and "output_token_details" in usage_input: usage_input["output_tokens_details"] = usage_input["output_token_details"] - total_tokens = usage_input.get("total_tokens") - if total_tokens is None: + if usage_input.get("total_tokens") is None: input_tokens: Final = usage_input.get("input_tokens") output_tokens: Final = usage_input.get("output_tokens") - if input_tokens is not None and output_tokens is not None: - total_tokens = input_tokens + output_tokens - usage_input["total_tokens"] = total_tokens + if isinstance(input_tokens, int) and isinstance(output_tokens, int): + usage_input["total_tokens"] = input_tokens + output_tokens response_api_usage = ResponseAPIUsage(**usage_input) else: response_api_usage = usage_input @@ -1089,12 +1096,27 @@ class ResponseAPILoggingUtils: audio_tokens=getattr(output_tokens_details, "audio_tokens", None), ) + extra_usage_fields: Final = { + key: value + for key, value in (response_api_usage.model_extra or {}).items() + if key + not in ( + "input_token_details", + "output_token_details", + "prompt_tokens", + "completion_tokens", + "total_tokens", + "prompt_tokens_details", + "completion_tokens_details", + ) + } 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, + **extra_usage_fields, ) # Preserve cost attribute if it exists on ResponseAPIUsage diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 24fd3c94ee3..7f735982129 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -339,7 +339,7 @@ def test_get_cost_for_vertex_ai_gemini_web_search(model, custom_llm_provider): for url_citation annotations, not usage.prompt_tokens_details.web_search_requests. This causes Vertex AI grounding costs to not be tracked. """ - from litellm.types.utils import PromptTokensDetailsWrapper, Usage, Choices, Message + from litellm.types.utils import Choices, Message, PromptTokensDetailsWrapper, Usage # Create a realistic ModelResponse like what Vertex AI returns response = ModelResponse( @@ -604,3 +604,66 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model( # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage + + +def test_response_includes_output_type_reads_dict_output_items(): + """ + Regression: output items that fail OpenAI SDK validation (e.g. xAI web_search_call + items without an "action" field) stay plain dicts in the output union. The gate must + read their "type" key instead of returning False and skipping the web search fee. + """ + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 1754900000, + "model": "grok-4", + "object": "response", + "status": "completed", + "output": [{"type": "web_search_call", "id": "ws_1", "status": "completed"}], + } + ) + + assert isinstance(response.output[0], dict) + assert StandardBuiltInToolCostTracking.response_includes_output_type( + response_object=response, output_type="web_search_call" + ) + assert not StandardBuiltInToolCostTracking.response_includes_output_type( + response_object=response, output_type="file_search_call" + ) + + +def test_web_search_gate_reads_server_side_tool_usage_details_without_citations(): + """ + Regression: xAI chat responses bridged from the Responses API only carry + usage.server_side_tool_usage_details; a searched answer with no url_citation + annotations must still be billed for its web search calls. + """ + from litellm.llms.xai.cost_calculator import _DEFAULT_WEB_SEARCH_COST_PER_CALL + from litellm.types.utils import Usage + + usage = Usage( + prompt_tokens=10, + completion_tokens=20, + total_tokens=30, + server_side_tool_usage_details={"web_search_calls": 3}, + ) + response = ModelResponse(model="xai/grok-4.5") + + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage + ) + assert not StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="xai/grok-4.5", + response_object=response, + usage=usage, + custom_llm_provider="xai", + standard_built_in_tools_params=None, + ) + assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index fb98dc0a917..871613c9c9a 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -9,14 +9,22 @@ Source: litellm/llms/xai/responses/transformation.py import os import sys +from unittest.mock import MagicMock sys.path.insert(0, os.path.abspath("../../../../..")) import pytest +import litellm from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams -from litellm.types.utils import LlmProviders +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) +from litellm.types.utils import LlmProviders, Usage from litellm.utils import ProviderConfigManager @@ -31,43 +39,29 @@ class TestXAIResponsesAPITransformation: ) assert config is not None, "Config should not be None for XAI provider" - assert isinstance( - config, XAIResponsesAPIConfig - ), f"Expected XAIResponsesAPIConfig, got {type(config)}" - assert ( - config.custom_llm_provider == LlmProviders.XAI - ), "custom_llm_provider should be XAI" + assert isinstance(config, XAIResponsesAPIConfig), f"Expected XAIResponsesAPIConfig, got {type(config)}" + assert config.custom_llm_provider == LlmProviders.XAI, "custom_llm_provider should be XAI" def test_code_interpreter_container_field_removed(self): """Test that container field is removed from code_interpreter tools""" config = XAIResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - tools=[{"type": "code_interpreter", "container": {"type": "auto"}}] - ) + params = ResponsesAPIOptionalRequestParams(tools=[{"type": "code_interpreter", "container": {"type": "auto"}}]) - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) + result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False) assert "tools" in result assert len(result["tools"]) == 1 assert result["tools"][0]["type"] == "code_interpreter" - assert ( - "container" not in result["tools"][0] - ), "Container field should be removed" + assert "container" not in result["tools"][0], "Container field should be removed" def test_instructions_parameter_dropped(self): """Test that instructions parameter is dropped for XAI""" config = XAIResponsesAPIConfig() - params = ResponsesAPIOptionalRequestParams( - instructions="You are a helpful assistant.", temperature=0.7 - ) + params = ResponsesAPIOptionalRequestParams(instructions="You are a helpful assistant.", temperature=0.7) - result = config.map_openai_params( - response_api_optional_params=params, model="grok-4-fast", drop_params=False - ) + result = config.map_openai_params(response_api_optional_params=params, model="grok-4-fast", drop_params=False) assert "instructions" not in result, "Instructions should be dropped" assert result.get("temperature") == 0.7, "Other params should be preserved" @@ -88,25 +82,15 @@ class TestXAIResponsesAPITransformation: # Test with default XAI API base url = config.get_complete_url(api_base=None, litellm_params={}) - assert ( - url == "https://api.x.ai/v1/responses" - ), f"Expected XAI responses endpoint, got {url}" + assert url == "https://api.x.ai/v1/responses", f"Expected XAI responses endpoint, got {url}" # Test with custom api_base - custom_url = config.get_complete_url( - api_base="https://custom.x.ai/v1", litellm_params={} - ) - assert ( - custom_url == "https://custom.x.ai/v1/responses" - ), f"Expected custom endpoint, got {custom_url}" + custom_url = config.get_complete_url(api_base="https://custom.x.ai/v1", litellm_params={}) + assert custom_url == "https://custom.x.ai/v1/responses", f"Expected custom endpoint, got {custom_url}" # Test with trailing slash - url_with_slash = config.get_complete_url( - api_base="https://api.x.ai/v1/", litellm_params={} - ) - assert ( - url_with_slash == "https://api.x.ai/v1/responses" - ), "Should handle trailing slash" + url_with_slash = config.get_complete_url(api_base="https://api.x.ai/v1/", litellm_params={}) + assert url_with_slash == "https://api.x.ai/v1/responses", "Should handle trailing slash" def test_web_search_tool_transformation(self): """Test that web_search tools are transformed to XAI format""" @@ -167,9 +151,7 @@ class TestXAIResponsesAPITransformation: config = XAIResponsesAPIConfig() params = ResponsesAPIOptionalRequestParams( - tools=[ - {"type": "web_search", "excluded_domains": ["example.com", "test.com"]} - ] + tools=[{"type": "web_search", "excluded_domains": ["example.com", "test.com"]}] ) result = config.map_openai_params( @@ -309,3 +291,115 @@ class TestXAIResponsesAPITransformation: # Verify function tool is unchanged assert result["tools"][3]["type"] == "function" assert result["tools"][3]["name"] == "get_weather" + + +class TestXAIResponsesWebSearchBilling: + """Web search billing must not change the client-visible Responses usage schema.""" + + _TOOL_DETAILS = { + "web_search_calls": 2, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + } + + def _raw_response_json(self, include_web_search: bool) -> dict: + web_search_output = ( + [{ + "type": "web_search_call", + "id": "ws_1", + "status": "completed", + "action": {"type": "search", "query": "grok"}, + }] if include_web_search else [] + ) + tool_usage = {"server_side_tool_usage_details": self._TOOL_DETAILS} if include_web_search else {} + return { + "id": "resp_1", + "object": "response", + "created_at": 1754900000, + "model": "grok-4", + "status": "completed", + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "top_p": 1.0, + "output": web_search_output + + [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "grok says hi", "annotations": []}], + } + ], + "usage": { + "input_tokens": 100, + "output_tokens": 20, + "total_tokens": 120, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + **tool_usage, + }, + } + + def _transform(self, include_web_search: bool) -> ResponsesAPIResponse: + raw_response = MagicMock() + raw_response.json.return_value = self._raw_response_json(include_web_search) + raw_response.text = "raw" + raw_response.headers = {} + return XAIResponsesAPIConfig().transform_response_api_response( + model="grok-4", raw_response=raw_response, logging_obj=MagicMock() + ) + + def test_response_usage_keeps_responses_api_schema(self): + response = self._transform(include_web_search=True) + + assert isinstance(response.usage, ResponseAPIUsage) + assert response.usage.input_tokens == 100 + assert response.usage.output_tokens == 20 + assert response.usage.model_extra["server_side_tool_usage_details"] == self._TOOL_DETAILS + + def test_bridged_usage_keeps_tool_details_for_billing(self): + response = self._transform(include_web_search=True) + + bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) + + assert isinstance(bridged, Usage) + assert bridged.prompt_tokens == 100 + assert bridged.completion_tokens == 20 + assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS + + def test_completion_cost_bills_web_search_calls(self): + with_search = litellm.completion_cost( + completion_response=self._transform(include_web_search=True), + model="xai/grok-4", + custom_llm_provider="xai", + ) + without_search = litellm.completion_cost( + completion_response=self._transform(include_web_search=False), + model="xai/grok-4", + custom_llm_provider="xai", + ) + + assert with_search - without_search == pytest.approx(2 * 5.0 / 1000.0) + + def test_streaming_terminal_event_keeps_schema_and_details(self): + parsed_chunk = { + "type": "response.completed", + "sequence_number": 7, + "response": self._raw_response_json(include_web_search=True), + } + + event = XAIResponsesAPIConfig().transform_streaming_response( + model="grok-4", parsed_chunk=parsed_chunk, logging_obj=MagicMock() + ) + + assert isinstance(event, ResponseCompletedEvent) + assert isinstance(event.response.usage, ResponseAPIUsage) + assert event.response.usage.input_tokens == 100 + + bridged = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(event.response.usage) + assert getattr(bridged, "server_side_tool_usage_details") == self._TOOL_DETAILS diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index 5c1f0f704d7..eac5b89e4f3 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -5,6 +5,9 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path +import pytest + +import litellm from litellm.llms.xai.chat.transformation import XAIChatConfig from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -135,3 +138,65 @@ class TestXAIUsageNormalization: XAIChatConfig._normalize_openai_compatible_usage_totals(usage) assert usage["total_tokens"] == 200 + + +class TestXAIChatWebSearchBilling: + _TOOL_DETAILS = { + "web_search_calls": 3, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + } + + @staticmethod + def _response_with_usage() -> ModelResponse: + response = ModelResponse(model="grok-4") + setattr( + response, + "usage", + Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120), + ) + return response + + def test_enhance_copies_details_and_mirrors_web_search_requests(self): + response = self._response_with_usage() + + XAIChatConfig()._enhance_usage_with_xai_web_search_fields( + response, + {"usage": {"server_side_tool_usage_details": self._TOOL_DETAILS}}, + ) + + usage = response.usage + assert getattr(usage, "server_side_tool_usage_details") == self._TOOL_DETAILS + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 3 + + def test_enhance_noop_without_details(self): + response = self._response_with_usage() + + XAIChatConfig()._enhance_usage_with_xai_web_search_fields( + response, {"usage": {"prompt_tokens": 100}} + ) + + assert response.usage.prompt_tokens_details is None + assert getattr(response.usage, "server_side_tool_usage_details", None) is None + + def test_completion_cost_bills_chat_web_search_calls(self): + billed = self._response_with_usage() + XAIChatConfig()._enhance_usage_with_xai_web_search_fields( + billed, + {"usage": {"server_side_tool_usage_details": self._TOOL_DETAILS}}, + ) + + with_search = litellm.completion_cost( + completion_response=billed, model="xai/grok-4", custom_llm_provider="xai" + ) + without_search = litellm.completion_cost( + completion_response=self._response_with_usage(), + model="xai/grok-4", + custom_llm_provider="xai", + ) + + assert with_search - without_search == pytest.approx(3 * 5.0 / 1000.0) diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 32141bead0e..df6f4d3edd8 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -17,7 +17,16 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -from litellm.llms.xai.cost_calculator import cost_per_token, cost_per_web_search_request +from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( + StandardBuiltInToolCostTracking, +) +from litellm.llms.xai.cost_calculator import ( + _DEFAULT_WEB_SEARCH_COST_PER_CALL, + _web_search_cost_per_call_from_model_info, + apply_server_side_tool_usage_details_to_usage, + cost_per_token, + cost_per_web_search_request, +) class TestXAICostCalculator: @@ -354,76 +363,53 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) - def test_web_search_cost_calculation(self): - """Test web search cost calculation for X.AI models.""" - # Test with web_search_requests in prompt_tokens_details (primary path) - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=100, - web_search_requests=3, # 3 sources used - ), + def test_web_search_cost_via_server_side_tool_usage_details(self): + """usage.server_side_tool_usage_details.web_search_calls at default $5/1k.""" + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + setattr( + usage, + "server_side_tool_usage_details", + { + "web_search_calls": 3, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + }, ) web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) + assert math.isclose(web_search_cost, 3 * (5.0 / 1000.0), rel_tol=1e-10) - # Expected cost: 3 sources * $0.025 per source = $0.075 - expected_cost = 3 * (25.0 / 1000.0) # 3 * $0.025 - - assert math.isclose(web_search_cost, expected_cost, rel_tol=1e-10) - assert math.isclose(web_search_cost, 0.075, rel_tol=1e-10) - - def test_web_search_cost_fallback_calculation(self): - """Test web search cost calculation using fallback num_sources_used.""" - # Test fallback: num_sources_used on usage object - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, + def test_web_search_cost_uses_model_info_search_context_pricing(self): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 2}) + model_info = { + "search_context_cost_per_query": { + "search_context_size_medium": 0.01, + } + } + web_search_cost = cost_per_web_search_request( + usage=usage, model_info=model_info ) - # Manually set num_sources_used (as done by transformation layer) - setattr(usage, "num_sources_used", 5) + assert math.isclose(web_search_cost, 0.02, rel_tol=1e-10) - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) + def test_web_search_cost_zero_without_details(self): + usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 - # Expected cost: 5 sources * $0.025 per source = $0.125 - expected_cost = 5 * (25.0 / 1000.0) # 5 * $0.025 - - assert math.isclose(web_search_cost, expected_cost, rel_tol=1e-10) - assert math.isclose(web_search_cost, 0.125, rel_tol=1e-10) - - def test_web_search_no_sources_used(self): - """Test web search cost calculation when no sources are used.""" - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=100, - web_search_requests=0, # No web search - ), + def test_apply_details_sets_web_search_requests_for_cost_gate(self): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + apply_server_side_tool_usage_details_to_usage( + usage, {"web_search_calls": 2, "x_search_calls": 0} ) - - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) - - # Expected cost: 0 sources * $0.025 per source = $0.0 - assert web_search_cost == 0.0 - - def test_web_search_cost_without_prompt_tokens_details(self): - """Test web search cost calculation when prompt_tokens_details is None.""" - usage = Usage( - prompt_tokens=100, - completion_tokens=50, - total_tokens=150, + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 2 + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=object(), usage=usage ) - web_search_cost = cost_per_web_search_request(usage=usage, model_info={}) - - # Expected cost: No web search data = $0.0 - assert web_search_cost == 0.0 - def test_grok_4_20_beta_reasoning_cost_calculation(self): """Test cost calculation for grok-4.20-beta-0309-reasoning model.""" usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) @@ -499,3 +485,112 @@ class TestXAICostCalculator: assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + +class TestXAIWebSearchCostHelpers: + """Focused coverage for web_search / tool-usage helpers in cost_calculator.py.""" + + def test_apply_details_noop_when_details_none(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + apply_server_side_tool_usage_details_to_usage(usage, None) + assert getattr(usage, "server_side_tool_usage_details", None) is None + + def test_apply_details_sets_attr_but_skips_mirror_when_web_search_zero(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + details = {"web_search_calls": 0, "x_search_calls": 3} + apply_server_side_tool_usage_details_to_usage(usage, details) + assert getattr(usage, "server_side_tool_usage_details") == details + assert ( + usage.prompt_tokens_details is None + or usage.prompt_tokens_details.web_search_requests is None + ) + + def test_apply_details_skips_mirror_when_web_search_calls_invalid(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + details = {"web_search_calls": "not-a-number"} + apply_server_side_tool_usage_details_to_usage(usage, details) + assert getattr(usage, "server_side_tool_usage_details") == details + assert usage.prompt_tokens_details is None + + def test_apply_details_updates_existing_prompt_tokens_details(self): + usage = Usage( + prompt_tokens=1, + completion_tokens=1, + total_tokens=2, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=7), + ) + apply_server_side_tool_usage_details_to_usage(usage, {"web_search_calls": 4}) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 7 + assert usage.prompt_tokens_details.web_search_requests == 4 + + def test_web_search_cost_per_call_default_when_model_info_empty(self): + assert ( + _web_search_cost_per_call_from_model_info({}) + == _DEFAULT_WEB_SEARCH_COST_PER_CALL + ) + + def test_web_search_cost_per_call_prefers_medium_over_low(self): + model_info = { + "search_context_cost_per_query": { + "search_context_size_low": 0.001, + "search_context_size_medium": 0.009, + } + } + assert _web_search_cost_per_call_from_model_info(model_info) == 0.009 + + def test_web_search_cost_per_call_falls_back_to_low_then_high(self): + assert ( + _web_search_cost_per_call_from_model_info( + {"search_context_cost_per_query": {"search_context_size_low": 0.003}} + ) + == 0.003 + ) + assert ( + _web_search_cost_per_call_from_model_info( + {"search_context_cost_per_query": {"search_context_size_high": 0.007}} + ) + == 0.007 + ) + + def test_web_search_cost_per_call_ignores_zero_and_invalid_values(self): + assert ( + _web_search_cost_per_call_from_model_info( + { + "search_context_cost_per_query": { + "search_context_size_medium": 0, + "search_context_size_low": "bad", + } + } + ) + == _DEFAULT_WEB_SEARCH_COST_PER_CALL + ) + + def test_cost_per_web_search_request_zero_when_details_not_mapping(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr(usage, "server_side_tool_usage_details", "invalid") + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_zero_when_web_search_calls_invalid(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": object()}, + ) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_zero_when_web_search_calls_zero(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": 0, "x_search_calls": 5}, + ) + assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + + def test_cost_per_web_search_request_uses_default_rate_without_model_pricing(self): + usage = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + setattr(usage, "server_side_tool_usage_details", {"web_search_calls": 4}) + cost = cost_per_web_search_request(usage=usage, model_info={}) + assert math.isclose(cost, 4 * _DEFAULT_WEB_SEARCH_COST_PER_CALL, rel_tol=1e-10) diff --git a/tests/test_litellm/responses/test_responses_utils.py b/tests/test_litellm/responses/test_responses_utils.py index 0141cf5d96a..2b9e6d34828 100644 --- a/tests/test_litellm/responses/test_responses_utils.py +++ b/tests/test_litellm/responses/test_responses_utils.py @@ -1,21 +1,16 @@ import base64 -import json import os import sys from unittest.mock import MagicMock, patch import pytest -from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path import litellm -from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils -from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIOptionalRequestParams from litellm.types.utils import Usage @@ -54,9 +49,7 @@ class TestResponsesAPIRequestUtils: # Setup model = "gpt-4o" config = OpenAIResponsesAPIConfig() - optional_params = ResponsesAPIOptionalRequestParams( - {"temperature": 0.7, "unsupported_param": "value"} - ) + optional_params = ResponsesAPIOptionalRequestParams({"temperature": 0.7, "unsupported_param": "value"}) # Execute and Assert with pytest.raises(litellm.UnsupportedParamsError) as excinfo: @@ -90,9 +83,7 @@ class TestResponsesAPIRequestUtils: assert result == {"temperature": 0.7} @pytest.mark.parametrize("request_drop_params", [None, False]) - def test_get_optional_params_responses_api_still_raises_without_drop( - self, monkeypatch, request_drop_params - ): + def test_get_optional_params_responses_api_still_raises_without_drop(self, monkeypatch, request_drop_params): """Absent or False request-level drop_params must not suppress the unsupported-param error""" monkeypatch.setattr(litellm, "drop_params", False) config = OpenAIResponsesAPIConfig() @@ -119,9 +110,7 @@ class TestResponsesAPIRequestUtils: } # Execute - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) # Assert assert "temperature" in result @@ -147,40 +136,31 @@ class TestResponsesAPIRequestUtils: ) # Execute - result = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( - encoded_id - ) + result = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(encoded_id) # Assert assert result == original_response_id # Test with a non-encoded ID plain_id = "resp_xyz789" - result_plain = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id( - plain_id - ) + result_plain = ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(plain_id) assert result_plain == plain_id def test_update_responses_api_response_id_with_model_id_handles_dict(self): """Ensure _update_responses_api_response_id_with_model_id works with dict input""" responses_api_response = {"id": "resp_abc123"} litellm_metadata = {"model_info": {"id": "gpt-4o"}} - updated = ( - ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( - responses_api_response=responses_api_response, - custom_llm_provider="openai", - litellm_metadata=litellm_metadata, - ) + updated = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( + responses_api_response=responses_api_response, + custom_llm_provider="openai", + litellm_metadata=litellm_metadata, ) assert updated["id"] != "resp_abc123" - decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id( - updated["id"] - ) + decoded = ResponsesAPIRequestUtils._decode_responses_api_response_id(updated["id"]) assert decoded.get("response_id") == "resp_abc123" assert decoded.get("model_id") == "gpt-4o" assert decoded.get("custom_llm_provider") == "openai" - def test_update_responses_api_response_id_with_model_id_is_idempotent_for_litellm_ids(self): raw = "resp_" + "a" * 48 litellm_metadata = {"model_info": {"id": "model-123"}} @@ -207,9 +187,7 @@ class TestResponsesAPIRequestUtils: model_id=None, container_id="cntr_upstream_abc", ) - assert "None" not in base64.b64decode( - encoded.replace("cntr_", "").encode("utf-8") - ).decode("utf-8") + assert "None" not in base64.b64decode(encoded.replace("cntr_", "").encode("utf-8")).decode("utf-8") decoded = ResponsesAPIRequestUtils._decode_container_id(encoded) assert decoded.get("custom_llm_provider") == "azure" assert decoded.get("model_id") is None @@ -217,12 +195,8 @@ class TestResponsesAPIRequestUtils: def test_decode_container_id_legacy_literal_none_model_id(self): """IDs encoded before the None fix should decode without a bogus model_id.""" - legacy_inner = ( - "litellm:custom_llm_provider:azure;model_id:None;container_id:cntr_x" - ) - legacy_id = "cntr_" + base64.b64encode(legacy_inner.encode("utf-8")).decode( - "utf-8" - ) + legacy_inner = "litellm:custom_llm_provider:azure;model_id:None;container_id:cntr_x" + legacy_id = "cntr_" + base64.b64encode(legacy_inner.encode("utf-8")).decode("utf-8") decoded = ResponsesAPIRequestUtils._decode_container_id(legacy_id) assert decoded.get("model_id") is None assert decoded.get("custom_llm_provider") == "azure" @@ -264,19 +238,14 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert isinstance(result, Usage) assert result.prompt_tokens == 10 assert result.completion_tokens == 20 assert result.total_tokens == 30 - assert ( - result.prompt_tokens_details - and result.prompt_tokens_details.cached_tokens == 2 - ) + assert result.prompt_tokens_details and result.prompt_tokens_details.cached_tokens == 2 def test_transform_response_api_usage_with_none_values(self): """Test transformation handles None values properly""" @@ -289,9 +258,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert result.prompt_tokens == 0 @@ -310,9 +277,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert assert result.prompt_tokens == 15 @@ -349,9 +314,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert - verify basic token counts assert isinstance(result, Usage) @@ -386,9 +349,7 @@ class TestResponseAPILoggingUtils: }, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.cache_write_tokens == 10059 @@ -417,9 +378,7 @@ class TestResponseAPILoggingUtils: } # Execute - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) # Assert - all token detail types should be preserved assert result.prompt_tokens_details is not None @@ -451,9 +410,7 @@ class TestResponseAPILoggingUtils: }, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.text_tokens == 8 @@ -475,9 +432,7 @@ class TestResponseAPILoggingUtils: "output_token_details": {"text_tokens": 2, "audio_tokens": 98}, } - result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) assert result.prompt_tokens_details is not None assert result.prompt_tokens_details.text_tokens == 10 @@ -487,6 +442,93 @@ class TestResponseAPILoggingUtils: assert result.completion_tokens_details.text_tokens == 20 assert result.completion_tokens_details.audio_tokens is None + def test_transform_response_api_usage_carries_extra_provider_fields(self): + """Non-standard usage fields (e.g. xAI tool details) must survive chat normalization.""" + details = {"web_search_calls": 2, "x_search_calls": 0} + usage = ResponseAPIUsage( + input_tokens=100, + output_tokens=20, + total_tokens=120, + server_side_tool_usage_details=details, + ) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert isinstance(result, Usage) + assert result.prompt_tokens == 100 + assert result.completion_tokens == 20 + assert getattr(result, "server_side_tool_usage_details") == details + + def test_transform_response_api_usage_ignores_chat_shaped_extras(self): + """Gemini image usage carries chat-shaped keys as extras; they must not collide with explicit kwargs.""" + usage = ResponseAPIUsage( + input_tokens=35, + output_tokens=1716, + total_tokens=1751, + prompt_tokens=35, + prompt_tokens_details={"image_tokens": 5, "text_tokens": 30}, + completion_tokens=1716, + completion_tokens_details={"image_tokens": 1120, "text_tokens": 596}, + server_side_tool_usage_details={"web_search_calls": 1}, + ) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result.prompt_tokens == 35 + assert result.completion_tokens == 1716 + assert getattr(result, "server_side_tool_usage_details") == {"web_search_calls": 1} + + def test_transform_already_chat_usage_passthrough_keeps_tool_details(self): + """Re-running the bridge on an already-converted chat Usage must not drop fields.""" + details = {"web_search_calls": 2, "x_search_calls": 0} + usage = Usage( + prompt_tokens=100, + completion_tokens=20, + total_tokens=120, + prompt_tokens_details={"web_search_requests": 2}, + ) + setattr(usage, "server_side_tool_usage_details", details) + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert result is usage + assert getattr(result, "server_side_tool_usage_details") == details + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.web_search_requests == 2 + + def test_transform_chat_shaped_usage_dict_keeps_tool_details(self): + """Streaming chat bridge dumps already-converted Usage as a prompt_tokens dict.""" + details = { + "web_search_calls": 3, + "x_search_calls": 0, + "code_interpreter_calls": 0, + "file_search_calls": 0, + "mcp_calls": 0, + "document_search_calls": 0, + "image_generation_calls": 0, + } + usage = { + "prompt_tokens": 50, + "completion_tokens": 10, + "total_tokens": 60, + "prompt_tokens_details": {"web_search_requests": 3, "cached_tokens": 8}, + "completion_tokens_details": {"reasoning_tokens": 4}, + "server_side_tool_usage_details": details, + } + + result = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + + assert isinstance(result, Usage) + assert result.prompt_tokens == 50 + assert result.completion_tokens == 10 + assert result.total_tokens == 60 + assert getattr(result, "server_side_tool_usage_details") == details + assert result.prompt_tokens_details is not None + assert result.prompt_tokens_details.web_search_requests == 3 + assert result.prompt_tokens_details.cached_tokens == 8 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.reasoning_tokens == 4 + class TestResponsesAPIProviderSpecificParams: """ @@ -503,9 +545,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result def test_provider_specific_params_no_crash_with_openai(self): @@ -517,9 +557,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result def test_provider_specific_params_no_crash_with_vertex_ai(self): @@ -531,9 +569,7 @@ class TestResponsesAPIProviderSpecificParams: } # Should not raise any exception - result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param( - params - ) + result = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(params) assert "temperature" in result