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,