diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index a59f1967fc1..0460a4df93c 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,5 +1,6 @@ import asyncio import hashlib +import inspect import json from collections import OrderedDict from collections.abc import Callable, Mapping @@ -782,6 +783,62 @@ async def route_request( ) +class UnroutableModelProbe: + """Asks ``route_request``, on a copy of the request, whether its model can be routed. A model the router + recognizes and an ``a2a/`` agent name count as routable without asking. + + ``refusal`` is the routing error a caller skipped a pre-call guardrail for, and the request must end with it. + """ + + __slots__ = ("_checked", "_llm_router", "_route_type", "_user_api_key_dict", "_user_model", "refusal") + + def __init__( + self, + llm_router: LitellmRouter | None, + user_model: str | None, + route_type: RouteType, + user_api_key_dict: UserAPIKeyAuth | None, + ) -> None: + self._llm_router: Final = llm_router + self._user_model: Final = user_model + self._route_type: Final = route_type + self._user_api_key_dict: Final = user_api_key_dict + self._checked: tuple[object, ProxyModelNotFoundError | None] | None = None + self.refusal: ProxyModelNotFoundError | None = None + + async def refuses(self, data: Mapping[str, object]) -> bool: + model: Final = data.get("model") + if self._checked is None or self._checked[0] != model: + self._checked = (model, await self._refusal(data, model)) + refusal: Final = self._checked[1] + if refusal is not None: + self.refusal = refusal + return refusal is not None + + async def _refusal(self, data: Mapping[str, object], model: object) -> ProxyModelNotFoundError | None: + if _is_a2a_agent_model(model) or ( + isinstance(model, str) + and self._llm_router is not None + and self._llm_router.registry().resolve(model, kind="recognized") + ): + return None + try: + call: Final[object] = await route_request( # pyright: ignore[reportUnknownVariableType] # unawaited provider call of any endpoint + data=dict(data), + llm_router=self._llm_router, + user_model=self._user_model, + route_type=self._route_type, + user_api_key_dict=self._user_api_key_dict, + ) + except ProxyModelNotFoundError as refusal: + return refusal + except Exception: # noqa: BLE001 # any other routing failure is left for the real dispatch to raise + return None + if inspect.iscoroutine(call): # pyright: ignore[reportUnknownArgumentType] # unawaited provider call of any endpoint + call.close() + return None + + async def _route_request_single_attempt( # noqa: ANN202 # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed data: dict, # mutable-ok: request body is the proxy-wide mutable dict contract shared with route_request llm_router: LitellmRouter | None, diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 71f816ea183..1aa03afe0d7 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,6 +1,7 @@ import asyncio import contextlib import gc +import inspect import threading import weakref import httpx @@ -30,6 +31,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.route_llm_request import ( JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, + UnroutableModelProbe, UserConfigRouters, endpoint_unsupported_by_model, route_request, @@ -3061,3 +3063,138 @@ def test_fallback_probe_skips_only_a_loaded_model_that_refuses_the_route() -> No assert skipped("embed-only") is True assert skipped("chat-fallback") is False assert skipped("added-on-another-worker") is False + + +@pytest.mark.asyncio +async def test_route_request_without_a_model_sends_to_the_team_deployment_by_its_public_name(): + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-gpt", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "from-team"}, + "model_info": {"id": "team-gpt-1", "team_id": "team-a", "team_public_model_name": "team-gpt"}, + } + ] + ) + data = { + "model": None, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_team_id": "team-a"}, + } + + llm_call = await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion") + response = await llm_call + + assert response.choices[0].message.content == "from-team" + assert data["model"] == "team-gpt" + + +@pytest.mark.asyncio +async def test_route_request_read_through_sends_through_the_router_it_built(monkeypatch): + """With no router loaded yet, a model the read-through finds is sent through the router the read-through built.""" + import litellm + from litellm.proxy import proxy_server + + model_name = "e2e-first-db-model" + built_router = litellm.Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "fake", + "mock_response": "hello-from-new-router", + }, + } + ] + ) + fake_prisma, _ = _fake_prisma_client_with_models([_db_model_row(model_name, "unused")]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", None) + + async def build_router_from_db(prisma_client, proxy_logging_obj): + monkeypatch.setattr(proxy_server, "llm_router", built_router) + + monkeypatch.setattr(proxy_server.proxy_config, "add_deployment", build_router_from_db) + + llm_call = await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=None, + user_model=None, + route_type="acompletion", + ) + response = await llm_call + + assert response.choices[0].message.content == "hello-from-new-router" + + +def _probe(llm_router: litellm.Router | None = None) -> UnroutableModelProbe: + return UnroutableModelProbe( + llm_router=llm_router, user_model=None, route_type="acompletion", user_api_key_dict=None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("routing_error", "refused"), + [ + pytest.param(ProxyModelNotFoundError(route="/chat/completions", model_name="typo"), True, id="model-not-found"), + pytest.param(HTTPException(status_code=403, detail="model blocked for this key"), False, id="blocked-model"), + pytest.param(ValueError("malformed body"), False, id="other-error"), + ], +) +async def test_unroutable_model_probe_refuses_only_a_model_not_found(monkeypatch, routing_error, refused): + async def failing_route_request(**kwargs): + raise routing_error + + monkeypatch.setattr("litellm.proxy.route_llm_request.route_request", failing_route_request) + probe: Final = _probe() + + assert await probe.refuses({"model": "typo"}) is refused + assert probe.refusal is (routing_error if refused else None) + + +@pytest.mark.asyncio +async def test_unroutable_model_probe_closes_the_call_and_routes_each_model_once(monkeypatch): + routed: Final[list[str]] = [] + calls: Final[list[Coroutine[object, object, None]]] = [] + + async def provider_call() -> None: + return None + + async def routing_route_request(*, data, **kwargs): + routed.append(data["model"]) + data["model"] = "rewritten-by-routing" + calls.append(provider_call()) + return calls[-1] + + monkeypatch.setattr("litellm.proxy.route_llm_request.route_request", routing_route_request) + probe: Final = _probe() + request: Final = {"model": "virtual-model"} + + assert [await probe.refuses(request), await probe.refuses(request)] == [False, False] + assert routed == ["virtual-model"] + assert request == {"model": "virtual-model"} + assert inspect.getcoroutinestate(calls[0]) == inspect.CORO_CLOSED + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-4o", "a2a/support-agent"]) +async def test_unroutable_model_probe_does_not_route_a_loaded_model_or_an_agent( + monkeypatch: pytest.MonkeyPatch, model: str +): + routed: Final[list[object]] = [] + + async def recording_route_request(*, data, **kwargs): + routed.append(data["model"]) + + monkeypatch.setattr("litellm.proxy.route_llm_request.route_request", recording_route_request) + router: Final = litellm.Router( + model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] + ) + + assert await _probe(router).refuses({"model": model}) is False + assert routed == []