diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 60100e8c2fd..815d77a9b52 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -841,7 +841,10 @@ class WebSearchInterceptionLogger(CustomLogger): return max_tokens @staticmethod - def _prepare_followup_kwargs(kwargs: Dict) -> Dict: + def _prepare_followup_kwargs( + kwargs: Dict, + internal_keys: Optional[set[str]] = None, + ) -> Dict: """Build kwargs for the follow-up call, excluding internal keys. ``litellm_logging_obj`` MUST be excluded so the follow-up call creates @@ -852,10 +855,24 @@ class WebSearchInterceptionLogger(CustomLogger): SpendLog / AWS billing mismatch. """ _internal_keys = {"litellm_logging_obj"} - return { + if internal_keys: + _internal_keys.update(internal_keys) + + kwargs_for_followup = { k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in _internal_keys } + # ``extra_headers`` is stored inside ``litellm_params`` on the initial + # request, but follow-up completion APIs consume it as a top-level kwarg. + # Promote it without assuming ``litellm_params`` is always a dict. + litellm_params = kwargs.get("litellm_params") + if isinstance(litellm_params, dict): + extra_headers = litellm_params.get("extra_headers") + if extra_headers: + kwargs_for_followup["extra_headers"] = extra_headers + + return kwargs_for_followup + async def _execute_agentic_loop( self, model: str, @@ -1203,9 +1220,10 @@ class WebSearchInterceptionLogger(CustomLogger): "stream_response", "custom_prompt_dict", } - kwargs_for_followup = { - k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in internal_params - } + kwargs_for_followup = self._prepare_followup_kwargs( + kwargs, + internal_keys=internal_params, + ) full_model_name = model if "custom_llm_provider" in 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 ea117380edf..96723ba7f93 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 @@ -56,9 +56,7 @@ def test_initialize_from_proxy_config_honors_dict_callback_specific_params(): """A valid dict under callback_settings.websearch_interception is applied.""" logger = WebSearchInterceptionLogger.initialize_from_proxy_config( litellm_settings={}, - callback_specific_params={ - "websearch_interception": {"search_tool_name": "ws-tool"} - }, + callback_specific_params={"websearch_interception": {"search_tool_name": "ws-tool"}}, ) assert logger.search_tool_name == "ws-tool" @@ -119,9 +117,7 @@ async def test_async_build_agentic_loop_plan_returns_request_patch(): "response_format": "anthropic", } logging_obj = MagicMock() - logging_obj.model_call_details = { - "agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"} - } + logging_obj.model_call_details = {"agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"}} kwargs = { "temperature": 0.2, "_websearch_interception_converted_stream": True, @@ -174,9 +170,7 @@ async def test_internal_flags_filtered_from_followup_kwargs(): # Apply the same filtering logic used in _execute_agentic_loop kwargs_for_followup = { - k: v - for k, v in kwargs_with_internal_flags.items() - if not k.startswith("_websearch_interception") + k: v for k, v in kwargs_with_internal_flags.items() if not k.startswith("_websearch_interception") } # Verify internal flags are filtered out @@ -216,15 +210,12 @@ async def test_async_pre_call_deployment_hook_provider_from_top_level_kwargs(): 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" + 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"] + t.get("type") == "function" and t.get("function", {}).get("name") == "other_tool" for t in result["tools"] ) @@ -261,8 +252,7 @@ async def test_async_pre_call_deployment_hook_returns_full_kwargs(): 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" + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" for t in result["tools"] ) @@ -323,8 +313,7 @@ async def test_async_pre_call_deployment_hook_nested_litellm_params_fallback(): assert result is not None assert any( - t.get("type") == "function" - and t.get("function", {}).get("name") == "litellm_web_search" + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" for t in result["tools"] ) # Full kwargs preserved @@ -357,8 +346,7 @@ async def test_async_pre_call_deployment_hook_provider_derived_from_model_name() # 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" + t.get("type") == "function" and t.get("function", {}).get("name") == "litellm_web_search" for t in result["tools"] ) # Full kwargs preserved @@ -366,6 +354,58 @@ async def test_async_pre_call_deployment_hook_provider_derived_from_model_name() assert result["api_key"] == "fake-key" +@pytest.mark.asyncio +async def test_anthropic_followup_promotes_extra_headers_from_litellm_params(): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + logger._execute_search = AsyncMock(return_value=("result", None)) + extra_headers = {"editor-version": "vscode/1.110.0"} + + patch, _ = await logger._build_anthropic_request_patch( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[], + tool_calls=[ + { + "id": "call_1", + "name": LITELLM_WEB_SEARCH_TOOL_NAME, + "input": {"query": "hello"}, + } + ], + thinking_blocks=[], + anthropic_messages_optional_request_params={}, + logging_obj=Mock(model_call_details={}), + kwargs={"litellm_params": {"extra_headers": extra_headers}}, + ) + + assert patch.kwargs["extra_headers"] == extra_headers + + +@pytest.mark.asyncio +async def test_chat_completion_followup_promotes_extra_headers_from_litellm_params(): + logger = WebSearchInterceptionLogger(enabled_providers=["openai"]) + logger._execute_search = AsyncMock(return_value=("result", None)) + extra_headers = {"editor-version": "vscode/1.110.0"} + + patch = await logger._build_chat_completion_request_patch( + model="gpt-4o", + messages=[], + tool_calls=[ + { + "id": "call_1", + "name": LITELLM_WEB_SEARCH_TOOL_NAME, + "input": {"query": "hello"}, + } + ], + optional_params={}, + kwargs={"litellm_params": {"extra_headers": extra_headers}}, + ) + + assert patch.kwargs["extra_headers"] == extra_headers + + +def test_prepare_followup_kwargs_handles_none_litellm_params(): + assert WebSearchInterceptionLogger._prepare_followup_kwargs({"litellm_params": None}) == {"litellm_params": None} + + @pytest.mark.asyncio async def test_deployment_hook_converts_stream_and_logging_obj_syncs(): """ @@ -478,9 +518,7 @@ def test_sync_forced_tool_choice_leaves_non_forced_untouched(tool_choice): and None pass through unchanged.""" converted_tools = [{"name": LITELLM_WEB_SEARCH_TOOL_NAME}] - result = WebSearchInterceptionLogger._sync_forced_tool_choice( - tool_choice, converted_tools - ) + result = WebSearchInterceptionLogger._sync_forced_tool_choice(tool_choice, converted_tools) assert result == tool_choice