From b1b7af884abe764c03dece34c8b621b5c0b19a55 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:19:48 -0700 Subject: [PATCH] fix(websearch): forward the deployment api_base to agentic follow-up calls on /v1/messages --- litellm/llms/custom_httpx/llm_http_handler.py | 25 +++--- .../custom_httpx/test_llm_http_handler.py | 90 +++++++++++++++++++ 2 files changed, 104 insertions(+), 11 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 477d10a3cbd..cd76f0d0b54 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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: """ 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 95dceccb2f5..6dc457d26ec 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 @@ -1956,6 +1956,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