From 339da4183d998b1ff2db54bc0468fa994a620095 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 12:34:54 -0700 Subject: [PATCH] test(cost): type the web search cost helpers and cover OpenAI-shaped tool_usage --- .../test_tool_call_cost_tracking.py | 59 ++++++++++--------- 1 file changed, 32 insertions(+), 27 deletions(-) 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 0d75301a57a..6cc3dcceebc 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 @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping, Sequence import pytest @@ -6,7 +7,7 @@ import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.types.llms.openai import FileSearchTool, WebSearchOptions +from litellm.types.llms.openai import FileSearchTool, ResponsesAPIResponse, WebSearchOptions from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams @@ -941,9 +942,9 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 -def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None): - from litellm.types.llms.openai import ResponsesAPIResponse - +def _responses_with_web_search( + model: str, actions: Sequence[Mapping[str, str]], tool_usage: Mapping[str, object] | None = None +) -> ResponsesAPIResponse: payload = { "id": "resp_1", "created_at": 1756900000, @@ -960,26 +961,21 @@ def _bedrock_mantle_responses_with_web_search(model, actions, tool_usage=None): ) -def _bedrock_mantle_web_search_cost(model, response): +def _web_search_cost(model: str, response: ResponsesAPIResponse, custom_llm_provider: str) -> float: from litellm.types.utils import Usage return StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( model=model, response_object=response, usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - custom_llm_provider="bedrock_mantle", + custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=None, ) @pytest.mark.parametrize("model", _BEDROCK_MANTLE_WEB_SEARCH_MODELS) def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model): - """ - Regression for LIT-6870: the bedrock_mantle GPT ids forward the web_search tool but carried no - search_context_cost_per_query, so every Bedrock web search (billed at $12 per 1,000 queries) - was costed at $0. Two reported queries must bill 2 x $0.012, whichever model prefix shape the - cost path resolves the deployment under. - """ + """Two Bedrock-reported web searches bill 2 x $0.012 under the prefixed and the bare model id alike.""" pricing = litellm.get_model_info(model)["search_context_cost_per_query"] assert pricing == { "search_context_size_low": _BEDROCK_MANTLE_WEB_SEARCH_RATE, @@ -987,13 +983,13 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model) "search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE, } - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], tool_usage={"web_search": {"num_requests": 2}}, ) for cost_model in (model, model.split("/", 1)[1]): - cost = _bedrock_mantle_web_search_cost(cost_model, response) + cost = _web_search_cost(cost_model, response, "bedrock_mantle") assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"{cost_model} must bill 2 x ${_BEDROCK_MANTLE_WEB_SEARCH_RATE} for 2 web searches, got ${cost}" ) @@ -1001,13 +997,9 @@ def test_bedrock_mantle_web_search_billed_per_query(local_model_cost_map, model) @pytest.mark.parametrize("num_requests", [1, 0]) def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests): - """ - Regression for LIT-6870: Bedrock bills one query per search and reports the billable count as - tool_usage.web_search.num_requests, while its open_page fetches share the web_search_call item - type. A search plus an open_page must bill the reported count, never the two items. - """ + """A search plus an open_page fetch bills tool_usage.web_search.num_requests, never the two items.""" model = "bedrock_mantle/openai.gpt-5.6-sol" - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[ {"type": "search", "query": "litellm"}, @@ -1016,7 +1008,7 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode tool_usage={"web_search": {"num_requests": num_requests}}, ) - cost = _bedrock_mantle_web_search_cost(model, response) + cost = _web_search_cost(model, response, "bedrock_mantle") assert cost == pytest.approx(num_requests * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"{num_requests} reported web search requests must bill {num_requests} x " @@ -1029,20 +1021,33 @@ def test_web_search_call_count_prefers_provider_reported_num_requests(local_mode [None, {}, {"web_search": None}, {"web_search": {"num_requests": "many"}}, {"web_search": {"num_requests": -1}}], ) def test_web_search_call_count_falls_back_to_items_without_reported_count(local_model_cost_map, tool_usage): - """ - Without a usable reported count (no tool_usage, no web_search block, or a malformed one) the - per-call path must keep counting web_search_call items instead of raising or billing zero. - """ + """Without a usable reported count the per-call path keeps counting web_search_call items.""" model = "bedrock_mantle/openai.gpt-5.6-sol" - response = _bedrock_mantle_responses_with_web_search( + response = _responses_with_web_search( model, actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}], tool_usage=tool_usage, ) - cost = _bedrock_mantle_web_search_cost(model, response) + cost = _web_search_cost(model, response, "bedrock_mantle") assert cost == pytest.approx(2 * _BEDROCK_MANTLE_WEB_SEARCH_RATE), ( f"2 web_search_call items with tool_usage={tool_usage!r} must bill 2 x " f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}" ) + + +def test_web_search_call_count_reads_reported_count_beside_other_tool_usage_entries(local_model_cost_map): + """OpenAI reports web_search.num_requests next to other tool entries, which must not disable the reported count.""" + response = _responses_with_web_search( + "gpt-5.6", + actions=[{"type": "search", "query": "S&P 500 close"}, {"type": "open_page", "url": "https://example.com/"}], + tool_usage={ + "image_gen": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, + "web_search": {"num_requests": 1}, + }, + ) + + cost = _web_search_cost("gpt-5.6", response, "openai") + + assert cost == pytest.approx(0.01), f"1 reported OpenAI web search must bill 1 x $0.01, not the 2 items, got ${cost}"