Merge pull request #39610 from BerriAI/litellm_bedrock_mantle_web_search_cost

fix(cost): bill bedrock_mantle web search at $12 per 1k queries using Bedrock's reported count
This commit is contained in:
Mateo Wang 2026-09-03 13:14:20 -07:00 committed by GitHub
commit cff2fd4f80
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 213 additions and 9 deletions

View file

@ -5,6 +5,8 @@ Helper utilities for tracking the cost of built-in tools.
from collections.abc import Mapping
from typing import Final, Literal
from pydantic import ValidationError
import litellm
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
from litellm.litellm_core_utils.llm_cost_calc.utils import (
@ -13,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
from litellm.types.llms.openai import (
FileSearchTool,
ResponsesAPIResponse,
ResponsesToolUsage,
WebSearchOptions,
)
from litellm.types.utils import (
@ -32,6 +35,17 @@ def _output_item_type(output_item: object) -> str | None:
return item_type if isinstance(item_type, str) else None
def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None:
tool_usage: Final = getattr(response_object, "tool_usage", None)
if tool_usage is None:
return None
try:
web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search
except ValidationError:
return None
return None if web_search is None else web_search.num_requests
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):
@ -182,15 +196,19 @@ class StandardBuiltInToolCostTracking:
Providers that report a request count in usage (gemini, anthropic, xai, vertex) are handled by
get_cost_for_web_search_request and never reach here. This path prices per call, so it must count
the web_search_call items. Chat-completions responses only expose url_citation annotations with no
count, so they floor to a single billable search.
the web_search_call items, unless the response reports the billable count itself
(Bedrock's tool_usage.web_search.num_requests, which excludes open_page fetches). Chat-completions
responses only expose url_citation annotations with no count, so they floor to a single billable search.
"""
if isinstance(response_object, ResponsesAPIResponse):
count = sum(
1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
)
return max(count, 1)
return 1
if not isinstance(response_object, ResponsesAPIResponse):
return 1
reported: Final = _reported_web_search_requests(response_object)
if reported is not None:
return reported
count: Final = sum(
1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"
)
return max(count, 1)
@staticmethod
def _handle_file_search_cost(

View file

@ -52911,6 +52911,11 @@
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -52945,6 +52950,11 @@
"cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53007,6 +53017,11 @@
"cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53195,6 +53210,11 @@
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53226,6 +53246,11 @@
"cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
"output_cost_per_token": 1.65e-05,
"output_cost_per_token_above_272k_tokens": 2.475e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,

View file

@ -66,6 +66,7 @@ from pydantic import (
ConfigDict,
Discriminator,
Field,
NonNegativeInt,
PrivateAttr,
SerializerFunctionWrapHandler,
field_serializer,
@ -1321,6 +1322,18 @@ class ResponseAPIUsage(BaseLiteLLMOpenAIResponseObject):
model_config = {"extra": "allow"}
class WebSearchToolUsage(BaseModel):
model_config = ConfigDict(frozen=True)
num_requests: NonNegativeInt
class ResponsesToolUsage(BaseModel):
model_config = ConfigDict(frozen=True)
web_search: WebSearchToolUsage | None = None
ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"]
"""
The status of the response generation.

View file

@ -52911,6 +52911,11 @@
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -52945,6 +52950,11 @@
"cache_read_input_token_cost_above_272k_tokens": 4.4e-07,
"output_cost_per_token": 1.32e-05,
"output_cost_per_token_above_272k_tokens": 1.98e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53007,6 +53017,11 @@
"cache_read_input_token_cost_above_272k_tokens": 4.4e-08,
"output_cost_per_token": 1.32e-06,
"output_cost_per_token_above_272k_tokens": 1.98e-06,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53195,6 +53210,11 @@
"cache_read_input_token_cost_above_272k_tokens": 1.1e-06,
"output_cost_per_token": 3.3e-05,
"output_cost_per_token_above_272k_tokens": 4.95e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,
@ -53226,6 +53246,11 @@
"cache_read_input_token_cost_above_272k_tokens": 5.5e-07,
"output_cost_per_token": 1.65e-05,
"output_cost_per_token_above_272k_tokens": 2.475e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.012,
"search_context_size_low": 0.012,
"search_context_size_medium": 0.012
},
"litellm_provider": "bedrock_mantle",
"max_input_tokens": 1050000,
"max_output_tokens": 128000,

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
@ -928,3 +929,125 @@ def test_web_search_gate_reads_server_side_tool_usage_details_without_citations(
standard_built_in_tools_params=None,
)
assert cost == 3 * _DEFAULT_WEB_SEARCH_COST_PER_CALL
_BEDROCK_MANTLE_WEB_SEARCH_MODELS = (
"bedrock_mantle/openai.gpt-5.6-sol",
"bedrock_mantle/openai.gpt-5.6-terra",
"bedrock_mantle/openai.gpt-5.6-luna",
"bedrock_mantle/openai.gpt-5.5",
"bedrock_mantle/openai.gpt-5.4",
)
_BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012
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,
"model": model.split("/", 1)[-1],
"object": "response",
"status": "completed",
"output": [
{"type": "web_search_call", "id": f"ws_{i}", "status": "completed", "action": action}
for i, action in enumerate(actions)
],
}
return ResponsesAPIResponse.model_validate(
payload if tool_usage is None else {**payload, "tool_usage": tool_usage}
)
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=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):
"""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,
"search_context_size_medium": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
"search_context_size_high": _BEDROCK_MANTLE_WEB_SEARCH_RATE,
}
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 = _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}"
)
@pytest.mark.parametrize("num_requests", [1, 0])
def test_web_search_call_count_prefers_provider_reported_num_requests(local_model_cost_map, num_requests):
"""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 = _responses_with_web_search(
model,
actions=[
{"type": "search", "query": "litellm"},
{"type": "open_page", "url": "https://docs.litellm.ai/"},
],
tool_usage={"web_search": {"num_requests": num_requests}},
)
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 "
f"${_BEDROCK_MANTLE_WEB_SEARCH_RATE}, got ${cost}"
)
@pytest.mark.parametrize(
"tool_usage",
[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 the per-call path keeps counting web_search_call items."""
model = "bedrock_mantle/openai.gpt-5.6-sol"
response = _responses_with_web_search(
model,
actions=[{"type": "search", "query": "litellm"}, {"type": "search", "query": "bedrock web search"}],
tool_usage=tool_usage,
)
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}"