diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 6be64c89828..3e24c64ad82 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1288,14 +1288,11 @@ class WebSearchInterceptionLogger(CustomLogger): 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: str | None = None + 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 {}) search_provider = search_litellm_params.get("search_provider") - selected_tool_name = search_tool.get("search_tool_name") - if isinstance(selected_tool_name, str) and selected_tool_name: - search_tool_name = selected_tool_name # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1318,14 +1315,20 @@ class WebSearchInterceptionLogger(CustomLogger): ) ) search_kwargs: Final = { - **{ - key: value - for key, value in search_litellm_params.items() - if key != "search_provider" and value is not None - }, - **({} if search_metadata is None else {"litellm_metadata": search_metadata}), + 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) @@ -1386,7 +1389,7 @@ class WebSearchInterceptionLogger(CustomLogger): def _build_search_request_metadata( user_api_key_auth: "UserAPIKeyAuth", search_tool_name: str | None, - ) -> dict[str, object]: + ) -> 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 @@ -1394,16 +1397,23 @@ class WebSearchInterceptionLogger(CustomLogger): """ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - user_api_key_metadata: StandardLoggingUserAPIKeyMetadata = ( + user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) ) - return { + return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches **user_api_key_metadata, - **({} if search_tool_name is None else {"model_group": search_tool_name}), + "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: