From 475dfaa15626b480dde86cfacd9e94cc4f327499 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 22 Mar 2025 16:56:32 -0700 Subject: [PATCH] fixes for web search cost tracking --- litellm/litellm_core_utils/litellm_logging.py | 7 ++ .../llm_cost_calc/tool_call_cost_tracking.py | 21 ++++- litellm/types/utils.py | 1 + .../test_token_counting.py | 77 ++++++++++++++----- 4 files changed, 83 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f437755771c..8ce6b156686 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1089,6 +1089,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, logging_obj=self, status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) elif isinstance(result, dict): # pass-through endpoints @@ -1101,6 +1102,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, logging_obj=self, status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) elif standard_logging_object is not None: @@ -1178,6 +1180,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, logging_obj=self, status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) callbacks = self.get_combined_callback_list( @@ -1718,6 +1721,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, logging_obj=self, status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) callbacks = self.get_combined_callback_list( @@ -1934,6 +1938,7 @@ class Logging(LiteLLMLoggingBaseClass): status="failure", error_str=str(exception), original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, ) ) return start_time, end_time @@ -3390,6 +3395,7 @@ def get_standard_logging_object_payload( status: StandardLoggingPayloadStatus, error_str: Optional[str] = None, original_exception: Optional[Exception] = None, + standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, ) -> Optional[StandardLoggingPayload]: try: kwargs = kwargs or {} @@ -3565,6 +3571,7 @@ def get_standard_logging_object_payload( guardrail_information=metadata.get( "standard_logging_guardrail_information", None ), + standard_built_in_tools_params=standard_built_in_tools_params, ) emit_standard_logging_payload(payload) 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 81bf4fd5e71..1717565831d 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 @@ -50,9 +50,12 @@ class StandardBuiltInToolCostTracking: ), model_info=model_info, ) - # elif isinstance(response_object, ResponsesAPIResponse): - # if response_object.web_search_options is not None: - # return 0.0 + elif StandardBuiltInToolCostTracking.response_includes_annotations( + response_object + ): + return StandardBuiltInToolCostTracking.get_default_cost_for_web_search( + model_info + ) return 0.0 @staticmethod @@ -80,6 +83,18 @@ class StandardBuiltInToolCostTracking: else: return 0.0 + @staticmethod + def get_default_cost_for_web_search(model_info: ModelInfo) -> float: + """ + If no web search options are provided, use the `search_context_size_medium` pricing. + + https://platform.openai.com/docs/pricing#web-search + """ + search_context_pricing: SearchContextCostPerQuery = ( + model_info.get("search_context_cost_per_query", {}) or {} + ) or {} + return search_context_pricing.get("search_context_size_medium", 0.0) + @staticmethod def response_includes_annotations(response_object: ModelResponse) -> bool: for _choice in response_object.choices: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6b027ab8b77..1c1ee24bfbf 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1752,6 +1752,7 @@ class StandardLoggingPayload(TypedDict): model_parameters: dict hidden_params: StandardLoggingHiddenParams guardrail_information: Optional[StandardLoggingGuardrailInformation] + standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] from typing import AsyncIterator, Iterator diff --git a/tests/logging_callback_tests/test_token_counting.py b/tests/logging_callback_tests/test_token_counting.py index 6821c7d345c..05e21efe0d9 100644 --- a/tests/logging_callback_tests/test_token_counting.py +++ b/tests/logging_callback_tests/test_token_counting.py @@ -248,34 +248,23 @@ async def test_stream_token_counting_anthropic_with_include_usage(): ) -@pytest.mark.asyncio -async def test_openai_web_search_logging_cost_tracking(): - """Makes a simple web search request and validates the response contains web search annotations and all expected fields are present""" +async def _setup_web_search_test(): + """Helper function to setup common test requirements""" litellm._turn_on_debug() test_custom_logger = TestCustomLogger() litellm.callbacks = [test_custom_logger] - response = await litellm.acompletion( - model="openai/gpt-4o-search-preview", - messages=[ - { - "role": "user", - "content": "What was a positive news story from today?", - } - ], - ) - print("litellm response: ", response.model_dump_json(indent=4)) + return test_custom_logger + +async def _verify_web_search_cost(test_custom_logger, expected_context_size): + """Helper function to verify web search costs""" await asyncio.sleep(1) - print( - "logged standard logging payload: ", - json.dumps(test_custom_logger.standard_logging_payload, indent=4), - ) standard_logging_payload = test_custom_logger.standard_logging_payload response_cost = standard_logging_payload.get("response_cost") assert response_cost is not None - # Assert the cost = Token Usage + Web Search Cost + # Calculate token cost model_map_information = standard_logging_payload["model_map_information"] model_map_value: ModelInfoBase = model_map_information["model_map_value"] total_token_cost = ( @@ -285,9 +274,57 @@ async def test_openai_web_search_logging_cost_tracking(): standard_logging_payload["completion_tokens"] * model_map_value["output_cost_per_token"] ) - print("total token cost:", total_token_cost) + + # Verify total cost assert ( response_cost == total_token_cost - + model_map_value["search_context_cost_per_query"]["search_context_size_low"] + + model_map_value["search_context_cost_per_query"][expected_context_size] ) + + +@pytest.mark.asyncio +async def test_openai_web_search_logging_cost_tracking_no_explicit_search_context_size(): + """Cost is tracked as `search_context_size_medium` when no `search_context_size` is passed in""" + test_custom_logger = await _setup_web_search_test() + + response = await litellm.acompletion( + model="openai/gpt-4o-search-preview", + messages=[ + {"role": "user", "content": "What was a positive news story from today?"} + ], + ) + + await _verify_web_search_cost(test_custom_logger, "search_context_size_medium") + + +@pytest.mark.asyncio +async def test_openai_web_search_logging_cost_tracking_explicit_search_context_size(): + """search_context_size=low passed in, so cost tracked as `search_context_size_low`""" + test_custom_logger = await _setup_web_search_test() + + response = await litellm.acompletion( + model="openai/gpt-4o-search-preview", + messages=[ + {"role": "user", "content": "What was a positive news story from today?"} + ], + web_search_options={"search_context_size": "low"}, + ) + + await _verify_web_search_cost(test_custom_logger, "search_context_size_low") + + +@pytest.mark.asyncio +async def test_openai_web_search_with_tool_call_logging_cost_tracking(): + """search_context_size=high passed in tool call, so cost tracked as `search_context_size_high`""" + test_custom_logger = await _setup_web_search_test() + + response = await litellm.aresponses( + model="openai/gpt-4o", + input=[ + {"role": "user", "content": "What was a positive news story from today?"} + ], + tools=[{"type": "web_search_preview", "search_context_size": "high"}], + ) + + await _verify_web_search_cost(test_custom_logger, "search_context_size_high")