scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(websearch): preserve request header precedence

Browse source
This commit is contained in:
glaziermag 2026-07-12 09:20:52 -07:00
parent 9eb2d06ff1
commit 8be9aed8a7
2 changed files with 74 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/integrations/websearch_interception/handler.py
View file

@ -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

86
tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py
View file

@ -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