Merge pull request #41918 from BerriAI/litellm_websearch_followup_api_base

fix(websearch): forward the deployment api_base to agentic follow-up calls on /v1/messages

litellm/llms/custom_httpx/llm_http_handler.py

@@ -2223,6 +2223,7 @@ class BaseLLMHTTPHandler:
 
         # Prepare headers
         kwargs = kwargs or {}
+        kwargs_for_agentic: Final = self._agentic_hook_kwargs(kwargs=kwargs, api_key=api_key, api_base=api_base)
         provider_specific_header: Final = cast(
             litellm.types.utils.ProviderSpecificHeader | Sequence[litellm.types.utils.ProviderSpecificHeader] | None,
             kwargs.get("provider_specific_header", None),
@@ -2410,7 +2411,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, "api_key": api_key} if api_key else kwargs,
+                kwargs=kwargs_for_agentic,
                 hold_back=bool(held_back_tool_names),
                 server_fulfilled_tool_names=held_back_tool_names,
             )
@@ -2433,8 +2434,7 @@ class BaseLLMHTTPHandler:
             anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
             logging_obj=logging_obj,
             custom_llm_provider=custom_llm_provider,
-            api_key=api_key,
-            kwargs=kwargs,
+            kwargs=kwargs_for_agentic,
         )
 
     async def _finalize_anthropic_messages_response(
@@ -2447,14 +2447,8 @@ class BaseLLMHTTPHandler:
         anthropic_messages_optional_request_params: dict,
         logging_obj: LiteLLMLoggingObj,
         custom_llm_provider: str,
-        api_key: str | None,
-        kwargs: dict,
+        kwargs: dict[str, object],
     ) -> AnthropicMessagesResponse | AsyncIterator:
-        # 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: Final = {**kwargs, "api_key": api_key} if api_key else kwargs
-        # Call agentic completion hooks (non-streaming path only)
         final_response: Final = await self._call_agentic_completion_hooks(
             response=initial_response,
             model=model,
@@ -2464,7 +2458,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
             stream=False,
             custom_llm_provider=custom_llm_provider,
-            kwargs=kwargs_for_agentic,
+            kwargs=kwargs,
         )
 
         return self._maybe_wrap_in_fake_stream(
@@ -5312,6 +5306,15 @@ class BaseLLMHTTPHandler:
         fingerprints: Final = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
         return depth, max_loops, fingerprints
 
+    @staticmethod
+    def _agentic_hook_kwargs(
+        kwargs: Mapping[str, object], api_key: str | None, api_base: str | None
+    ) -> dict[str, object]:
+        """``api_key`` and ``api_base`` are named parameters of ``anthropic_messages`` rather than kwargs, so the
+        follow-up call an agentic hook makes only reaches the same deployment if they are re-added here."""
+        deployment_params: Final = {"api_key": api_key, "api_base": api_base}
+        return {**kwargs, **{key: value for key, value in deployment_params.items() if value}}
+
     @staticmethod
     def _has_agentic_completion_hook(logging_obj: LiteLLMLoggingObj) -> bool:
         """

tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py

@@ -1957,6 +1957,96 @@ async def test_async_anthropic_messages_handler_passes_api_key_to_agentic_hooks(
     )
 
 
+_FOUNDRY_API_BASE: Final = "https://lit5418.services.ai.azure.com/anthropic"
+_FOUNDRY_SSE_BODY: Final = (
+    b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1", "type": "message", '
+    b'"role": "assistant", "model": "claude-fable-5-1", "content": [], "stop_reason": null, '
+    b'"usage": {"input_tokens": 1, "output_tokens": 0}}}\n\n'
+    b'event: content_block_start\ndata: {"type": "content_block_start", "index": 0, '
+    b'"content_block": {"type": "text", "text": ""}}\n\n'
+    b'event: content_block_delta\ndata: {"type": "content_block_delta", "index": 0, '
+    b'"delta": {"type": "text_delta", "text": "ready"}}\n\n'
+    b'event: content_block_stop\ndata: {"type": "content_block_stop", "index": 0}\n\n'
+    b'event: message_delta\ndata: {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, '
+    b'"usage": {"output_tokens": 1}}\n\n'
+    b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+)
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.asyncio
+async def test_async_anthropic_messages_handler_passes_deployment_api_base_to_agentic_hooks(stream, monkeypatch):
+    """
+    Regression for LIT-5418: an azure_ai deployment carries its Foundry endpoint as
+    ``api_base``, a named parameter that never lands in kwargs. The agentic hooks
+    (websearch interception's follow-up call after the search) must receive it on
+    both the non-streaming and the streaming path, or the follow-up fails with
+    "Missing Azure API Base" and the client gets the dangling tool_use back.
+    """
+    from litellm.integrations.custom_logger import CustomLogger
+    from litellm.llms.azure_ai.anthropic.messages_transformation import AzureAnthropicMessagesConfig
+
+    monkeypatch.delenv("AZURE_API_BASE", raising=False)
+
+    class CapturingAgenticCallback(CustomLogger):
+        def __init__(self):
+            super().__init__()
+            self.hook_kwargs: dict | None = None
+
+        async def async_should_run_agentic_loop(self, response, model, messages, tools, stream, custom_llm_provider, kwargs):
+            self.hook_kwargs = dict(kwargs)
+            return False, {}
+
+    callback = CapturingAgenticCallback()
+    handler = BaseLLMHTTPHandler()
+    upstream_request = httpx.Request("POST", f"{_FOUNDRY_API_BASE}/v1/messages")
+    upstream_response = (
+        httpx.Response(200, content=_FOUNDRY_SSE_BODY, request=upstream_request)
+        if stream
+        else httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "type": "message",
+                "role": "assistant",
+                "model": "claude-fable-5-1",
+                "content": [{"type": "text", "text": "ready"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 1, "output_tokens": 1},
+            },
+            request=upstream_request,
+        )
+    )
+    mock_client = AsyncMock(spec=AsyncHTTPHandler)
+    mock_client.post = AsyncMock(return_value=upstream_response)
+
+    mock_logging_obj = Mock()
+    mock_logging_obj.model_call_details = {}
+    mock_logging_obj.dynamic_success_callbacks = [callback]
+
+    result = await handler.async_anthropic_messages_handler(
+        model="claude-fable-5-1",
+        messages=[{"role": "user", "content": "Say ready"}],
+        anthropic_messages_provider_config=AzureAnthropicMessagesConfig(),
+        anthropic_messages_optional_request_params={"max_tokens": 32},
+        custom_llm_provider="azure_ai",
+        litellm_params=GenericLiteLLMParams(api_key="foundry-key", api_base=_FOUNDRY_API_BASE),
+        logging_obj=mock_logging_obj,
+        client=mock_client,
+        api_key="foundry-key",
+        api_base=_FOUNDRY_API_BASE,
+        stream=stream,
+        kwargs={},
+    )
+    if stream:
+        _ = [chunk async for chunk in result]
+
+    assert mock_client.post.call_args.kwargs["url"] == f"{_FOUNDRY_API_BASE}/v1/messages"
+    assert callback.hook_kwargs is not None, "agentic hook never ran"
+    assert callback.hook_kwargs.get("api_base") == _FOUNDRY_API_BASE
+    assert callback.hook_kwargs.get("api_key") == "foundry-key"
+
+
 class _FakeWSExceptions:
     class WebSocketException(Exception):
         pass