From 8687d7372ad3f2f40586fa40102b7571b678de52 Mon Sep 17 00:00:00 2001 From: Yang Yang Date: Fri, 19 Jun 2026 22:55:35 -0700 Subject: [PATCH] fix(xai): gate web search cost on server_side_tool_usage_details Treat positive web_search_calls as a web-search signal in built-in tool cost gating, and mirror counts onto prompt_tokens_details.web_search_requests when attaching xAI tool usage details so charges are not skipped. --- .../llm_cost_calc/tool_call_cost_tracking.py | 18 +++++++ litellm/llms/xai/chat/transformation.py | 5 +- litellm/llms/xai/cost_calculator.py | 25 +++++++++- litellm/llms/xai/responses/transformation.py | 7 ++- .../llms/xai/test_xai_cost_calculator.py | 50 ++++++++++++++++++- 5 files changed, 99 insertions(+), 6 deletions(-) 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..eb9473438f1 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 @@ -311,6 +311,22 @@ class StandardBuiltInToolCostTracking: return Usage(server_tool_use=server_tool_use) return usage.model_copy(update={"server_tool_use": server_tool_use}) + @staticmethod + def _usage_has_server_side_web_search_calls(usage: Usage | None) -> bool: + """True when usage.server_side_tool_usage_details.web_search_calls > 0.""" + if usage is None: + return False + details = getattr(usage, "server_side_tool_usage_details", None) + if details is None: + return False + try: + web_search_calls = ( + details.get("web_search_calls") if isinstance(details, dict) else getattr(details, "web_search_calls", None) + ) + return int(web_search_calls or 0) > 0 + except (TypeError, ValueError): + return False + @staticmethod def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool: """ @@ -328,6 +344,8 @@ class StandardBuiltInToolCostTracking: if get_anthropic_web_search_requests_from_response(response_object) is not None: return True + if StandardBuiltInToolCostTracking._usage_has_server_side_web_search_calls(usage): + return True if isinstance(response_object, ModelResponse): # chat completions only include url_citation annotations when a web search call is made diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index e9e9f205f94..1129872163c 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -12,6 +12,9 @@ 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 ( @@ -362,7 +365,7 @@ class XAIChatConfig(OpenAIGPTConfig): return details = response_usage.get("server_side_tool_usage_details") if details is not None: - setattr(usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(usage, details) verbose_logger.debug("X.AI server_side_tool_usage_details: %s", details) @staticmethod diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index dcc96625975..981cdb7ac88 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -5,10 +5,10 @@ Helper util for handling XAI-specific cost calculation """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Any, 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 @@ -17,6 +17,27 @@ if TYPE_CHECKING: _WEB_SEARCH_COST_PER_CALL = 5.0 / 1000.0 +def apply_server_side_tool_usage_details_to_usage( + usage: Usage, details: Mapping[str, Any] | 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 + setattr(usage, "server_side_tool_usage_details", details) + try: + web_search_calls = int(details.get("web_search_calls") or 0) + except (TypeError, ValueError): + return + if web_search_calls <= 0: + return + if usage.prompt_tokens_details is None: + usage.prompt_tokens_details = PromptTokensDetailsWrapper() + usage.prompt_tokens_details.web_search_requests = web_search_calls + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 7cf2f48f956..983f20fcf03 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -8,6 +8,9 @@ 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.llms.xai.cost_calculator import ( + apply_server_side_tool_usage_details_to_usage, +) from litellm.responses.utils import ResponseAPILoggingUtils from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -106,11 +109,11 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return if isinstance(response.usage, Usage): - setattr(response.usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(response.usage, details) return chat_usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response.usage) - setattr(chat_usage, "server_side_tool_usage_details", details) + apply_server_side_tool_usage_details_to_usage(chat_usage, details) response.usage = chat_usage # type: ignore[assignment] def _transform_web_search_tool(self, tool: dict[str, Any]) -> XAIWebSearchTool | dict[str, Any]: 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 b166778deaf..6435cf6218b 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -16,7 +16,15 @@ 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 ( + apply_server_side_tool_usage_details_to_usage, + cost_per_token, + cost_per_web_search_request, +) +from litellm.types.llms.openai import ResponsesAPIResponse class TestXAICostCalculator: @@ -376,6 +384,46 @@ class TestXAICostCalculator: usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) assert cost_per_web_search_request(usage=usage, model_info={}) == 0.0 + 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} + ) + 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 + ) + + def test_gate_detects_server_side_tool_usage_details_without_web_search_output( + self, + ): + usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + setattr( + usage, + "server_side_tool_usage_details", + {"web_search_calls": 1}, + ) + response = ResponsesAPIResponse.model_construct( + id="resp_test", + created_at=0, + output=[{"type": "message", "role": "assistant", "content": []}], + usage=None, + ) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=response, usage=usage + ) + assert ( + StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="grok-4.3", + response_object=response, + usage=usage, + standard_built_in_tools_params={}, + custom_llm_provider="xai", + ) + == 5.0 / 1000.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)