fix(search): bill an unpriced fallback search tool at its own price, not the failed primary's (#46734)

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

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(

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"]

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)