diff --git a/litellm/search/main.py b/litellm/search/main.py index 27bcece898b..fb19db5a483 100644 --- a/litellm/search/main.py +++ b/litellm/search/main.py @@ -30,8 +30,9 @@ _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}) + """Always carries the key so an unpriced fallback tool clears the price a failed primary left in the + shared logging object (`Logging.update_environment_variables` merges litellm_params).""" + return MappingProxyType({"input_cost_per_query": kwargs.get("input_cost_per_query")}) def _build_search_optional_params( diff --git a/tests/integration/providers/test_search_wire.py b/tests/integration/providers/test_search_wire.py index b60d92cb6a6..5ee80f7f5db 100644 --- a/tests/integration/providers/test_search_wire.py +++ b/tests/integration/providers/test_search_wire.py @@ -57,13 +57,21 @@ _EXPECTED_RESPONSE: Final = _SearchResponse( ) -def _create_exa_tool(gateway: Gateway, scenario: Scenario, name: str, api_base: str) -> None: +def _create_exa_tool( + gateway: Gateway, scenario: Scenario, name: str, api_base: str, input_cost_per_query: float | None = None +) -> None: + pricing: Final = {} if input_cost_per_query is None else {"input_cost_per_query": input_cost_per_query} created: Final = gateway.post( "/search_tools", { "search_tool": { "search_tool_name": name, - "litellm_params": {"search_provider": "exa_ai", "api_key": _EXA_KEY, "api_base": api_base}, + "litellm_params": { + "search_provider": "exa_ai", + "api_key": _EXA_KEY, + "api_base": api_base, + **pricing, + }, } }, ) @@ -189,3 +197,28 @@ def test_unknown_search_tool_is_a_client_error_without_any_upstream_call(gateway response.text ) assert wire.drain() == () + + +def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed_primarys(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + assert "input_cost_per_query" not in json.loads(request.body), request.body + if request.target == "/down/search": + return Reply(status=500, body=json.dumps({"error": "primary down"}).encode()) + assert (request.method, request.target) == ("POST", "/exa/search"), request.target + return Reply(body=json.dumps(_EXA_RESULTS).encode()) + + primary: Final = f"exa-priced-{uuid.uuid4().hex}" + backup: Final = f"exa-backup-{uuid.uuid4().hex}" + primary_price: Final = 0.05 + with wire_server(respond) as wire, gateway.scenario() as scenario: + _create_exa_tool(gateway, scenario, primary, f"{wire.url}/down", input_cost_per_query=primary_price) + _create_exa_tool(gateway, scenario, backup, f"{wire.url}/exa") + direct: Final = gateway.request("POST", f"/v1/search/{backup}", {"query": _QUERY}) + assert direct.status_code == 200, direct.text + backup_price: Final = float(direct.headers["x-litellm-response-cost"]) + assert backup_price not in (0.0, primary_price), direct.headers + + fallback: Final = gateway.request("POST", f"/v1/search/{primary}", {"query": _QUERY, "fallbacks": [backup]}) + assert fallback.status_code == 200, fallback.text + assert float(fallback.headers["x-litellm-response-cost"]) == pytest.approx(backup_price), fallback.headers + assert [request.target for request in wire.drain()] == ["/exa/search", "/down/search", "/exa/search"] diff --git a/tests/unit/search/test_main.py b/tests/unit/search/test_main.py index 62e28b6e7fa..72266321715 100644 --- a/tests/unit/search/test_main.py +++ b/tests/unit/search/test_main.py @@ -1,6 +1,7 @@ import datetime from typing import Final +import httpx import pytest import respx @@ -56,3 +57,36 @@ def test_search_bills_search_tool_input_cost_per_query(respx_mock: respx.MockRou 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 + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_unpriced_fallback_search_tool_is_billed_at_its_own_price_not_the_failed_primarys( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(EXA_SEARCH_URL).mock( + side_effect=( + httpx.Response(500, json={"error": "primary down"}), + httpx.Response(200, json={"results": [{"title": "t", "url": "https://example.com"}]}), + ) + ) + map_price, _ = search_provider_cost_per_query(model="exa_ai/search", custom_llm_provider="exa_ai") + primary_price: Final = 0.05 + assert map_price != primary_price + shared_by_every_fallback_attempt: Final = _proxy_style_logging_obj() + + with pytest.raises(litellm.InternalServerError): + litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + input_cost_per_query=primary_price, + litellm_logging_obj=shared_by_every_fallback_attempt, + ) + fallback: Final = litellm.search( + query="litellm release notes", + search_provider="exa_ai", + api_key="test-exa-key", + litellm_logging_obj=shared_by_every_fallback_attempt, + ) + + assert fallback._hidden_params["response_cost"] == pytest.approx(map_price)