mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge 53d7e124b7 into 11a02b9581
This commit is contained in:
commit
b75b9e18cd
9 changed files with 640 additions and 46 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue