mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
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:
commit
cff2fd4f80
5 changed files with 213 additions and 9 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 (
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue