This commit is contained in:
Traci Lim 2026-09-03 23:08:33 +08:00 committed by GitHub
commit b75b9e18cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 640 additions and 46 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 (
@ -175,16 +178,36 @@ class StandardBuiltInToolCostTracking:
)
return per_call_cost * StandardBuiltInToolCostTracking._count_web_search_calls(response_object)
@staticmethod
def _reported_web_search_requests(response_object: object) -> int | None:
"""The provider's own billable search count off a Responses payload, when it reports one.
Bedrock returns ``tool_usage.web_search.num_requests``, which counts only the searches it
charges for and excludes the cached-page fetches that share the ``web_search_call`` item
type. Counting items instead overcharges every request that opened a page.
"""
tool_usage: Final = getattr(response_object, "tool_usage", None)
if tool_usage is None:
return None
try:
return ResponsesToolUsage.model_validate(tool_usage).web_search.num_requests
except ValidationError:
return None
@staticmethod
def _count_web_search_calls(response_object: object) -> int:
"""
Number of web searches to bill for on the per-call pricing path.
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.
get_cost_for_web_search_request and never reach here. Of the rest, some report the count on the
response itself and are believed over the item count; otherwise 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.
"""
reported: Final = StandardBuiltInToolCostTracking._reported_web_search_requests(response_object)
if reported is not None:
return reported
if isinstance(response_object, ResponsesAPIResponse):
count = sum(
1 for output_item in response_object.output if _output_item_type(output_item) == "web_search_call"

View file

@ -16,7 +16,7 @@ BaseAWSLLM._sign_request after the request body is finalized.
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Any, Final
from typing_extensions import ReadOnly, TypedDict
@ -48,7 +48,18 @@ _BASE_SUFFIXES_TO_STRIP: Final = (
)
# Per Bedrock Mantle Responses API validation errors.
_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES = frozenset({"function", "mcp", "custom", "namespace", "tool_search"})
_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
{"function", "mcp", "custom", "namespace", "tool_search"}
)
# Enabled per model, not per provider: models AWS has not enabled reject the whole request
# with "Tool type 'web_search' is not supported for model `<id>`", so the cost map's
# supports_web_search flag decides rather than a provider-wide allowlist entry.
_BEDROCK_MANTLE_WEB_SEARCH_TOOL_TYPE: Final = "web_search"
_BEDROCK_MANTLE_RESPONSE_TOOL_TYPES_WITH_WEB_SEARCH: Final = _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES | frozenset(
(_BEDROCK_MANTLE_WEB_SEARCH_TOOL_TYPE,)
)
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
@ -131,9 +142,20 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
def supports_native_websocket(self) -> bool:
return False
@staticmethod
def _filter_unsupported_tools(tools: list[Any]) -> list[Any]:
"""Keep only tool types Mantle's Responses API accepts."""
def _supported_response_tool_types(self, tools: "Sequence[Any]", model: str) -> frozenset[str]:
"""The tool types `tools` may keep: the provider-wide set, widened by `web_search`
only when this request asks for it and the model is one AWS enabled it for."""
if not any(
isinstance(tool, dict) and tool.get("type") == _BEDROCK_MANTLE_WEB_SEARCH_TOOL_TYPE for tool in tools
):
return _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES
if litellm.supports_web_search(model=model, custom_llm_provider=self.custom_llm_provider.value):
return _BEDROCK_MANTLE_RESPONSE_TOOL_TYPES_WITH_WEB_SEARCH
return _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES
def _filter_unsupported_tools(self, tools: "Sequence[Any]", model: str) -> list[Any]:
"""Keep only tool types Mantle's Responses API accepts for this model."""
supported_tool_types: Final = self._supported_response_tool_types(tools=tools, model=model)
kept: Final[list[Any]] = []
dropped_types: Final[list[str]] = []
for tool in tools:
@ -141,16 +163,17 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
kept.append(tool)
continue
tool_type = tool.get("type")
if tool_type in _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES:
if tool_type in supported_tool_types:
kept.append(tool)
else:
dropped_types.append(str(tool_type))
if dropped_types:
verbose_logger.warning(
"Bedrock Mantle Responses API: dropping unsupported tool type(s) %s (supported: %s).",
"Bedrock Mantle Responses API: dropping unsupported tool type(s) %s for model %s (supported: %s).",
sorted(set(dropped_types)),
sorted(_BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES),
model,
sorted(supported_tool_types),
)
return kept
@ -185,7 +208,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input=input, model=model)
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
request_params: Final = (
{
@ -215,10 +238,10 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
tools: Final = item.get("tools")
return tools if isinstance(tools, list) else []
@classmethod
def _hoist_codex_additional_tools(
cls,
self,
input: "str | ResponseInputParam",
model: str,
) -> "tuple[str | ResponseInputParam, list[Any]]":
"""Codex's "responses lite" wire mode ships tool definitions inside
`input` as {"type": "additional_tools", "role": "developer",
@ -229,18 +252,18 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
"""
if not isinstance(input, list):
return input, []
additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)]
additional_tools_items: Final = [item for item in input if self._is_codex_additional_tools_item(item)]
if not additional_tools_items:
return input, []
remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)]
hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)]
remaining_input: Final = [item for item in input if not self._is_codex_additional_tools_item(item)]
hoisted_tools = [tool for item in additional_tools_items for tool in self._tools_of_additional_tools_item(item)]
verbose_logger.debug(
"Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) "
"into the top-level tools param (Mantle rejects that input item type).",
len(hoisted_tools),
len(additional_tools_items),
)
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
return remaining_input, self._filter_unsupported_tools(tools=hoisted_tools, model=model)
@staticmethod
def _agent_message_text(item: "Mapping[str, object]") -> str:
@ -347,7 +370,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
return params
tools_list: Final = tools if isinstance(tools, list) else [tools]
filtered: Final = self._filter_unsupported_tools(tools_list)
filtered: Final = self._filter_unsupported_tools(tools=tools_list, model=model)
if filtered:
params["tools"] = filtered
else:

View file

@ -262,18 +262,45 @@ class OpenAIResponsesHandler(BaseTranslation):
return data
@staticmethod
def _request_tool_name(tool: object) -> str | None:
"""The name a Responses request tool acts under: ``name`` for function and custom,
``server_label`` for mcp, and the bare ``type`` for built-in server-side tools
(web_search, code_interpreter, ...) that carry no name of their own. Those bill and
reach the internet, so allowlist checks must see them under some name."""
if not isinstance(tool, dict):
return None
tool_type: Final = tool.get("type")
if tool_type in ("function", "custom"):
name: Final = tool.get("name")
return str(name) if name else None
if tool_type == "mcp":
server_label: Final = tool.get("server_label")
return str(server_label) if server_label else None
return str(tool_type) if tool_type else None
@staticmethod
def _tools_nested_in_input_item(item: object) -> tuple[object, ...]:
"""Tools declared inside an ``additional_tools`` input item. Codex's responses-lite wire
mode ships tool definitions there instead of in top-level ``tools``, and providers hoist
them back out before dispatch, so reading only ``tools`` would miss them."""
if not isinstance(item, dict) or item.get("type") != "additional_tools":
return ()
tools: Final = item.get("tools")
return tuple(tools) if isinstance(tools, list) else ()
def extract_request_tool_names(self, data: dict) -> list[str]:
"""Extract tool names from Responses API request (tools[].name for function
and custom, tools[].server_label for mcp)."""
names: Final[list[str]] = []
for tool in data.get("tools") or []:
if not isinstance(tool, dict):
continue
if tool.get("type") in ("function", "custom") and tool.get("name"):
names.append(str(tool["name"]))
elif tool.get("type") == "mcp" and tool.get("server_label"):
names.append(str(tool["server_label"]))
return names
input_items: Final = data.get("input")
nested: Final = (
tuple(tool for item in input_items for tool in self._tools_nested_in_input_item(item))
if isinstance(input_items, list)
else ()
)
return [
name
for tool in (*(data.get("tools") or []), *nested)
if (name := self._request_tool_name(tool)) is not None
]
def _apply_guardrailed_tools_to_data(
self,

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,
@ -52933,7 +52938,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.2e-06,
@ -52944,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,
@ -52966,7 +52977,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-cyber": {
"input_cost_per_token": 1.375e-05,
@ -53005,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,
@ -53027,7 +53044,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"us.openai.gpt-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -53192,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,
@ -53213,7 +53236,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.4": {
"input_cost_per_token": 2.75e-06,
@ -53222,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,
@ -53243,7 +53272,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,

View file

@ -1328,6 +1328,26 @@ One of: completed, failed, in_progress, cancelled, queued, or incomplete.
"""
class WebSearchToolUsage(BaseModel):
num_requests: int
model_config = ConfigDict(extra="allow")
class ResponsesToolUsage(BaseModel):
"""A Responses payload's own `tool_usage`, used to bill server-side tools off the provider's
count rather than off the returned items. Bedrock populates `web_search.num_requests`.
Deliberately not a declared field on ``ResponsesAPIResponse``: it arrives as an extra and is
validated only where the cost path reads it, so a payload reporting a tool we do not model, or
a shape we do not expect, still parses instead of failing the whole response.
"""
web_search: WebSearchToolUsage
model_config = ConfigDict(extra="allow")
class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject):
id: str
created_at: int

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,
@ -52933,7 +52938,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-terra": {
"input_cost_per_token": 2.2e-06,
@ -52944,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,
@ -52966,7 +52977,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.6-cyber": {
"input_cost_per_token": 1.375e-05,
@ -53005,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,
@ -53027,7 +53044,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"us.openai.gpt-5.6-sol": {
"input_cost_per_token": 4.4e-06,
@ -53192,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,
@ -53213,7 +53236,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/openai.gpt-5.4": {
"input_cost_per_token": 2.75e-06,
@ -53222,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,
@ -53243,7 +53272,8 @@
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
"supports_vision": true,
"supports_web_search": true
},
"bedrock_mantle/google.gemma-4-31b": {
"input_cost_per_token": 1.4e-07,

View file

@ -709,6 +709,185 @@ def _openai_responses_with_web_search_calls(model, num_calls):
)
def _responses_with_reported_web_search(model, search_items, fetch_items, tool_usage):
"""A Bedrock-shaped Responses payload: `web_search_call` items for both operations, plus the
raw `tool_usage` block, which Bedrock populates with the count it charges for."""
from litellm.types.llms.openai import ResponsesAPIResponse
output = [
{
"id": f"ws_s{i}",
"type": "web_search_call",
"status": "completed",
"action": {"type": "search", "queries": ["a", "b", "c"]},
}
for i in range(search_items)
] + [
{
"id": f"ws_f{i}",
"type": "web_search_call",
"status": "completed",
"action": {"type": "open_page", "url": "https://example.invalid/page"},
}
for i in range(fetch_items)
]
return ResponsesAPIResponse(
id="resp_1",
created_at=0,
model=model,
object="response",
output=output,
parallel_tool_calls=False,
tool_choice="auto",
tools=[],
tool_usage=tool_usage,
)
def test_bedrock_mantle_web_search_bills_the_count_bedrock_reports(local_model_cost_map):
"""Bedrock reports its billable search count as ``tool_usage.web_search.num_requests`` and it
excludes cached-page fetches, which arrive as ``web_search_call`` items indistinguishable from
searches by item type. Measured against bedrock-mantle.us-east-1.api.aws: a response with five
search items and one ``open_page`` item reports 5, and one with one search and one
``open_page`` reports 1. Counting items would bill 6 and 2, overcharging every request that
opened a page. Three queries inside a single search item still count as one request."""
from litellm.types.utils import Usage
model = "bedrock_mantle/openai.gpt-5.6-terra"
per_query = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"]
usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
cases = (
{"searches": 5, "fetches": 1, "reported": 5},
{"searches": 1, "fetches": 1, "reported": 1},
{"searches": 2, "fetches": 1, "reported": 2},
# A follow-up turn can fetch a cached page without searching, and Bedrock then charges
# nothing. The item path floors at one search; a reported count must not be floored.
{"searches": 0, "fetches": 2, "reported": 0},
)
for case in cases:
response = _responses_with_reported_web_search(
model,
search_items=case["searches"],
fetch_items=case["fetches"],
tool_usage={"web_search": {"num_requests": case["reported"]}},
)
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
response_object=response,
usage=usage,
custom_llm_provider="bedrock_mantle",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(case["reported"] * per_query), (
f"{case['searches'] + case['fetches']} items reporting {case['reported']} requests must "
f"bill {case['reported']} x ${per_query}, got ${cost}"
)
def test_web_search_falls_back_to_counting_items_when_no_count_is_reported(local_model_cost_map):
"""Providers that report nothing must keep the item-count behaviour, so adding the Bedrock path
cannot change what OpenAI and Azure bill."""
model = "gpt-5-nano"
response = _openai_responses_with_web_search_calls(model, num_calls=3)
assert not hasattr(response, "tool_usage")
assert StandardBuiltInToolCostTracking._count_web_search_calls(response) == 3
def test_streamed_web_search_bills_the_reported_count(local_model_cost_map):
"""On the streaming path the cost calculator is handed the response unwrapped out of the
terminal `ResponseCompletedEvent`, not the event. Pin that: the reported count has to survive
the unwrap, because the event itself carries no output items and would floor the bill at one
search instead of charging what Bedrock reported."""
import datetime
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import ResponseCompletedEvent
from litellm.types.utils import Usage
model = "bedrock_mantle/openai.gpt-5.6-terra"
per_query = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"]
response = _responses_with_reported_web_search(
model, search_items=2, fetch_items=3, tool_usage={"web_search": {"num_requests": 2}}
)
logging_obj = LiteLLMLoggingObj(
model=model, messages=[], stream=True, call_type="aresponses",
start_time=0, litellm_call_id="1", function_id="1",
)
now = datetime.datetime.now()
assembled = logging_obj._get_assembled_streaming_response(
result=ResponseCompletedEvent(type="response.completed", response=response),
start_time=now, end_time=now, is_async=True, streaming_chunks=[],
)
assert assembled is response, "the completed event must be unwrapped to the response itself"
cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools(
model=model,
response_object=assembled,
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
custom_llm_provider="bedrock_mantle",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(2 * per_query), (
f"5 items reporting 2 requests must bill 2 x ${per_query} after the stream is assembled, got ${cost}"
)
@pytest.mark.parametrize(
"tool_usage",
[
{"web_search": {"num_requests": "not-a-number"}},
{"unexpected_shape": True},
"not-an-object",
None,
],
)
def test_unreadable_reported_web_search_count_falls_back_to_items(tool_usage, local_model_cost_map):
"""A tool_usage block we cannot read must neither raise nor bill zero: fall back to the items.
Reading it defensively here, rather than declaring it on the response model, is what keeps a
payload that reports an unmodelled tool from failing the whole response."""
response = _responses_with_reported_web_search(
"bedrock_mantle/openai.gpt-5.6-terra", search_items=2, fetch_items=1, tool_usage=tool_usage
)
assert StandardBuiltInToolCostTracking._count_web_search_calls(response) == 3
@pytest.mark.parametrize(
"model",
[
"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",
],
)
def test_bedrock_mantle_native_web_search_priced_per_query(model, local_model_cost_map):
"""Bedrock's server-side Web Search bills $12.00 per 1,000 queries, the same in every Region
that offers it (AWS Pricing API usage type ``Bedrock-Websearch-Queries``, effective
2026-08-01). Bedrock counts each search rather than each request, so a Responses output
carrying three ``web_search_call`` items must bill three queries. Without
``search_context_cost_per_query`` the default fallback silently returns $0."""
from litellm.types.utils import Usage
per_query = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"]
assert per_query == 0.012
for num_calls in (1, 3):
response = _openai_responses_with_web_search_calls(model, num_calls=num_calls)
cost = 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",
standard_built_in_tools_params=None,
)
assert cost == pytest.approx(num_calls * per_query), (
f"{model} must bill {num_calls} x ${per_query} for {num_calls} web search(es), got ${cost}"
)
def test_openai_responses_web_search_priced_per_call(local_model_cost_map):
"""
Regression for LIT-5013 bug 1: OpenAI reasoning models (gpt-5 family, o-series, deep-research)

View file

@ -11,6 +11,7 @@ import copy
import json
import logging
from pathlib import Path
from typing import Final
import pytest
from botocore.exceptions import (
@ -332,7 +333,7 @@ class TestBedrockMantleResponsesTools:
params = cfg.map_openai_params(
response_api_optional_params={
"tools": [
{"type": "web_search"},
{"type": "file_search"},
{"type": "function", "name": "exec_command"},
]
},
@ -344,7 +345,7 @@ class TestBedrockMantleResponsesTools:
def test_map_openai_params_removes_tools_when_all_unsupported(self):
cfg = BedrockMantleResponsesAPIConfig()
params = cfg.map_openai_params(
response_api_optional_params={"tools": [{"type": "web_search"}]},
response_api_optional_params={"tools": [{"type": "file_search"}]},
model="openai.gpt-5.5",
drop_params=False,
)
@ -358,12 +359,167 @@ class TestBedrockMantleResponsesTools:
"litellm.llms.bedrock_mantle.responses.transformation.verbose_logger.warning"
) as mock_warning:
cfg.map_openai_params(
response_api_optional_params={"tools": [{"type": "web_search"}]},
response_api_optional_params={"tools": [{"type": "file_search"}]},
model="openai.gpt-5.5",
drop_params=False,
)
assert mock_warning.call_count == 1
assert "web_search" in str(mock_warning.call_args)
assert "file_search" in str(mock_warning.call_args)
def _web_search_tool(external_web_access: bool = False) -> dict[str, object]:
return {"type": "web_search", "external_web_access": external_web_access}
def _function_tool() -> dict[str, object]:
return {"type": "function", "name": "exec_command"}
def _codex_additional_tools_input(tools: list[dict[str, object]]) -> list[dict[str, object]]:
return [
{"type": "additional_tools", "role": "developer", "tools": tools},
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
]
_WEB_SEARCH_CAPABLE_MODELS: Final = (
"openai.gpt-5.6-sol",
"openai.gpt-5.6-terra",
"openai.gpt-5.6-luna",
"openai.gpt-5.5",
"openai.gpt-5.4",
)
class TestBedrockMantleResponsesNativeWebSearch:
"""Verified against bedrock-mantle.us-east-1.api.aws: every id in
_WEB_SEARCH_CAPABLE_MODELS returns `web_search_call` items and `url_citation`
annotations, while google.gemma-4-31b answers "Tool type 'web_search' is not
supported for model `google.gemma-4-31b`"."""
@pytest.mark.parametrize("model", [*_WEB_SEARCH_CAPABLE_MODELS, "bedrock_mantle/openai.gpt-5.6-terra"])
def test_web_search_survives_for_capable_models(self, model, local_model_cost_map):
cfg = BedrockMantleResponsesAPIConfig()
params = cfg.map_openai_params(
response_api_optional_params={"tools": [_web_search_tool(), _function_tool()]},
model=model,
drop_params=False,
)
assert params["tools"] == [_web_search_tool(), _function_tool()]
def test_external_web_access_true_is_forwarded_verbatim(self, local_model_cost_map):
cfg = BedrockMantleResponsesAPIConfig()
params = cfg.map_openai_params(
response_api_optional_params={"tools": [_web_search_tool(external_web_access=True)]},
model="openai.gpt-5.6-terra",
drop_params=False,
)
assert params["tools"] == [_web_search_tool(external_web_access=True)]
@pytest.mark.parametrize("model", ["google.gemma-4-31b", "openai.gpt-oss-120b"])
def test_web_search_still_dropped_for_models_without_the_capability(self, model, local_model_cost_map):
cfg = BedrockMantleResponsesAPIConfig()
params = cfg.map_openai_params(
response_api_optional_params={"tools": [_web_search_tool(), _function_tool()]},
model=model,
drop_params=False,
)
assert params["tools"] == [_function_tool()]
def test_web_search_is_not_advertised_unless_the_request_asks_for_it(self, local_model_cost_map):
"""The tool type set widens only for requests that carry a Web Search tool, so a capable
model's other requests keep the provider-wide set and never reach the cost map."""
cfg = BedrockMantleResponsesAPIConfig()
supported = cfg._supported_response_tool_types(tools=[_function_tool()], model="openai.gpt-5.6-terra")
assert "web_search" not in supported, "a request carrying no web_search tool must not widen the set"
assert "function" in supported
def test_web_search_hoisted_out_of_codex_additional_tools(self, local_model_cost_map):
cfg = BedrockMantleResponsesAPIConfig()
body = cfg.transform_responses_api_request(
model="openai.gpt-5.6-terra",
input=_codex_additional_tools_input([_web_search_tool()]),
response_api_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert body["tools"] == [_web_search_tool()]
def test_web_search_hoisted_from_codex_is_dropped_for_incapable_model(self, local_model_cost_map):
cfg = BedrockMantleResponsesAPIConfig()
body = cfg.transform_responses_api_request(
model="google.gemma-4-31b",
input=_codex_additional_tools_input([_web_search_tool()]),
response_api_optional_request_params={},
litellm_params=GenericLiteLLMParams(),
headers={},
)
assert "tools" not in body
def test_cost_map_advertises_web_search(self, local_model_cost_map):
assert all(
litellm.supports_web_search(model=f"bedrock_mantle/{model}", custom_llm_provider="bedrock_mantle")
for model in _WEB_SEARCH_CAPABLE_MODELS
)
def test_search_results_and_citations_survive_the_response_transform(self):
"""Mantle returns Web Search results in the OpenAI Responses shape, so the config adds no
response-side handling. Lock that: a `web_search_call` item and the `url_citation`
annotations that make the answer attributable must both reach the caller intact."""
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
citation: Final = {
"type": "url_citation",
"title": "Web Search - Amazon Bedrock",
"url": "https://docs.aws.amazon.com/bedrock/latest/userguide/web-search.html",
"start_index": 0,
"end_index": 12,
}
upstream: Final = {
"id": "resp_1",
"created_at": 0,
"model": "openai.gpt-5.6-terra",
"object": "response",
"output": [
{
"id": "ws_1",
"type": "web_search_call",
"status": "completed",
"action": {"type": "search", "queries": ["bedrock web search regions"]},
},
{
"id": "msg_1",
"type": "message",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "Three Regions", "annotations": [citation]}],
},
],
"parallel_tool_calls": False,
"tool_choice": "auto",
"tools": [],
}
cfg = BedrockMantleResponsesAPIConfig()
out = cfg.transform_response_api_response(
model="openai.gpt-5.6-terra",
raw_response=httpx.Response(200, json=upstream),
logging_obj=LiteLLMLoggingObj(
model="openai.gpt-5.6-terra",
messages=[],
stream=False,
call_type="aresponses",
start_time=0,
litellm_call_id="1",
function_id="1",
),
)
dumped = out.model_dump()
assert [item["type"] for item in dumped["output"]] == ["web_search_call", "message"]
assert dumped["output"][1]["content"][0]["annotations"] == [citation]
def _codex_exec_tool():
@ -551,7 +707,7 @@ class TestBedrockMantleCodexAdditionalTools:
"type": "additional_tools",
"role": "developer",
"tools": [
{"type": "web_search"},
{"type": "file_search"},
{"type": "function", "name": "wait"},
],
},
@ -563,7 +719,7 @@ class TestBedrockMantleCodexAdditionalTools:
def test_item_stripped_even_when_no_hoisted_tool_survives(self):
body = self._transform(
input=[
{"type": "additional_tools", "role": "developer", "tools": [{"type": "web_search"}]},
{"type": "additional_tools", "role": "developer", "tools": [{"type": "file_search"}]},
self._USER_MESSAGE,
]
)

View file

@ -86,6 +86,66 @@ class TestExtractRequestToolNames:
"get_current_weather",
]
def test_openai_responses_builtin_tools(self):
"""Built-in server-side tools carry no name of their own, so they act under their
type; without that a restricted key could still reach the internet and bill for
web_search while its allowlist named nothing of the sort (VERIA finding on PR #37995)."""
data = {
"tools": [
{"type": "web_search"},
{"type": "code_interpreter", "container": {"type": "auto"}},
{"type": "function", "name": "get_current_weather"},
{"type": "mcp", "server_label": "dmcp", "server_url": "http://x"},
]
}
assert extract_request_tool_names("/v1/responses", data) == [
"web_search",
"code_interpreter",
"get_current_weather",
"dmcp",
]
def test_openai_responses_tools_nested_in_additional_tools_input_item(self):
"""Codex's responses-lite wire mode declares tools inside an `additional_tools` input
item, and providers hoist them into top-level `tools` before dispatch. Reading only
`tools` would let a restricted key smuggle any tool through `input`
(VERIA finding on PR #37995)."""
data = {
"input": [
{"type": "message", "role": "user", "content": "hi"},
{
"type": "additional_tools",
"role": "developer",
"tools": [{"type": "web_search"}, {"type": "function", "name": "run_sql"}],
},
],
"tools": [{"type": "function", "name": "declared_up_front"}],
}
assert extract_request_tool_names("/v1/responses", data) == [
"declared_up_front",
"web_search",
"run_sql",
]
def test_openai_responses_malformed_additional_tools_yields_no_name(self):
"""An `additional_tools` item with a missing or non-list `tools` slot, and a plain string
input, must not raise on the auth hot path."""
data = {
"input": [
{"type": "additional_tools"},
{"type": "additional_tools", "tools": "not-a-list"},
"junk",
]
}
assert extract_request_tool_names("/v1/responses", data) == []
assert extract_request_tool_names("/v1/responses", {"input": "plain string"}) == []
def test_openai_responses_unnamed_tool_yields_no_name(self):
"""A function or custom tool missing its name must not fall back to the bare type:
that would let "function" satisfy an allowlist that never granted the real tool."""
data = {"tools": [{"type": "function"}, {"type": "custom", "name": ""}, {"type": "mcp"}, "junk"]}
assert extract_request_tool_names("/v1/responses", data) == []
def test_anthropic_tools(self):
data = {"tools": [{"name": "get_weather"}, {"name": "run_sql"}]}
assert extract_request_tool_names("/v1/messages", data) == [
@ -227,6 +287,52 @@ class TestCheckToolsAllowlist:
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
assert "restricted_tool" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_disallowed_builtin_web_search_raises_on_responses_route(self):
token = _token(metadata={"allowed_tools": ["run_sql"]})
body = {"tools": [{"type": "web_search"}]}
with pytest.raises(ProxyException) as exc_info:
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/responses",
)
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
assert "web_search" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_allowlisted_builtin_web_search_passes_on_responses_route(self):
token = _token(metadata={"allowed_tools": ["web_search", "run_sql"]})
tools = [{"type": "web_search"}, {"type": "function", "name": "run_sql"}]
body = {"tools": tools}
assert extract_request_tool_names("/v1/responses", body) == ["web_search", "run_sql"]
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/responses",
)
assert body["tools"] == tools
@pytest.mark.asyncio
async def test_disallowed_tool_nested_in_input_raises_on_responses_route(self):
token = _token(metadata={"allowed_tools": ["run_sql"]})
body = {
"input": [
{"type": "additional_tools", "role": "developer", "tools": [{"type": "web_search"}]},
]
}
with pytest.raises(ProxyException) as exc_info:
await check_tools_allowlist(
request_body=body,
valid_token=token,
team_object=None,
route="/v1/responses",
)
assert exc_info.value.type == ProxyErrorTypes.tool_access_denied
assert "web_search" in str(exc_info.value.message)
@pytest.mark.asyncio
async def test_team_allowlist_used_when_key_empty(self):
token = _token(