diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index f521b2c8edd..41823f51280 100644 --- a/.github/workflows/test-unit.yml +++ b/.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" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9915d76433d..918b95280be 100644 --- a/litellm/cost_calculator.py +++ b/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, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 04f391d4d20..0faab05ecf8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/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": { diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 9086c8417ae..f644d08738d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/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": [ { diff --git a/litellm/search/main.py b/litellm/search/main.py index e3db8bfb0ac..27bcece898b 100644 --- a/litellm/search/main.py +++ b/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, ) diff --git a/litellm/types/router.py b/litellm/types/router.py index de993d46e44..285a97c3221 100644 --- a/litellm/types/router.py +++ b/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): diff --git a/litellm/types/search.py b/litellm/types/search.py index b4180d47a51..32a92d945e5 100644 --- a/litellm/types/search.py +++ b/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): diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 04f391d4d20..0faab05ecf8 100644 --- a/model_prices_and_context_window.json +++ b/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": { diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 4526c428f1a..664dbf87fe9 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/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: diff --git a/tests/unit/search/__init__.py b/tests/unit/search/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py new file mode 100644 index 00000000000..62e28b6e7fa --- /dev/null +++ b/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 diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index b6fa2054514..afbc73ab280 100644 --- a/tests/unit/test_cost_calculator.py +++ b/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) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 87213359f2d..eb74003c103 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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 */