fix(cost): price AgentCore web search at $7 per 1,000 queries (#46571)

* fix(cost): price AgentCore web search at $7 per 1,000 queries

* feat(search): honor input_cost_per_query set on a search tool (#46560)

* feat(search): honor input_cost_per_query set on a search tool

* feat(search): record the search tool price on proxy-created logging objects

* chore(search): type the search tool price lookups and shard tests/unit/search

* chore(openapi): regenerate snapshot and ui types for search tool input_cost_per_query

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Kerry Lu <klu@berri.ai>

.github/workflows/test-unit.yml

@@ -417,6 +417,7 @@ jobs:
               tests/unit/rag
               tests/unit/realtime_api
               tests/unit/rerank_api
+              tests/unit/search
               tests/unit/vector_stores
               tests/unit/videos
             python-version: "3.12"

litellm/cost_calculator.py

@@ -10,7 +10,7 @@ from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, cast
 
 from httpx import Response
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, TypedDict, assert_never
 
 import litellm
@@ -1516,6 +1516,9 @@ def completion_cost(
 
         explicit_pricing: Final = custom_pricing is True or base_model is not None
         custom_priced_request_model: Final = model if custom_pricing is True else None
+        configured_cost_per_query: Final = (
+            _configured_search_cost_per_query(litellm_logging_obj) if custom_pricing else None
+        )
         base_model_provider: Final = (
             _base_model_pricing_provider(
                 base_model,
@@ -1762,11 +1765,15 @@ def completion_cost(
                     (
                         prompt_cost,
                         completion_cost_result,
-                    ) = search_provider_cost_per_query(
-                        model=search_model,
-                        custom_llm_provider=custom_llm_provider,
-                        number_of_queries=number_of_queries,
-                        optional_params=optional_params,
+                    ) = (
+                        (number_of_queries * configured_cost_per_query, 0.0)
+                        if configured_cost_per_query is not None
+                        else search_provider_cost_per_query(
+                            model=search_model,
+                            custom_llm_provider=custom_llm_provider,
+                            number_of_queries=number_of_queries,
+                            optional_params=optional_params,
+                        )
                     )
 
                     # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost)
@@ -2152,6 +2159,17 @@ def response_cost_calculator(
         raise e
 
 
+_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
+
+
+def _configured_search_cost_per_query(litellm_logging_obj: LitellmLoggingObject | None) -> float | None:
+    logged_params: Final = getattr(litellm_logging_obj, "litellm_params", None)
+    if logged_params is None:
+        return None
+    cost_per_query: Final = _LOGGED_LITELLM_PARAMS.validate_python(logged_params).get("input_cost_per_query")
+    return float(cost_per_query) if isinstance(cost_per_query, (int, float)) else None
+
+
 def _deployment_model_info(
     litellm_logging_obj: LitellmLoggingObject | None,
     custom_pricing: bool | None,

litellm/model_prices_and_context_window_backup.json

@@ -24280,7 +24280,7 @@
         }
     },
     "agentcore/search": {
-        "input_cost_per_query": 0.0,
+        "input_cost_per_query": 0.007,
         "litellm_provider": "agentcore",
         "mode": "search",
         "metadata": {

litellm/proxy/_lazy_openapi_snapshot.json

@@ -54325,6 +54325,17 @@
               ],
               "title": "Api Key"
             },
+            "input_cost_per_query": {
+              "anyOf": [
+                {
+                  "type": "number"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Input Cost Per Query"
+            },
             "max_retries": {
               "anyOf": [
                 {

litellm/search/main.py

@@ -4,11 +4,13 @@ Main Search function for LiteLLM.
 
 import asyncio
 import contextvars
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Mapping
 from functools import partial
+from types import MappingProxyType
 from typing import Any, Final
 
 import httpx
+from pydantic import TypeAdapter
 
 import litellm
 from litellm._logging import verbose_logger
@@ -24,6 +26,14 @@ base_llm_http_handler = BaseLLMHTTPHandler()
 #################################################
 
 
+_SEARCH_KWARGS: Final = TypeAdapter(Mapping[str, object])
+
+
+def _search_pricing_params(kwargs: Mapping[str, object]) -> Mapping[str, object]:
+    cost_per_query: Final = kwargs.get("input_cost_per_query")
+    return MappingProxyType({} if cost_per_query is None else {"input_cost_per_query": cost_per_query})
+
+
 def _build_search_optional_params(
     max_results: int | None = None,
     search_domain_filter: list[str] | None = None,
@@ -258,7 +268,8 @@ def search(
         )
 
         # Filter out internal LiteLLM parameters from kwargs
-        filtered_kwargs: Final = filter_out_litellm_params(kwargs=kwargs)
+        search_kwargs: Final = _SEARCH_KWARGS.validate_python(kwargs)
+        filtered_kwargs: Final = filter_out_litellm_params(kwargs=search_kwargs)
 
         # Add remaining kwargs to optional_params (for provider-specific params)
         for key, value in filtered_kwargs.items():
@@ -290,6 +301,7 @@ def search(
             litellm_params={
                 "litellm_call_id": litellm_call_id,
                 "api_base": complete_url,
+                **_search_pricing_params(search_kwargs),
             },
             custom_llm_provider=search_provider,
         )

litellm/types/router.py

@@ -963,6 +963,7 @@ class SearchToolLiteLLMParams(TypedDict, total=False):
     api_base: str | None
     timeout: float | str | httpx.Timeout | None
     max_retries: int | None
+    input_cost_per_query: ReadOnly[float | None]
 
 
 class SearchToolInfoTypedDict(TypedDict, total=False):

litellm/types/search.py

@@ -4,7 +4,7 @@ LiteLLM Search API Types
 This module defines types for the unified search API across different providers.
 """
 
-from typing_extensions import Required, TypedDict
+from typing_extensions import ReadOnly, Required, TypedDict
 
 from litellm.types.utils import SearchProviders
 
@@ -24,6 +24,7 @@ class SearchToolLiteLLMParams(TypedDict, total=False):
     api_base: str | None
     timeout: float | None
     max_retries: int | None
+    input_cost_per_query: ReadOnly[float | None]
 
 
 class SearchTool(TypedDict, total=False):

model_prices_and_context_window.json

@@ -24280,7 +24280,7 @@
         }
     },
     "agentcore/search": {
-        "input_cost_per_query": 0.0,
+        "input_cost_per_query": 0.007,
         "litellm_provider": "agentcore",
         "mode": "search",
         "metadata": {

tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py

@@ -636,7 +636,11 @@ class TestAgentCoreSearchEdgeCases:
         from litellm.search.cost_calculator import search_provider_cost_per_query
 
         monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
-        assert search_provider_cost_per_query(model="agentcore/search", custom_llm_provider="agentcore") == (0.0, 0.0)
+        price_per_query = litellm.model_cost["agentcore/search"]["input_cost_per_query"]
+        assert price_per_query > 0
+        assert search_provider_cost_per_query(
+            model="agentcore/search", custom_llm_provider="agentcore", number_of_queries=3
+        ) == (3 * price_per_query, 0.0)
 
 
 def _transform_text(text: str) -> SearchResponse:

tests/unit/search/test_main.py

@@ -0,0 +1,58 @@
+import datetime
+from typing import Final
+
+import pytest
+import respx
+
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging
+from litellm.search.cost_calculator import search_provider_cost_per_query
+
+EXA_SEARCH_URL: Final = "https://api.exa.ai/search"
+
+
+@pytest.fixture
+def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
+
+    monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map())
+
+
+def _proxy_style_logging_obj() -> Logging:
+    """The proxy builds the logging object from the request body before the router adds the search tool's
+    litellm_params, so the per-tool price only reaches cost tracking if the search call itself records it."""
+    return Logging(
+        model="exa-search",
+        messages=[],
+        stream=False,
+        call_type="search",
+        start_time=datetime.datetime.now(),
+        litellm_call_id="search-tool-price",
+        function_id="f",
+    )
+
+
+@pytest.mark.usefixtures("_local_model_cost_map")
+def test_search_bills_search_tool_input_cost_per_query(respx_mock: respx.MockRouter) -> None:
+    respx_mock.post(EXA_SEARCH_URL).respond(json={"results": [{"title": "t", "url": "https://example.com"}]})
+    map_price, _ = search_provider_cost_per_query(model="exa_ai/search", custom_llm_provider="exa_ai")
+    configured_price: Final = 0.05
+    assert map_price not in (0.0, configured_price)
+
+    map_priced: Final = litellm.search(
+        query="litellm release notes",
+        search_provider="exa_ai",
+        api_key="test-exa-key",
+        litellm_logging_obj=_proxy_style_logging_obj(),
+    )
+    tool_priced: Final = litellm.search(
+        query="litellm release notes",
+        search_provider="exa_ai",
+        api_key="test-exa-key",
+        input_cost_per_query=configured_price,
+        litellm_logging_obj=_proxy_style_logging_obj(),
+    )
+
+    assert map_priced._hidden_params["response_cost"] == pytest.approx(map_price)
+    assert tool_priced._hidden_params["response_cost"] == pytest.approx(configured_price)
+    assert respx_mock.calls.last.request.content.find(b"input_cost_per_query") == -1

tests/unit/test_cost_calculator.py

@@ -9054,3 +9054,28 @@ def test_live_session_uses_latest_snapshot_when_closed_usage_is_missing(_local_m
     )
 
     assert cost == pytest.approx(15 * rate)
+
+
+@pytest.mark.parametrize("query, expected_queries", [("litellm release notes", 1), (["a", "b", "c"], 3)])
+def test_completion_cost_search_uses_search_tool_input_cost_per_query(
+    _local_model_cost_map: None, query: str | list[str], expected_queries: int
+) -> None:
+    map_cost: Final = completion_cost(
+        completion_response=None,
+        model="agentcore/search",
+        custom_llm_provider="agentcore",
+        call_type="asearch",
+        optional_params={"query": query},
+    )
+    configured_cost: Final = completion_cost(
+        completion_response=None,
+        model="agentcore/search",
+        custom_llm_provider="agentcore",
+        call_type="asearch",
+        optional_params={"query": query},
+        custom_pricing=True,
+        litellm_logging_obj=SimpleNamespace(litellm_params={"input_cost_per_query": 0.05}),
+    )
+
+    assert map_cost == pytest.approx(expected_queries * litellm.model_cost["agentcore/search"]["input_cost_per_query"])
+    assert configured_cost == pytest.approx(expected_queries * 0.05)

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -44711,6 +44711,8 @@ export interface components {
             api_base?: string | null;
             /** Api Key */
             api_key?: string | null;
+            /** Input Cost Per Query */
+            input_cost_per_query?: number | null;
             /** Max Retries */
             max_retries?: number | null;
             /** Search Provider */