diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 1277cac51d..d7858d71eb 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -16,6 +16,7 @@ from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import ( get_litellm_web_search_tool, + get_litellm_web_search_tool_openai, is_web_search_tool, is_web_search_tool_chat_completion, ) @@ -77,7 +78,13 @@ class WebSearchInterceptionLogger(CustomLogger): that we can intercept and execute ourselves. """ # Check if this is for an enabled provider - custom_llm_provider = kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + # Try top-level kwargs first, then nested litellm_params, then derive from model name + custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get("custom_llm_provider", "") + if not custom_llm_provider: + try: + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) + except Exception: + custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: return None @@ -101,7 +108,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool in tools: if is_web_search_tool(tool): # Convert to LiteLLM standard web search tool - converted_tool = get_litellm_web_search_tool() + converted_tool = get_litellm_web_search_tool_openai() converted_tools.append(converted_tool) verbose_logger.debug( f"WebSearchInterception: Converted {tool.get('name', 'unknown')} " @@ -111,8 +118,9 @@ class WebSearchInterceptionLogger(CustomLogger): # Keep other tools as-is converted_tools.append(tool) - # Return modified kwargs with converted tools - return {"tools": converted_tools} + # Update tools in-place and return full kwargs + kwargs["tools"] = converted_tools + return kwargs @classmethod def from_config_yaml( diff --git a/litellm/integrations/websearch_interception/tools.py b/litellm/integrations/websearch_interception/tools.py index c39d150fb1..7ef2b35004 100644 --- a/litellm/integrations/websearch_interception/tools.py +++ b/litellm/integrations/websearch_interception/tools.py @@ -49,6 +49,39 @@ def get_litellm_web_search_tool() -> Dict[str, Any]: } +def get_litellm_web_search_tool_openai() -> Dict[str, Any]: + """ + Get the standard LiteLLM web search tool definition in OpenAI format. + + Used by async_pre_call_deployment_hook which runs in the chat completions + path where tools must be in OpenAI format (type: "function" with + function.parameters). + + Returns: + Dict containing the OpenAI-style tool definition. + """ + return { + "type": "function", + "function": { + "name": LITELLM_WEB_SEARCH_TOOL_NAME, + "description": ( + "Search the web for information. Use this when you need current " + "information or answers to questions that require up-to-date data." + ), + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to execute" + } + }, + "required": ["query"] + } + } + } + + def is_web_search_tool_chat_completion(tool: Dict[str, Any]) -> bool: """ Check if a tool is a web search tool for Chat Completions API (strict check). 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 5abecb46c9..85ed5dddee 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 @@ -100,3 +100,176 @@ async def test_internal_flags_filtered_from_followup_kwargs(): # Verify regular kwargs are preserved assert kwargs_for_followup["temperature"] == 0.7 assert kwargs_for_followup["max_tokens"] == 1024 + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_provider_from_top_level_kwargs(): + """Test that async_pre_call_deployment_hook finds custom_llm_provider at top-level kwargs. + + Regression test for bug where the hook only checked kwargs["litellm_params"]["custom_llm_provider"] + but the router places custom_llm_provider at the top level of kwargs. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + # Simulate kwargs as they arrive from the router path: + # custom_llm_provider is at the TOP LEVEL (not nested under litellm_params) + kwargs = { + "model": "anthropic.claude-3-5-sonnet-20241022-v2:0", + "messages": [{"role": "user", "content": "Search the web for LiteLLM"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + {"type": "function", "function": {"name": "other_tool", "parameters": {}}}, + ], + "custom_llm_provider": "bedrock", + "api_key": "fake-key", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + # Should NOT be None — the hook should have triggered + assert result is not None + # The web_search tool should be converted to litellm_web_search (OpenAI format) + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # The non-web-search tool should be preserved + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "other_tool" + for t in result["tools"] + ) + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_returns_full_kwargs(): + """Test that async_pre_call_deployment_hook returns the full kwargs dict, not a partial one. + + Regression test for bug where the hook returned {"tools": converted_tools} instead of + the full kwargs dict, causing model/messages/api_key/etc. to be lost. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Search for something"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search"}, + ], + "custom_llm_provider": "openai", + "api_key": "sk-fake", + "temperature": 0.7, + "metadata": {"user": "test"}, + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + assert result is not None + # All original keys must be preserved + assert result["model"] == "gpt-4o" + assert result["messages"] == [{"role": "user", "content": "Search for something"}] + assert result["api_key"] == "sk-fake" + assert result["temperature"] == 0.7 + assert result["metadata"] == {"user": "test"} + assert result["custom_llm_provider"] == "openai" + # Tools should be converted + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_skips_disabled_provider(): + """Test that the hook returns None for providers not in enabled_providers.""" + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "test"}], + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + "custom_llm_provider": "openai", # Not in enabled_providers + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + assert result is None + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_skips_no_websearch_tools(): + """Test that the hook returns None when no web search tools are present.""" + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "test"}], + "tools": [ + {"type": "function", "function": {"name": "calculator", "parameters": {}}}, + ], + "custom_llm_provider": "openai", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + assert result is None + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_nested_litellm_params_fallback(): + """Test that the hook still works when custom_llm_provider is in nested litellm_params. + + This is the Anthropic experimental pass-through path where litellm_params is + explicitly constructed with custom_llm_provider inside it. + """ + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + kwargs = { + "model": "anthropic.claude-3-5-sonnet-20241022-v2:0", + "messages": [{"role": "user", "content": "test"}], + "tools": [{"type": "web_search_20250305", "name": "web_search"}], + "litellm_params": { + "custom_llm_provider": "bedrock", + }, + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + assert result is not None + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # Full kwargs preserved + assert result["model"] == "anthropic.claude-3-5-sonnet-20241022-v2:0" + + +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_provider_derived_from_model_name(): + """Test that async_pre_call_deployment_hook derives custom_llm_provider from the model name. + + Regression test for the router _acompletion path where custom_llm_provider is NOT + in kwargs at all — neither at top-level nor in litellm_params. The hook must derive + the provider from the model name (e.g., "openai/gpt-4o-mini" → "openai"). + """ + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + + # Simulate kwargs as they arrive from router._acompletion: + # NO custom_llm_provider key anywhere — only model name contains the provider + kwargs = { + "model": "openai/gpt-4o-mini", + "messages": [{"role": "user", "content": "Search the web for LiteLLM"}], + "tools": [ + {"type": "web_search_20250305", "name": "web_search", "max_uses": 3}, + ], + "api_key": "fake-key", + } + + result = await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=None) + + # Should NOT be None — the hook should derive "openai" from "openai/gpt-4o-mini" + assert result is not None + assert any( + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" + for t in result["tools"] + ) + # Full kwargs preserved + assert result["model"] == "openai/gpt-4o-mini" + assert result["api_key"] == "fake-key"