diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 7abbf0c96e5..2c5f7484dac 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -42,7 +42,12 @@ WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import AgenticLoopParams, CallTypes, LlmProviders +from litellm.types.utils import ( + AgenticLoopParams, + CallTypes, + LlmProviders, + StandardLoggingUserAPIKeyMetadata, +) from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -1318,6 +1323,7 @@ async def _execute_search( search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: dict[str, Any] = {} + search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) @@ -1334,12 +1340,30 @@ async def _execute_search( verbose_logger.debug( "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider ) + user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs) + search_metadata: Final = ( + None + if user_api_key_auth is None + else self._build_search_request_metadata( + user_api_key_auth=user_api_key_auth, + search_tool_name=search_tool_name, + ) + ) search_kwargs: Final = { key: value for key, value in search_litellm_params.items() if key != "search_provider" and value is not None } - result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + result: Final = ( + await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + if search_metadata is None + else await litellm.asearch( + query=query, + search_provider=search_provider, + litellm_metadata=search_metadata, + **search_kwargs, + ) + ) # Format using transformation function search_result_text: Final = WebSearchTransformation.format_search_response(result) @@ -1396,6 +1420,35 @@ async def _authorize_search_tool( team_object=team_object, ) + @staticmethod + def _build_search_request_metadata( + user_api_key_auth: "UserAPIKeyAuth", + search_tool_name: str | None, + ) -> Mapping[str, object]: + """ + Spend-tracking metadata for the intercepted search, so its provider cost is logged + and billed against the key/user/team that made the originating LLM request instead + of being dropped by the proxy's spend hook for lack of an owner. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) + ) + return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches + **user_api_key_metadata, + "model_group": search_tool_name, + "user_api_key": user_api_key_auth.api_key, + "user_api_key_auth": user_api_key_auth, + } + + @staticmethod + def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None: + if search_tool is None: + return None + search_tool_name: Final = search_tool.get("search_tool_name") + return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None + @staticmethod def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index b6ff3b70a4d..f39f41a6d12 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, ) - mock_asearch.assert_awaited_once_with( - query="what is litellm", - search_provider="tavily", - api_key="fake-ui-key", - api_base="https://api.tavily.com", - timeout=10.0, - max_retries=2, + forwarded_kwargs = mock_asearch.await_args.kwargs + assert forwarded_kwargs["query"] == "what is litellm" + assert forwarded_kwargs["search_provider"] == "tavily" + assert forwarded_kwargs["api_key"] == "fake-ui-key" + assert forwarded_kwargs["api_base"] == "https://api.tavily.com" + assert forwarded_kwargs["timeout"] == 10.0 + assert forwarded_kwargs["max_retries"] == 2 + + +@pytest.mark.asyncio +async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): + """An intercepted search is billed and logged against the key that made the LLM request. + + Without the forwarded attribution metadata the proxy's spend hook skips the search + entirely, so its provider cost never reaches SpendLogs or any budget. + """ + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + user_api_key_auth = UserAPIKeyAuth( + api_key="hashed-sk-1234", + key_alias="alice-key", + user_id="user-alice", + org_id="org-1", + ) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search( + "what is litellm", + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, + ) + + forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"] + assert forwarded_metadata["user_api_key"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_alias"] == "alice-key" + assert forwarded_metadata["user_api_key_user_id"] == "user-alice" + assert forwarded_metadata["user_api_key_org_id"] == "org-1" + assert forwarded_metadata["model_group"] == "perplexity-sonar-pro" + assert ( + _should_track_cost_callback( + user_api_key=forwarded_metadata["user_api_key"], + user_id=forwarded_metadata["user_api_key_user_id"], + team_id=forwarded_metadata["user_api_key_team_id"], + end_user_id=None, + call_type="asearch", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch): + """SDK callers have no key to attribute the search to, so no proxy metadata is invented.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", + ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm", kwargs={"litellm_params": {}}) + + assert "litellm_metadata" not in mock_asearch.await_args.kwargs @pytest.mark.asyncio