scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

test(cost): type the web search cost helpers and cover OpenAI-shaped tool_usage

Browse source
This commit is contained in:
mateo-berri 2026-09-03 12:34:54 -07:00
parent 51d821ae45
commit 339da4183d
1 changed files with 32 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py
View file

@ -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}"