From 8be9aed8a75924a54ad79aafe4becd8f5ceed725 Mon Sep 17 00:00:00 2001 From: glaziermag <130600081+glaziermag@users.noreply.github.com> Date: Sun, 12 Jul 2026 09:20:52 -0700 Subject: [PATCH] fix(websearch): preserve request header precedence --- .../websearch_interception/handler.py | 11 +-- .../test_websearch_interception_handler.py | 86 +++++++++++++++---- 2 files changed, 74 insertions(+), 23 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 815d77a9b52..16ce24bc565 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -862,13 +862,14 @@ class WebSearchInterceptionLogger(CustomLogger): 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. + # Deployment ``extra_headers`` is stored inside ``litellm_params`` on + # the initial request, but follow-up APIs consume it as a top-level + # kwarg. Keep a request-level value when both are present, matching the + # Router's request-over-deployment precedence. litellm_params = kwargs.get("litellm_params") - if isinstance(litellm_params, dict): + if "extra_headers" not in kwargs_for_followup and isinstance(litellm_params, dict): extra_headers = litellm_params.get("extra_headers") - if extra_headers: + if extra_headers is not None: kwargs_for_followup["extra_headers"] = extra_headers return kwargs_for_followup 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 96723ba7f93..4380068b647 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,7 +56,9 @@ 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" @@ -117,7 +119,9 @@ 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, @@ -170,7 +174,9 @@ 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 @@ -210,12 +216,15 @@ 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"] ) @@ -252,7 +261,8 @@ 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"] ) @@ -313,7 +323,8 @@ 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 @@ -346,7 +357,8 @@ 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 @@ -354,11 +366,28 @@ async def test_async_pre_call_deployment_hook_provider_derived_from_model_name() assert result["api_key"] == "fake-key" +@pytest.mark.parametrize( + ("kwargs", "expected_headers"), + [ + ( + {"litellm_params": {"extra_headers": {"source": "deployment"}}}, + {"source": "deployment"}, + ), + ( + { + "extra_headers": {"source": "request"}, + "litellm_params": {"extra_headers": {"source": "deployment"}}, + }, + {"source": "request"}, + ), + ], +) @pytest.mark.asyncio -async def test_anthropic_followup_promotes_extra_headers_from_litellm_params(): +async def test_anthropic_followup_preserves_extra_headers_precedence( + kwargs, expected_headers +): 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", @@ -373,17 +402,34 @@ async def test_anthropic_followup_promotes_extra_headers_from_litellm_params(): thinking_blocks=[], anthropic_messages_optional_request_params={}, logging_obj=Mock(model_call_details={}), - kwargs={"litellm_params": {"extra_headers": extra_headers}}, + kwargs=kwargs, ) - assert patch.kwargs["extra_headers"] == extra_headers + assert patch.kwargs["extra_headers"] == expected_headers +@pytest.mark.parametrize( + ("kwargs", "expected_headers"), + [ + ( + {"litellm_params": {"extra_headers": {"source": "deployment"}}}, + {"source": "deployment"}, + ), + ( + { + "extra_headers": {"source": "request"}, + "litellm_params": {"extra_headers": {"source": "deployment"}}, + }, + {"source": "request"}, + ), + ], +) @pytest.mark.asyncio -async def test_chat_completion_followup_promotes_extra_headers_from_litellm_params(): +async def test_chat_completion_followup_preserves_extra_headers_precedence( + kwargs, expected_headers +): 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", @@ -396,14 +442,16 @@ async def test_chat_completion_followup_promotes_extra_headers_from_litellm_para } ], optional_params={}, - kwargs={"litellm_params": {"extra_headers": extra_headers}}, + kwargs=kwargs, ) - assert patch.kwargs["extra_headers"] == extra_headers + assert patch.kwargs["extra_headers"] == expected_headers def test_prepare_followup_kwargs_handles_none_litellm_params(): - assert WebSearchInterceptionLogger._prepare_followup_kwargs({"litellm_params": None}) == {"litellm_params": None} + assert WebSearchInterceptionLogger._prepare_followup_kwargs( + {"litellm_params": None} + ) == {"litellm_params": None} @pytest.mark.asyncio @@ -518,7 +566,9 @@ 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