From 67ae3fb527350caad8bd4463a5a37c9fbd8a67f4 Mon Sep 17 00:00:00 2001 From: Priyansh Nandwana <136072927+msdianprince-7@users.noreply.github.com> Date: Thu, 29 Oct 2026 17:17:18 +0530 Subject: [PATCH] fix(websearch)!: keep the provider prefix on the agentic follow-up model (#38885) * fix(websearch): keep the provider prefix on the agentic follow-up model The agentic follow-up rebuilt its model string with "a slash means it is already qualified". Handlers are handed the provider-stripped model, and for a provider that routes through a sub-path the remainder still holds a slash, so bedrock/mantle/anthropic.claude-sonnet-5 arrived as mantle/anthropic.claude-sonnet-5 and stayed that way. The follow-up acompletion could not resolve a provider from it and raised "LLM Provider NOT provided". Prefix unless the model is already qualified with that same provider, in one shared helper used by both the handler and the websearch patch builder, which carried copies of the same heuristic. This also fixes openrouter, whose stripped model keeps a slash for the same reason, and the old startswith test that read openai as a prefix of openai_like. Fixes #38829 * test(websearch): cover the agentic follow-up call site The helper had unit tests but the call site did not, so the two changed lines in _execute_chat_completion_agentic_plan were uncovered. Drive the real follow-up with mock_response instead of patching litellm, which keeps the test on the behaviour: against the old code the sub-path case raises "LLM Provider NOT provided", and the ordinary-model case passes either way as a guard. * refactor(websearch): trim the comments on the provider-qualify helper * refactor(websearch): cut the helper and test docstrings to one line * fix(websearch): leave a hook-qualified follow-up model on its own provider A hook that patches in its own model owns that choice, so re-qualifying it sent a cross-provider follow-up to the original request's provider instead. Only a bare patched model and the request's provider-stripped model take the prefix now, and the sibling agentic-loop path resolves the model the same way. * fix(websearch): fall back to the request model when a patch carries no model An empty patched model dispatched a provider-only name instead of falling back, and a provider-stripped sub-path copied into the patch lost its prefix. A patched model now keeps its own provider only when its leading segment names one. * fix(websearch): keep the provider prefix on the responses follow-up model too The responses patch builder still used the old "a slash means already qualified" test, so a provider-stripped sub-path model reached the follow-up unqualified on that surface. It now resolves the model the same way the chat completion builder does. * fix(websearch): read custom_llm_provider without a member call on untyped kwargs * fix(websearch): type the chat patch builder kwargs instead of working around them * test(websearch): run the real search path and flatten the qualify helper tests * fix(websearch): route a hook's model on a JSON-configured provider and keep the request provider on an echoed model * fix(websearch): leave a follow-up model alone when its prefix already routes to the request provider * fix(websearch): send the agentic follow-up to the model the caller requested completion() and responses() record the caller's model before routing strips it, and the chat and Responses follow-up builders prefer that model. This keeps bedrock/converse_like/, bedrock/claude_platform/ and openrouter/openrouter/* deployments on their own route for the follow-up. * test(websearch): cover the recorded caller model in Headroom's retrieval follow-up * test(websearch): cover the recorded caller model in Compresr's retrieval follow-up Also merges the repeated provider-prefix check in core_helpers into one helper and drops a test docstring * test(websearch): drive the real search path through a mocked Perplexity HTTP route * refactor(websearch): read the provider from kwargs inside _caller_model * test(http-handler): type the agentic follow-up provider tests * refactor(websearch): record agentic loop params through a Logging method Moving the recording out of completion() keeps it under basedpyright's code-flow complexity limit after the merge with main. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../websearch_interception/handler.py | 30 +++-- .../chat_completion_agentic_loop.py | 5 +- litellm/litellm_core_utils/core_helpers.py | 43 +++++++ litellm/litellm_core_utils/litellm_logging.py | 4 + litellm/llms/custom_httpx/llm_http_handler.py | 10 +- litellm/main.py | 18 +++ litellm/responses/main.py | 2 + litellm/types/utils.py | 4 +- .../websearch_interception/conftest.py | 18 +++ .../test_websearch_chat_completion.py | 60 +++++++++- .../test_websearch_responses.py | 111 ++++++++++++++++++ .../test_websearch_rich_query_shape.py | 1 + .../test_chat_completion_agentic_loop.py | 43 +++++++ .../litellm_core_utils/test_core_helpers.py | 83 +++++++++++++ .../custom_httpx/test_llm_http_handler.py | 57 +++++++++ .../guardrail_hooks/test_compresr.py | 57 ++++++++- .../guardrail_hooks/test_headroom.py | 68 +++++++++++ 17 files changed, 592 insertions(+), 22 deletions(-) create mode 100644 tests/unit/integrations/websearch_interception/conftest.py diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 6e1ba3eea08..c25fcfe3e70 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -37,6 +37,7 @@ from litellm.integrations.websearch_interception.transformation import ( from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) +from litellm.litellm_core_utils.core_helpers import qualify_provider_stripped_model from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, @@ -114,6 +115,18 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +def _caller_model(model: str, logging_obj: "LiteLLMLoggingObj | None", kwargs: Mapping[str, object]) -> str: + agentic_view: Final[_AgenticLoopParamsView] = { + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) if logging_obj else {} + } + recorded_model: Final = agentic_view["agentic_loop_params"].get("model") + custom_llm_provider: Final = kwargs.get("custom_llm_provider") + return qualify_provider_stripped_model( + recorded_model if isinstance(recorded_model, str) and recorded_model else model, + custom_llm_provider if isinstance(custom_llm_provider, str) else "", + ) + + class _SearchToolLitellmParams(TypedDict, total=False): search_provider: ReadOnly[str | None] @@ -1155,6 +1168,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls=tool_calls, optional_params=optional_params, kwargs=kwargs, + logging_obj=logging_obj, response_format=response_format, ) return AgenticLoopPlan( @@ -1181,6 +1195,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls=tool_calls, optional_params=optional_params, kwargs=kwargs, + logging_obj=logging_obj, ) return AgenticLoopPlan( run_agentic_loop=True, @@ -1195,6 +1210,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: list[dict], optional_params: dict, kwargs: dict, + logging_obj: "LiteLLMLoggingObj | None", ) -> AgenticLoopRequestPatch: """Execute litellm.asearch() and build a Responses API rerun patch.""" search_tasks: Final = [ @@ -1254,9 +1270,7 @@ class WebSearchInterceptionLogger(CustomLogger): } } - full_model_name = model - if "/" not in model and isinstance(kwargs.get("custom_llm_provider"), str): - full_model_name = f"{kwargs['custom_llm_provider']}/{model}" + full_model_name: Final = _caller_model(model, logging_obj, kwargs) verbose_logger.debug( "WebSearchInterception: Built responses request patch model=%s input_items=%d searches=%d", @@ -1814,6 +1828,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls=tool_calls, optional_params=dict[str, object](optional_params), kwargs=dict[str, object](kwargs), + logging_obj=logging_obj, response_format=response_format, ) if request_patch.messages is None: @@ -1836,7 +1851,8 @@ class WebSearchInterceptionLogger(CustomLogger): messages: list[dict], tool_calls: list[dict], optional_params: dict, - kwargs: dict, + kwargs: dict[str, object], + logging_obj: "LiteLLMLoggingObj | None", response_format: str = "openai", ) -> AgenticLoopRequestPatch: """Execute litellm.search() and build chat-completion rerun patch.""" @@ -1922,11 +1938,7 @@ class WebSearchInterceptionLogger(CustomLogger): k: v for k, v in kwargs.items() if not k.startswith("_websearch_interception") and k not in internal_params } - full_model_name = model - if "custom_llm_provider" in kwargs: - custom_llm_provider: Final = kwargs["custom_llm_provider"] - if not model.startswith(custom_llm_provider) and "/" not in model: - full_model_name = f"{custom_llm_provider}/{model}" + full_model_name: Final = _caller_model(model, logging_obj, kwargs) verbose_logger.debug( "WebSearchInterception: Built chat completion request patch model=%s messages=%d", diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 628486f377f..2ba47cbb400 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) +from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.types.integrations.custom_logger import ( @@ -182,9 +183,7 @@ async def _execute_chat_completion_agentic_plan( if patch.messages is None: raise ValueError("Agentic loop plan missing patched messages") - full_model_name = patch.model or model - if "/" not in full_model_name: - full_model_name = f"{custom_llm_provider}/{full_model_name}" + full_model_name: Final = qualify_agentic_followup_model(patch.model, model, custom_llm_provider) optional_params_for_followup: Final = {**optional_params, **patch.optional_params} if patch.tools is not None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 55025508e40..8a53f0cbb42 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -34,6 +34,49 @@ def is_codex_user_agent(user_agent: str) -> bool: return bool(_CODEX_CLIENT_PREFIX_RE.match(user_agent)) +def qualify_provider_stripped_model(model: str, custom_llm_provider: str) -> str: + """Put the provider prefix back on a provider-stripped model.""" + if not custom_llm_provider or model.startswith(f"{custom_llm_provider}/"): + return model + if _already_routes_to(model, custom_llm_provider): + return model + return f"{custom_llm_provider}/{model}" + + +def _already_routes_to(model: str, custom_llm_provider: str) -> bool: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + return _leads_with_provider(model) and inferred_provider(model) == custom_llm_provider + + +def _leads_with_provider(model: str) -> bool: + leading_segment, separator, _ = model.partition("/") + return bool(separator) and names_a_provider(leading_segment) + + +def names_a_provider(candidate: str) -> bool: + """Whether a model name's leading segment is a provider rather than part of the model.""" + import litellm + from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_custom_provider + from litellm.llms.openai_like.json_loader import JSONProviderRegistry + + provider_names: Final = {getattr(provider, "value", provider) for provider in litellm.provider_list} + return ( + candidate in provider_names + or JSONProviderRegistry.exists(candidate) + or is_registered_custom_provider(candidate) + ) + + +def qualify_agentic_followup_model(patch_model: str | None, model: str, custom_llm_provider: str) -> str: + """Resolve the model an agentic follow-up re-dispatches, keeping a hook's own provider.""" + if not patch_model or patch_model == model: + return qualify_provider_stripped_model(model, custom_llm_provider) + if _leads_with_provider(patch_model): + return patch_model + return qualify_provider_stripped_model(patch_model, custom_llm_provider) + + def safe_divide_seconds(seconds: float, denominator: float, default: float | None = None) -> float | None: """ Safely divide seconds by denominator, handling zero division. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7f4b069ded6..77c4136efda 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -165,6 +165,7 @@ from litellm.types.rerank import RerankResponse from litellm.types.router import RouterRateLimitError from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, + AgenticLoopParams, CachingDetails, CallTypes, CostBreakdown, @@ -1158,6 +1159,9 @@ class Logging(LiteLLMLoggingBaseClass): if "custom_llm_provider" in self.model_call_details: self.custom_llm_provider = self.model_call_details["custom_llm_provider"] + def record_agentic_loop_params(self, params: AgenticLoopParams) -> None: + self.model_call_details["agentic_loop_params"] = params + def record_applied_timeout(self, timeout: float | Timeout) -> None: self.applied_timeout = timeout self.litellm_params["timeout"] = timeout diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 9700465ebc3..9e125db4211 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -50,7 +50,11 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) -from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + qualify_agentic_followup_model, + set_provider_response_headers_in_hidden_params, +) from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header @@ -6275,9 +6279,7 @@ class BaseLLMHTTPHandler: if patch.messages is None: raise ValueError("Agentic loop plan missing patched messages") - full_model_name = patch.model or model - if "/" not in full_model_name: - full_model_name = f"{custom_llm_provider}/{full_model_name}" + full_model_name: Final = qualify_agentic_followup_model(patch.model, model, custom_llm_provider) optional_params_for_followup: Final = dict(optional_params) optional_params_for_followup.update(patch.optional_params) diff --git a/litellm/main.py b/litellm/main.py index a3a63d5a83a..d1af2de3a13 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -155,6 +155,7 @@ from litellm.types.litellm_params import ControlOptions, RetryStrategy from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + AgenticLoopParams, BaseUrlRoutingHint, CustomPricingLiteLLMParams, ModelResponseStream, @@ -1272,6 +1273,16 @@ def _get_router_deployment_id(kwargs: dict) -> str | None: return None +def _record_agentic_loop_params( + logging_obj: object, requested_model: str, custom_llm_provider: str, owns_agentic_loop: bool +) -> None: + if not owns_agentic_loop or not isinstance(logging_obj, LiteLLMLoggingObj): + return + logging_obj.record_agentic_loop_params( + AgenticLoopParams(model=requested_model, custom_llm_provider=custom_llm_provider) + ) + + def _register_custom_pricing_for_request( model: str, custom_llm_provider: str, @@ -5691,6 +5702,7 @@ def completion( **{k: kwargs[k] for k in OPTIONAL_KWARGS_KEYS | {"provider_specific_header"} if k in kwargs}, "extra_headers": extra_headers, } + requested_model: Final = model model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, @@ -5702,6 +5714,12 @@ def completion( else None ), ) + _record_agentic_loop_params( + logging_obj=logging, + requested_model=requested_model, + custom_llm_provider=custom_llm_provider, + owns_agentic_loop=not skip_responses_api_bridge, + ) deployment_model_info: Final = str_keyed_mapping( cast(object, model_info) # cast-ok: model_info is the untyped dict kwarg diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5bed6037fa4..11e1602e9f8 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -90,6 +90,7 @@ from litellm.types.responses.streaming_websocket import ResponsesWebSocketReques from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( ADDRESSED_RESPONSE_ID_FIELD, + AgenticLoopParams, CallTypes, CostPerToken, CustomPricingLiteLLMParams, @@ -1580,6 +1581,7 @@ def responses( local_vars=local_vars, use_chat_completions_api=use_chat_completions_api, ) + litellm_logging_obj.record_agentic_loop_params(AgenticLoopParams(model=model)) from litellm.litellm_core_utils.core_helpers import attribute_or_none diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 266a6a20131..9c3329d6387 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -263,10 +263,10 @@ class AgenticLoopParams(TypedDict, total=False): """ model: str - """The model string with provider prefix (e.g., 'bedrock/invoke/...')""" + """The model as the caller requested it, usually with its provider prefix (e.g., 'bedrock/invoke/...')""" custom_llm_provider: str - """The LLM provider name (e.g., 'bedrock', 'anthropic')""" + """The LLM provider name (e.g., 'bedrock', 'anthropic'), when the recorder resolved one""" class OffPeakWindow(TypedDict, total=False): diff --git a/tests/unit/integrations/websearch_interception/conftest.py b/tests/unit/integrations/websearch_interception/conftest.py new file mode 100644 index 00000000000..a9f79a60bfe --- /dev/null +++ b/tests/unit/integrations/websearch_interception/conftest.py @@ -0,0 +1,18 @@ +from typing import Final + +import pytest +import respx + +import litellm + +PERPLEXITY_SEARCH_URL: Final = "https://api.perplexity.ai/search" + + +@pytest.fixture +def perplexity_search(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setenv("PERPLEXITYAI_API_KEY", "test-perplexity-key") + monkeypatch.delenv("PERPLEXITY_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post(PERPLEXITY_SEARCH_URL).respond( + json={"results": [{"title": "AI news", "url": "https://example.com/ai-news", "snippet": "AI news."}]} + ) diff --git a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py index 21e50561f57..fc415374096 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_chat_completion.py @@ -5,15 +5,19 @@ Tests the end-to-end flow of websearch_interception callback with litellm.acompletion() for transparent server-side web search execution. """ -from unittest.mock import MagicMock +import datetime +from typing import Final +from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx import litellm from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) -from litellm.types.utils import LlmProviders, ModelResponse +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import CallTypes, LlmProviders, ModelResponse @pytest.fixture @@ -378,7 +382,6 @@ async def test_execute_chat_completion_agentic_loop_strips_tool_choice(): 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]) @@ -422,3 +425,54 @@ async def test_execute_chat_completion_agentic_loop_strips_tool_choice(): if __name__ == "__main__": # Run with: pytest test_websearch_chat_completion.py -v -s pytest.main([__file__, "-v", "-s"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "requested_model,handler_model,custom_llm_provider", + [ + ("openrouter/openrouter/free", "openrouter/free", "openrouter"), + ("bedrock/converse_like/my-endpoint-model", "my-endpoint-model", "bedrock"), + ], +) +async def test_chat_followup_reuses_the_model_the_caller_requested( + requested_model: str, handler_model: str, custom_llm_provider: str, perplexity_search: respx.Route +) -> None: + """The provider handler sees the model after its routing prefixes are stripped, which + re-qualifying with the provider alone cannot restore.""" + logging_obj: Final = Logging( + model=requested_model, + messages=[], + stream=False, + call_type=CallTypes.completion.value, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-chat-followup", + function_id="fn-chat-followup", + ) + litellm.completion( + model=requested_model, + messages=[{"role": "user", "content": "What's the latest AI news?"}], + mock_response="searching", + litellm_logging_obj=logging_obj, + ) + logger: Final = WebSearchInterceptionLogger() + + plan: Final = await logger.async_build_chat_completion_agentic_loop_plan( + tools={ + "tool_calls": [{"id": "call_abc", "name": "litellm_web_search", "input": {"query": "latest ai news"}}], + "tool_type": "websearch", + "provider": custom_llm_provider, + "response_format": "openai", + }, + model=handler_model, + messages=[{"role": "user", "content": "What's the latest AI news?"}], + response=MagicMock(), + optional_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={"custom_llm_provider": custom_llm_provider}, + ) + + assert plan.request_patch is not None + assert plan.request_patch.model == requested_model + assert perplexity_search.call_count == 1 diff --git a/tests/unit/integrations/websearch_interception/test_websearch_responses.py b/tests/unit/integrations/websearch_interception/test_websearch_responses.py index 00b7299cc31..6ca812a5e40 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_responses.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_responses.py @@ -6,17 +6,22 @@ tool calls returned by /v1/responses, executes the search server-side, and builds a Responses-format follow-up request. """ +import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm + from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) from litellm.types.integrations.custom_logger import ( RESPONSES_AGENTIC_SURFACE, ) +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import CallTypes, LlmProviders @@ -252,6 +257,112 @@ async def test_build_responses_plan_produces_responses_input(): assert patch_obj.model == "openai/gpt-4o" +_WEB_SEARCH_FUNCTION_CALL: Final = { + "tool_calls": [ + { + "id": "fc_1", + "call_id": "fc_1", + "type": "function_call", + "name": "litellm_web_search", + "arguments": '{"query": "latest ai news"}', + "input": {"query": "latest ai news"}, + } + ], + "tool_type": "websearch", + "response_format": "responses", +} + + +async def _responses_followup_model(model: str, custom_llm_provider: str, logging_obj: Logging | None) -> str | None: + logger: Final = WebSearchInterceptionLogger() + plan: Final = await logger.async_build_responses_agentic_loop_plan( + tools={**_WEB_SEARCH_FUNCTION_CALL, "provider": custom_llm_provider}, + model=model, + messages=[{"role": "user", "content": "What's the latest AI news?"}], + response=_responses_output_with_web_search(), + optional_params={"tools": [{"type": "function", "name": "litellm_web_search"}]}, + logging_obj=logging_obj, + stream=False, + kwargs={ + "custom_llm_provider": custom_llm_provider, + "_agentic_loop_api_surface": RESPONSES_AGENTIC_SURFACE, + }, + ) + assert plan.request_patch is not None + return plan.request_patch.model + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("perplexity_search") +async def test_build_responses_plan_keeps_the_provider_on_a_sub_path_model() -> None: + """#38829 on the Responses surface: a provider-stripped model whose remainder still holds a + slash was left unqualified, so the follow-up could not resolve a provider.""" + followup_model: Final = await _responses_followup_model("mantle/anthropic.claude-sonnet-5", "bedrock", None) + + assert followup_model == "bedrock/mantle/anthropic.claude-sonnet-5" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("perplexity_search") +async def test_responses_followup_reuses_the_model_the_caller_requested() -> None: + """An OpenRouter model whose own id starts with "openrouter/" reaches the handler as + "openrouter/free", which re-qualifying alone cannot tell apart from a request for model "free".""" + requested_model: Final = "openrouter/openrouter/free" + logging_obj: Final = Logging( + model=requested_model, + messages=[], + stream=False, + call_type=CallTypes.responses.value, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-responses-followup", + function_id="fn-responses-followup", + ) + litellm.responses( + model=requested_model, + input="What's the latest AI news?", + mock_response="searching", + litellm_logging_obj=logging_obj, + ) + + followup_model: Final = await _responses_followup_model("openrouter/free", "openrouter", logging_obj) + + assert followup_model == requested_model + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("perplexity_search") +async def test_responses_followup_keeps_the_caller_model_across_the_chat_completions_bridge() -> None: + """A Responses request bridged to chat completions re-enters completion() with the stripped model.""" + requested_model: Final = "openrouter/openrouter/free" + logging_obj: Final = Logging( + model=requested_model, + messages=[], + stream=False, + call_type=CallTypes.responses.value, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-bridged-followup", + function_id="fn-bridged-followup", + ) + litellm.responses( + model=requested_model, + input="What's the latest AI news?", + mock_response="searching", + litellm_logging_obj=logging_obj, + ) + litellm.completion( + model="openrouter/free", + custom_llm_provider="openrouter", + messages=[{"role": "user", "content": "What's the latest AI news?"}], + mock_response="searching", + litellm_logging_obj=logging_obj, + _skip_responses_api_bridge=True, + ) + + followup_model: Final = await _responses_followup_model("openrouter/free", "openrouter", logging_obj) + + assert followup_model == requested_model + + @pytest.mark.asyncio async def test_deployment_hook_converts_native_responses_web_search_tool(): """async_pre_call_deployment_hook converts a native Responses web_search tool diff --git a/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py b/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py index 72149e8a435..1020848ec89 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_rich_query_shape.py @@ -239,6 +239,7 @@ class TestCallSiteWiring: tool_calls=tool_calls, optional_params={}, kwargs={}, + logging_obj=None, ) call_kwargs = mock_asearch.await_args.kwargs diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 3cbc3591009..6167e91c37f 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -364,6 +364,49 @@ async def test_dispatcher_runs_followup_with_incremented_depth_and_patched_messa assert logger.cleanup_calls == 1 +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,custom_llm_provider,routed_model", + [ + ("cohere/command-a-03-2025", "cohere_chat", "command-a-03-2025"), + ("anthropic/claude-2", "anthropic_text", "claude-2"), + ], +) +async def test_sdk_followup_keeps_a_model_whose_prefix_already_routes_to_its_provider( + model: str, custom_llm_provider: str, routed_model: str +) -> None: + plan: Final = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages()), + ) + gate: Final = _GateOnlyLogger(plan=plan, tool_calls={"tool_calls": [{"id": "call_abc"}]}) + logging_obj: Final = Logging( + model=model, + messages=[{"role": "user", "content": "what is 6*7?"}], + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-prefix", + function_id="fn-prefix", + dynamic_success_callbacks=[gate], + ) + + followup: Final = await maybe_run_chat_completion_agentic_loop( + response=_tool_call_model_response(), + model=model, + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={"mock_response": "done"}, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + stream=False, + ) + + assert isinstance(followup, ModelResponse) + assert followup.model == routed_model + assert followup._hidden_params["custom_llm_provider"] == custom_llm_provider + + @pytest.mark.asyncio async def test_dispatcher_followup_does_not_repeat_a_request_param_found_in_request_kwargs( restore_callbacks, diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 66a8e6bbda6..a4a3ae285b1 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -26,6 +26,8 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, preserve_upstream_non_openai_attributes, process_response_headers, + qualify_agentic_followup_model, + qualify_provider_stripped_model, reconstruct_model_name, redact_nested_match_and_regex_keys, set_provider_response_headers_in_hidden_params, @@ -637,6 +639,87 @@ class TestProviderResponseHeadersInHiddenParams: assert "additional_headers" not in sibling._hidden_params +@pytest.mark.parametrize( + "model,provider,expected", + [ + ("mantle/anthropic.claude-sonnet-5", "bedrock", "bedrock/mantle/anthropic.claude-sonnet-5"), + ("invoke/anthropic.claude-v2", "bedrock", "bedrock/invoke/anthropic.claude-v2"), + ("openai/gpt-4o", "openrouter", "openrouter/openai/gpt-4o"), + ("gpt-4o", "openai", "openai/gpt-4o"), + ("claude-sonnet-4-5", "anthropic", "anthropic/claude-sonnet-4-5"), + ], +) +def test_the_provider_prefix_is_restored(model: str, provider: str, expected: str) -> None: + assert qualify_provider_stripped_model(model, provider) == expected + + +@pytest.mark.parametrize( + "model,provider", + [ + ("bedrock/mantle/anthropic.claude-sonnet-5", "bedrock"), + ("openai/gpt-4o", "openai"), + ("cohere/command-a-03-2025", "cohere_chat"), + ("ai21/jamba-1.5-large", "ai21_chat"), + ("anthropic/claude-2", "anthropic_text"), + ], +) +def test_an_already_qualified_model_is_left_alone(model: str, provider: str) -> None: + assert qualify_provider_stripped_model(model, provider) == model + + +def test_a_provider_that_only_shares_a_prefix_is_still_qualified() -> None: + assert qualify_provider_stripped_model("openai_like/foo", "openai") == "openai/openai_like/foo" + + +def test_no_provider_leaves_the_model_untouched() -> None: + assert qualify_provider_stripped_model("gpt-4o", "") == "gpt-4o" + + +def test_a_cross_provider_patched_model_is_dispatched_as_the_hook_asked() -> None: + assert qualify_agentic_followup_model("anthropic/claude-sonnet-4-5", "gpt-4o", "openai") == ( + "anthropic/claude-sonnet-4-5" + ) + + +def test_a_bare_patched_model_takes_the_request_provider() -> None: + assert qualify_agentic_followup_model("gpt-4o-mini", "gpt-4o", "openai") == "openai/gpt-4o-mini" + + +def test_a_sub_path_request_model_keeps_its_provider() -> None: + assert qualify_agentic_followup_model(None, "mantle/anthropic.claude-sonnet-5", "bedrock") == ( + "bedrock/mantle/anthropic.claude-sonnet-5" + ) + + +def test_an_unpatched_ordinary_model_takes_the_request_provider() -> None: + assert qualify_agentic_followup_model(None, "gpt-4o", "openai") == "openai/gpt-4o" + + +def test_a_patched_model_already_holding_the_request_provider_is_left_alone() -> None: + assert qualify_agentic_followup_model("openai/gpt-4o", "gpt-4o", "openai") == "openai/gpt-4o" + + +@pytest.mark.parametrize("patched", ["", None]) +def test_a_patch_without_a_model_falls_back_to_the_request_model(patched: str | None) -> None: + assert qualify_agentic_followup_model(patched, "gpt-4o", "openai") == "openai/gpt-4o" + + +def test_a_patched_sub_path_model_still_takes_the_request_provider() -> None: + assert qualify_agentic_followup_model("mantle/anthropic.claude-sonnet-5", "gpt-4o", "bedrock") == ( + "bedrock/mantle/anthropic.claude-sonnet-5" + ) + + +def test_a_patched_model_on_a_json_configured_provider_reaches_that_provider() -> None: + assert qualify_agentic_followup_model("aihubmix/gpt-4o", "gpt-4o", "openai") == "aihubmix/gpt-4o" + + +def test_a_patch_echoing_the_request_model_keeps_the_request_provider() -> None: + assert qualify_agentic_followup_model("openai/gpt-4o", "openai/gpt-4o", "openrouter") == ( + "openrouter/openai/gpt-4o" + ) + + def test_spend_metadata_merge_survives_the_logging_thread_inserting_into_the_request_metadata() -> None: metadata: Final[dict[str, object]] = {} diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 90d1f6772a0..d8326536d03 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5634,6 +5634,63 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) +class TestAgenticFollowUpKeepsTheProviderPrefix: + @staticmethod + async def _run_followup( + model: str, custom_llm_provider: str, patched_model: str | None = None + ) -> ModelResponse: + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + from litellm.types.integrations.custom_logger import ( + AgenticLoopPlan, + AgenticLoopRequestPatch, + ) + + response: Final = await BaseLLMHTTPHandler()._execute_chat_completion_agentic_plan( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[{"role": "user", "content": "hi"}], model=patched_model + ), + ), + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"mock_response": "ok from followup"}, + kwargs={}, + custom_llm_provider=custom_llm_provider, + depth=0, + max_loops=2, + fingerprints=[], + fingerprint="fp", + ) + assert isinstance(response, ModelResponse) + return response + + @pytest.mark.asyncio + async def test_a_sub_path_model_still_resolves_a_provider(self) -> None: + response: Final = await self._run_followup("mantle/anthropic.claude-sonnet-5", "bedrock") + + assert response.choices[0].message.content == "ok from followup" + + @pytest.mark.asyncio + async def test_an_ordinary_model_still_resolves_a_provider(self) -> None: + response: Final = await self._run_followup("gpt-4o-mini", "openai") + + assert response.choices[0].message.content == "ok from followup" + + @pytest.mark.asyncio + async def test_a_hook_that_patches_another_providers_model_reaches_that_provider(self) -> None: + response: Final = await self._run_followup("gpt-4o", "openai", patched_model="anthropic/claude-sonnet-4-5") + + assert response.model == "claude-sonnet-4-5" + assert response._hidden_params["custom_llm_provider"] == "anthropic" + + @pytest.mark.asyncio + async def test_a_hook_echoing_an_org_style_model_stays_on_the_request_provider(self) -> None: + response: Final = await self._run_followup("openai/gpt-4o", "openrouter", patched_model="openai/gpt-4o") + + assert response._hidden_params["custom_llm_provider"] == "openrouter" + + @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_compresr.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_compresr.py index feb7090c2e3..cc5ad10b645 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_compresr.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_compresr.py @@ -15,16 +15,20 @@ Tests cover: - fail_closed raises HTTPException; fail_open forwards uncompressed """ +import datetime import hashlib import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, create_autospec, patch import httpx import pytest from fastapi import HTTPException +import litellm from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import ( COMPRESR_RETRIEVE_TOOL_NAME, CompresrGuardrail, @@ -33,7 +37,7 @@ from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import ( _scoped_store_key, has_compresr_retrieve_tool, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs FAKE_API_BASE = "https://compresr.example.com" FAKE_API_KEY = "cmp_test-key" @@ -1435,6 +1439,57 @@ async def test_agentic_plan_builds_responses_followup_shape( assert output_item["output"] == TOOL_OUTPUT +@pytest.mark.asyncio +async def test_responses_retrieval_followup_reuses_the_model_the_caller_requested( + guardrail: CompresrGuardrail, +) -> None: + requested_model: Final = "bedrock/converse/us.anthropic.claude-sonnet-5-5" + hash_value: Final = "f" * 24 + logging_obj: Final = Logging( + model=requested_model, + messages=[], + stream=False, + call_type=CallTypes.responses.value, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-1", + function_id="fn-compresr-followup", + ) + litellm.responses( + model=requested_model, + input=USER_QUESTION, + mock_response="retrieving", + litellm_logging_obj=logging_obj, + ) + guardrail._store_originals(_scoped_store_key(logging_obj), {hash_value: TOOL_OUTPUT}) + response: Final = MagicMock() + response.output = [{"type": "function_call", "call_id": "fc_1"}] + response.content = None + + plan: Final = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "fc_1", + "type": "function", + "name": COMPRESR_RETRIEVE_TOOL_NAME, + "arguments": {"hash": hash_value}, + } + ] + }, + model="converse/us.anthropic.claude-sonnet-5-5", + messages=[{"role": "user", "content": USER_QUESTION}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={}, + ) + + assert plan.request_patch is not None + assert plan.request_patch.model == requested_model + + @pytest.mark.asyncio async def test_agentic_plan_chat_parallel_tool_calls_echoes_only_retrieve( guardrail: CompresrGuardrail, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 3336b20f5eb..741fea5f2b9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -21,6 +21,7 @@ Tests cover: loop resolves the retrieve tool call, then fake-streamed back to the client """ +import datetime import json import time from typing import Final, Optional @@ -33,6 +34,7 @@ from fastapi import HTTPException import litellm +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( HeadroomGuardrail, has_headroom_retrieve_tool, @@ -771,6 +773,72 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ assert output_item["output"] == original_content +@pytest.mark.asyncio +async def test_responses_retrieval_followup_reuses_the_model_the_caller_requested( + guardrail: HeadroomGuardrail, +) -> None: + requested_model: Final = "bedrock/converse/us.anthropic.claude-sonnet-5-5" + retrieve_hash: Final = "b573993006976af767214fac" + logging_obj: Final = Logging( + model=requested_model, + messages=[], + stream=False, + call_type=CallTypes.responses.value, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-1", + function_id="fn-headroom-followup", + ) + litellm.responses( + model=requested_model, + input=f"What does it say? hash={retrieve_hash}", + mock_response="retrieving", + litellm_logging_obj=logging_obj, + ) + response: Final = MagicMock() + response.choices = None + response.content = None + response.output = [ + { + "type": "function_call", + "id": "fc_abc123", + "call_id": "call_abc123", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": retrieve_hash}), + } + ] + guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({retrieve_hash}), time.monotonic() + 999) + + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response("The full compressed content."), + ): + plan: Final = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": retrieve_hash}, + } + ] + }, + model="converse/us.anthropic.claude-sonnet-5-5", + messages=[{"role": "user", "content": f"What does it say? hash={retrieve_hash}"}], + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + kwargs={"litellm_call_id": "call-1"}, + ) + + assert plan.request_patch is not None + assert plan.request_patch.model == requested_model + + @pytest.mark.asyncio async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messages( guardrail: HeadroomGuardrail,