diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2e11405af3f..bfae6d5b7b0 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -31,6 +31,7 @@ from litellm.types.integrations.websearch_interception import ( WebSearchInterceptionConfig, ) from litellm.types.integrations.custom_logger import ( + CHAT_COMPLETION_AGENTIC_SURFACE, AgenticLoopPlan, AgenticLoopRequestPatch, ) @@ -440,12 +441,16 @@ class WebSearchInterceptionLogger(CustomLogger): custom_llm_provider: str, kwargs: Dict, ) -> Tuple[bool, Dict]: - """ - Check if WebSearch tool interception is needed for Anthropic Messages API. - - This is the legacy method for Anthropic-style responses. - For chat completions, use async_should_run_chat_completion_agentic_loop instead. - """ + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return await self.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) verbose_logger.debug(f"WebSearchInterception: Hook called! provider={custom_llm_provider}, stream={stream}") verbose_logger.debug(f"WebSearchInterception: Response type: {type(response)}") @@ -629,6 +634,18 @@ class WebSearchInterceptionLogger(CustomLogger): stream: bool, kwargs: Dict, ) -> AgenticLoopPlan: + if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return await self.async_build_chat_completion_agentic_loop_plan( + tools=tools, + model=model, + messages=messages, + response=response, + optional_params=anthropic_messages_optional_request_params, + logging_obj=logging_obj, + stream=stream, + kwargs=kwargs, + ) + tool_calls = tools["tool_calls"] thinking_blocks = tools.get("thinking_blocks", []) request_patch, structured_results = await self._build_anthropic_request_patch( @@ -1088,6 +1105,7 @@ class WebSearchInterceptionLogger(CustomLogger): raise ValueError("WebSearchInterception: missing follow-up messages") params = dict(optional_params) params.update(request_patch.optional_params) + params.pop("tool_choice", None) return await litellm.acompletion( model=request_patch.model or model, messages=request_patch.messages, @@ -1203,6 +1221,7 @@ class WebSearchInterceptionLogger(CustomLogger): if k not in { "tools", + "tool_choice", "extra_body", "model_alias_map", "stream_response", diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 828605d5ef8..b7262a42324 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -137,8 +137,8 @@ async def _execute_chat_completion_agentic_plan( optional_params_for_followup = {**optional_params, **patch.optional_params} if patch.tools is not None: optional_params_for_followup["tools"] = patch.tools - if "tool_choice" not in patch.optional_params: - optional_params_for_followup.pop("tool_choice", None) + if "tool_choice" not in patch.optional_params: + optional_params_for_followup.pop("tool_choice", None) kwargs_for_followup = _filter_followup_kwargs(kwargs) kwargs_for_followup.update( @@ -206,10 +206,11 @@ async def maybe_run_chat_completion_agentic_loop( for callback in callbacks: if not isinstance(callback, CustomLogger): continue + if not _gate_overridden(callback): continue - gate_kwargs = { + hook_kwargs = { **kwargs, "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, "custom_llm_provider": custom_llm_provider, @@ -222,7 +223,7 @@ async def maybe_run_chat_completion_agentic_loop( tools=tools, stream=stream, custom_llm_provider=custom_llm_provider, - kwargs=gate_kwargs, + kwargs=hook_kwargs, ) except Exception as e: verbose_logger.exception( @@ -243,11 +244,6 @@ async def maybe_run_chat_completion_agentic_loop( ) try: - plan_kwargs = { - **kwargs, - "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, - "custom_llm_provider": custom_llm_provider, - } if not _build_plan_overridden(callback): return await callback.async_run_agentic_loop( tools=tool_calls, @@ -258,7 +254,7 @@ async def maybe_run_chat_completion_agentic_loop( anthropic_messages_optional_request_params=optional_params, logging_obj=logging_obj, stream=stream, - kwargs=plan_kwargs, + kwargs=hook_kwargs, ) plan = await callback.async_build_agentic_loop_plan( @@ -270,7 +266,7 @@ async def maybe_run_chat_completion_agentic_loop( anthropic_messages_optional_request_params=optional_params, logging_obj=logging_obj, stream=stream, - kwargs=plan_kwargs, + kwargs=hook_kwargs, ) if plan.response_override is not None: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9bb956d0808..3c10239f868 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2112,7 +2112,7 @@ class BaseLLMHTTPHandler: anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs={**kwargs, "api_key": api_key} if api_key else kwargs, ) return initial_response else: @@ -2122,6 +2122,10 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + # Inject api_key into kwargs so follow-up calls in agentic hooks can + # authenticate. api_key is a named param here (not in kwargs), so + # _prepare_followup_kwargs would miss it otherwise. + kwargs_for_agentic = {**kwargs, "api_key": api_key} if api_key else kwargs # Call agentic completion hooks (non-streaming path only) final_response = await self._call_agentic_completion_hooks( response=initial_response, @@ -2132,7 +2136,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=False, custom_llm_provider=custom_llm_provider, - kwargs=kwargs, + kwargs=kwargs_for_agentic, ) return self._maybe_wrap_in_fake_stream( diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py index 34555d76554..7ef43e2eadf 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_chat_completion.py @@ -6,7 +6,7 @@ litellm.acompletion() for transparent server-side web search execution. """ import os -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -34,9 +34,7 @@ def mock_search_response(): @pytest.fixture def websearch_logger(): """Create a WebSearchInterceptionLogger instance""" - return WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.OPENAI, LlmProviders.MINIMAX] - ) + return WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI, LlmProviders.MINIMAX]) @pytest.mark.asyncio @@ -55,9 +53,7 @@ async def test_websearch_chat_completion_with_openai(): """ # Configure WebSearch interception original_callbacks = litellm.callbacks.copy() if litellm.callbacks else [] - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.OPENAI] - ) + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) litellm.callbacks = [websearch_logger] try: @@ -100,9 +96,7 @@ async def test_websearch_chat_completion_with_openai(): if hasattr(response.choices[0].message, "tool_calls"): # If tool_calls exist, it means agentic loop didn't run # This could happen if search tool is not configured - pytest.skip( - "Agentic loop did not execute - search tool may not be configured" - ) + pytest.skip("Agentic loop did not execute - search tool may not be configured") # Verify we got a meaningful response assert response.choices[0].finish_reason in ["stop", "end_turn"] @@ -122,9 +116,7 @@ async def test_websearch_chat_completion_hook_detection(): Message, ) - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.OPENAI] - ) + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) # Mock response with litellm_web_search tool call mock_response = ModelResponse( @@ -155,21 +147,19 @@ async def test_websearch_chat_completion_hook_detection(): ) # Test should_run_chat_completion_agentic_loop - should_run, tools_dict = ( - await websearch_logger.async_should_run_chat_completion_agentic_loop( - response=mock_response, - model="gpt-4o", - messages=[{"role": "user", "content": "What's the weather?"}], - tools=[ - { - "type": "function", - "function": {"name": "litellm_web_search"}, - } - ], - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) + should_run, tools_dict = await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "What's the weather?"}], + tools=[ + { + "type": "function", + "function": {"name": "litellm_web_search"}, + } + ], + stream=False, + custom_llm_provider="openai", + kwargs={}, ) # Verify hook detected the tool call @@ -185,9 +175,7 @@ async def test_websearch_not_triggered_without_tool(): """Test that websearch hook is NOT triggered when no web search tool in request.""" from litellm.types.utils import Choices, Message - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.OPENAI] - ) + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) mock_response = ModelResponse( id="test-123", @@ -208,21 +196,19 @@ async def test_websearch_not_triggered_without_tool(): ) # Test without web search tool - should_run, tools_dict = ( - await websearch_logger.async_should_run_chat_completion_agentic_loop( - response=mock_response, - model="gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - tools=[ - { - "type": "function", - "function": {"name": "some_other_tool"}, - } - ], - stream=False, - custom_llm_provider="openai", - kwargs={}, - ) + should_run, tools_dict = await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + tools=[ + { + "type": "function", + "function": {"name": "some_other_tool"}, + } + ], + stream=False, + custom_llm_provider="openai", + kwargs={}, ) # Verify hook did NOT trigger @@ -241,9 +227,7 @@ async def test_websearch_not_triggered_for_disabled_provider(): ) # Only enable bedrock - websearch_logger = WebSearchInterceptionLogger( - enabled_providers=[LlmProviders.BEDROCK] - ) + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.BEDROCK]) mock_response = ModelResponse( id="test-123", @@ -273,21 +257,19 @@ async def test_websearch_not_triggered_for_disabled_provider(): ) # Test with OpenAI provider (not enabled) - should_run, tools_dict = ( - await websearch_logger.async_should_run_chat_completion_agentic_loop( - response=mock_response, - model="gpt-4o", - messages=[{"role": "user", "content": "test"}], - tools=[ - { - "type": "function", - "function": {"name": "litellm_web_search"}, - } - ], - stream=False, - custom_llm_provider="openai", # Not in enabled_providers - kwargs={}, - ) + should_run, tools_dict = await websearch_logger.async_should_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "test"}], + tools=[ + { + "type": "function", + "function": {"name": "litellm_web_search"}, + } + ], + stream=False, + custom_llm_provider="openai", # Not in enabled_providers + kwargs={}, ) # Verify hook did NOT trigger @@ -341,8 +323,7 @@ async def test_websearch_json_serialization_fix(): @pytest.mark.asyncio @pytest.mark.skipif( - os.environ.get("OPENAI_API_KEY") is None - or os.environ.get("PERPLEXITY_API_KEY") is None, + os.environ.get("OPENAI_API_KEY") is None or os.environ.get("PERPLEXITY_API_KEY") is None, reason="OPENAI_API_KEY or PERPLEXITY_API_KEY not set", ) async def test_websearch_streaming_conversion(): @@ -395,6 +376,174 @@ async def test_websearch_streaming_conversion(): litellm.callbacks = [] +@pytest.mark.asyncio +async def test_maybe_run_chat_completion_agentic_loop_calls_chat_completion_hook(): + """Regression test: maybe_run_chat_completion_agentic_loop must call + async_should_run_chat_completion_agentic_loop, not async_should_run_agentic_loop. + + Before the fix, the function used the wrong gate check and wrong hook, + causing WebSearchInterceptionLogger to never intercept chat completion requests + even when the LLM returned a litellm_web_search tool call. + """ + from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + maybe_run_chat_completion_agentic_loop, + ) + from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Function, + Message, + ) + + mock_response = ModelResponse( + id="test-regression-123", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_abc", + type="function", + function=Function( + name="litellm_web_search", + arguments='{"query": "latest news"}', + ), + ) + ], + ), + ) + ], + model="gpt-4o", + object="chat.completion", + created=1234567890, + ) + + sentinel = ModelResponse( + id="sentinel-final", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(role="assistant", content="Here is the news."), + ) + ], + model="gpt-4o", + object="chat.completion", + created=1234567890, + ) + + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) + + chat_completion_hook_called = False + + async def fake_should_run_chat_completion(response, model, messages, tools, stream, custom_llm_provider, kwargs): + nonlocal chat_completion_hook_called + chat_completion_hook_called = True + return True, { + "tool_calls": [{"id": "call_abc", "name": "litellm_web_search", "input": {"query": "latest news"}}], + "tool_type": "websearch", + "provider": "openai", + "response_format": "openai", + } + + async def fake_build_plan(tools, model, messages, response, optional_params, logging_obj, stream, kwargs): + from litellm.types.integrations.custom_logger import AgenticLoopPlan + + return AgenticLoopPlan(run_agentic_loop=False, response_override=sentinel) + + websearch_logger.async_should_run_chat_completion_agentic_loop = fake_should_run_chat_completion + websearch_logger.async_build_chat_completion_agentic_loop_plan = fake_build_plan + + import litellm as _litellm + + original_callbacks = _litellm.callbacks[:] + _litellm.callbacks = [websearch_logger] + + mock_logging_obj = MagicMock() + mock_logging_obj.dynamic_success_callbacks = None + + try: + result = await maybe_run_chat_completion_agentic_loop( + response=mock_response, + model="gpt-4o", + messages=[{"role": "user", "content": "Latest news?"}], + optional_params={ + "tools": [ + { + "type": "function", + "function": {"name": "litellm_web_search"}, + } + ] + }, + kwargs={}, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + stream=False, + ) + finally: + _litellm.callbacks = original_callbacks + + assert chat_completion_hook_called, ( + "async_should_run_chat_completion_agentic_loop was never called; " + "maybe_run_chat_completion_agentic_loop used the wrong hook" + ) + assert result is sentinel, "Expected agentic loop to return sentinel final response" + + +@pytest.mark.asyncio +async def test_execute_chat_completion_agentic_loop_strips_tool_choice(): + """Regression: _execute_chat_completion_agentic_loop must not forward tool_choice + from the original request into the follow-up synthesis call. + + When the original request forces tool_choice to litellm_web_search, merging + optional_params into the follow-up params without explicit removal causes the + model to call the search tool again instead of synthesizing an answer. + """ + from unittest.mock import patch + + websearch_logger = WebSearchInterceptionLogger(enabled_providers=[LlmProviders.OPENAI]) + + captured_kwargs: dict = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return ModelResponse(id="followup", model="gpt-4o", object="chat.completion") + + async def fake_search(query): + return ("Bitcoin price is $60,000", None) + + with patch.object(websearch_logger, "_execute_search", side_effect=fake_search): + with patch("litellm.acompletion", side_effect=fake_acompletion): + await websearch_logger._execute_chat_completion_agentic_loop( + model="gpt-4o", + messages=[{"role": "user", "content": "What is Bitcoin price?"}], + tool_calls=[ + { + "id": "call_1", + "name": "litellm_web_search", + "input": {"query": "bitcoin price"}, + } + ], + optional_params={ + "tools": [{"type": "function", "function": {"name": "litellm_web_search"}}], + "tool_choice": {"type": "function", "function": {"name": "litellm_web_search"}}, + "max_tokens": 512, + }, + logging_obj=MagicMock(), + stream=False, + kwargs={}, + ) + + assert "tool_choice" not in captured_kwargs, ( + "tool_choice must not appear in follow-up acompletion kwargs; " + "it would force the model to call the search tool again instead of synthesizing" + ) + + if __name__ == "__main__": # Run with: pytest test_websearch_chat_completion.py -v -s pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py index f1196ab4692..cc16ad558e4 100644 --- a/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/test_litellm/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -181,8 +181,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal # The loop must have actually fired (sanity: two provider calls). assert create.await_count == 2, ( - "expected the agentic loop to issue a follow-up provider call; " - f"got {create.await_count} call(s)" + f"expected the agentic loop to issue a follow-up provider call; got {create.await_count} call(s)" ) for idx, call in enumerate(create.await_args_list): @@ -194,8 +193,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal f"top-level request body: {sorted(body.keys())}" ) assert field not in extra_body, ( - f"provider call #{idx}: internal field {field!r} leaked into " - f"extra_body: {sorted(extra_body.keys())}" + f"provider call #{idx}: internal field {field!r} leaked into extra_body: {sorted(extra_body.keys())}" ) # The native code_interpreter tool must have been swapped for the # function tool, never sent raw to OpenAI as a chat-completions request. @@ -254,9 +252,7 @@ class _GateOnlyLogger(CustomLogger): ) -> AgenticLoopPlan: return self._plan - async def async_agentic_loop_cleanup_hook( - self, plan: AgenticLoopPlan, kwargs: Dict[str, Any] - ) -> None: + async def async_agentic_loop_cleanup_hook(self, plan: AgenticLoopPlan, kwargs: Dict[str, Any]) -> None: self.cleanup_calls += 1 @@ -343,9 +339,7 @@ async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messa assert call_kwargs["max_agentic_loops"] >= 1 assert "_agentic_loop_fingerprints" in call_kwargs # Interception markers are mirrored into litellm_metadata for the follow-up. - assert ( - call_kwargs["litellm_metadata"]["_code_interpreter_interception_active"] is True - ) + assert call_kwargs["litellm_metadata"]["_code_interpreter_interception_active"] is True # The transient surface marker is NOT forwarded to the follow-up call. assert "_agentic_loop_api_surface" not in call_kwargs # Cleanup hook always runs. @@ -390,9 +384,7 @@ async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callb # The dispatcher fingerprints the whole value the gate returns as its second # tuple element, so the seeded fingerprint must mirror that dict exactly. - gate_tool_calls = { - "tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}] - } + gate_tool_calls = {"tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}]} fingerprint = json.dumps(gate_tool_calls, sort_keys=True, default=str) logger = _GateOnlyLogger( diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 8c934f9c21e..b18af060a20 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1212,6 +1212,73 @@ def test_async_compact_handler_sends_json_when_not_signed(): assert "data" not in kwargs +@pytest.mark.asyncio +async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(): + """ + Regression: async_anthropic_messages_handler must inject api_key into the + kwargs dict forwarded to _call_agentic_completion_hooks. + + Without this, follow-up calls made by agentic hooks (e.g. websearch + interception's second LLM call after executing searches) have no api_key + and fail with "x-api-key header is required". + """ + handler = BaseLLMHTTPHandler() + + mock_config = Mock() + mock_config.validate_anthropic_messages_environment = Mock( + return_value=({"x-api-key": "sk-test"}, "https://api.anthropic.com") + ) + mock_config.transform_anthropic_messages_request = Mock( + return_value={"model": "claude-haiku", "messages": [], "max_tokens": 16} + ) + mock_config.sign_request = Mock(return_value=({}, None)) + + fake_raw_response = {"id": "msg_1", "type": "message", "role": "assistant", "content": [], "stop_reason": "end_turn"} + mock_config.transform_anthropic_messages_response = Mock(return_value=fake_raw_response) + + mock_logging_obj = Mock() + mock_logging_obj.update_environment_variables = Mock() + mock_logging_obj.model_call_details = {} + mock_logging_obj.stream = False + mock_logging_obj.dynamic_success_callbacks = None + + captured_kwargs: dict = {} + sentinel_response = object() + + async def fake_agentic_hooks(**call_kwargs): + captured_kwargs.update(call_kwargs) + return sentinel_response + + mock_httpx_response = Mock() + mock_httpx_response.status_code = 200 + + with ( + patch.object(handler, "_async_post_anthropic_messages_with_http_error_retry", new=AsyncMock(return_value=mock_httpx_response)), + patch.object(handler, "_call_agentic_completion_hooks", side_effect=fake_agentic_hooks), + patch("litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client"), + patch("litellm.litellm_core_utils.get_provider_specific_headers.ProviderSpecificHeaderUtils.get_provider_specific_headers", return_value=None), + ): + result = await handler.async_anthropic_messages_handler( + model="claude-haiku", + messages=[{"role": "user", "content": "hi"}], + anthropic_messages_provider_config=mock_config, + anthropic_messages_optional_request_params={"stream": False}, + custom_llm_provider="anthropic", + litellm_params=GenericLiteLLMParams(api_key="sk-real-anthropic-key"), + logging_obj=mock_logging_obj, + api_key="sk-real-anthropic-key", + stream=False, + ) + + assert result is sentinel_response + assert "kwargs" in captured_kwargs, "_call_agentic_completion_hooks not called" + forwarded = captured_kwargs["kwargs"] + assert forwarded.get("api_key") == "sk-real-anthropic-key", ( + "api_key must be injected into kwargs passed to _call_agentic_completion_hooks " + "so follow-up calls in agentic hooks (e.g. websearch) can authenticate" + ) + + class _FakeWSExceptions: class WebSocketException(Exception): pass