From fd150ad5991b5be69134288f7da7ade3cde3454d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:39:16 -0700 Subject: [PATCH 01/21] feat(proxy): reject chat routes on models declared non-chat via model_info.supported_endpoints A model whose every deployment sets a non-chat model_info.mode and a supported_endpoints list with no chat endpoint now answers /v1/chat/completions, /v1/messages and /v1/responses with a 400 before any provider call. Models without both fields keep today's behavior. --- litellm/proxy/common_request_processing.py | 10 + litellm/proxy/route_llm_request.py | 63 +++++- .../llm_nonconversational.yaml | 1 + ..._embedding_only_supported_endpoints_e2e.py | 180 ++++++++++++++++++ tests/e2e/models.py | 2 +- .../proxy/test_common_request_processing.py | 62 ++++++ tests/unit/proxy/test_route_llm_request.py | 158 ++++++++++++++- 7 files changed, 473 insertions(+), 3 deletions(-) create mode 100644 tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6b87960fcc0..f7540f5d687 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -119,6 +119,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( + endpoint_unsupported_by_model, route_request, ) from litellm.proxy.utils import ( # noqa: F401 # legacy module exports @@ -2289,6 +2290,15 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj + unsupported_endpoint: Final = endpoint_unsupported_by_model( + route_type=route_type, + model=self.data.get("model"), + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if unsupported_endpoint is not None: + raise unsupported_endpoint + # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set # via litellm_params.guardrails are only honored on post_call paths diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index d8b8a68c8f9..d2593a0f66e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal import httpx from fastapi import HTTPException, status -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -332,6 +332,67 @@ def raise_if_required_body_param_missing( ) +CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( + { + "acompletion": "/v1/chat/completions", + "anthropic_messages": "/v1/messages", + "aresponses": "/v1/responses", + } +) + +CHAT_SHAPED_ENDPOINTS: Final[frozenset[str]] = frozenset( + {"/v1/chat/completions", "/chat/completions", "/v1/messages", "/v1/responses", "/responses"} +) + +CHAT_SHAPED_MODES: Final[frozenset[str]] = frozenset({"chat", "completion", "responses"}) + + +class _DeploymentEndpointDeclaration(BaseModel): + model_config = ConfigDict(frozen=True) + + mode: str | None = None + supported_endpoints: tuple[str, ...] | None = None + + +class ProxyUnsupportedEndpointError(ProxyException): + pass + + +def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None: + try: + declared: Final = _DeploymentEndpointDeclaration.model_validate(deployment.get("model_info") or {}) + except ValidationError: + return None + if declared.mode is None or declared.mode in CHAT_SHAPED_MODES or declared.supported_endpoints is None: + return None + endpoints: Final = frozenset(declared.supported_endpoints) + return endpoints if endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS) else None + + +def endpoint_unsupported_by_model( + route_type: str, + model: object, + llm_router: LitellmRouter | None, + team_id: str | None, +) -> ProxyUnsupportedEndpointError | None: + endpoint: Final = CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE.get(route_type) + if endpoint is None or not isinstance(model, str) or not model or llm_router is None: + return None + deployments: Final = tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) + non_chat: Final = tuple( + endpoints for deployment in deployments if (endpoints := _declared_non_chat_endpoints(deployment)) is not None + ) + if not deployments or len(non_chat) < len(deployments): + return None + supported: Final = ", ".join(sorted(frozenset[str]().union(*non_chat))) or "none" + return ProxyUnsupportedEndpointError( + message=f"{endpoint} is not supported by model={model}. Supported endpoints: {supported}", + type="invalid_request_error", + param="model", + code=status.HTTP_400_BAD_REQUEST, + ) + + class MockTestingParamsDisabledError(HTTPException): def __init__(self, params: tuple[str, ...]): super().__init__( diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index b05e9c916d2..bef555c6374 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -2,6 +2,7 @@ - {id: llm.completions.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: completions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_completions_endpoint_e2e.py", rationale: "Legacy text /completions endpoint, second-highest production request volume"} - {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"} - {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy ยง9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"} +- {id: llm.embeddings.openai.supported_endpoints_gate.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], fail_before_fix: proven, source: "test_embedding_only_supported_endpoints_e2e.py", rationale: "A model whose every deployment lists only non-chat supported_endpoints answers each chat-shaped route 400 at the proxy, naming the endpoint and the declared list, while /embeddings keeps serving"} - {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"} - {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (GitHub issue #36646)"} - {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"} diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py new file mode 100644 index 00000000000..2c496dca822 --- /dev/null +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -0,0 +1,180 @@ +"""Live e2e: an embedding model whose deployments declare supported_endpoints refuses chat-shaped calls. + +Each test registers an OpenAI embedding deployment whose model_info.supported_endpoints lists only +/v1/embeddings, or a control deployment on the same backend that declares nothing. /embeddings serves +real vectors on both. Every chat-shaped route to the declared model answers 400 with the proxy's own +message, while the same chat call to the control goes to OpenAI and comes back with whatever OpenAI +answers, so getting the proxy's message is the proof that no provider call was made. +""" + +from __future__ import annotations + +import dataclasses +from typing import Final + +import anthropic +import openai +import pytest +from e2e_config import unique_marker +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import AnthropicErrorBody, AnthropicErrorEvent, LiteLLMParamsBody, ModelInfoBody, ModelNewBody +from openai import OpenAI +from proxy_client import ProxyClient +from pydantic import BaseModel, ConfigDict +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +EMBEDDING_BACKEND: Final = "openai/text-embedding-3-small" +EMBEDDINGS_ONLY: Final = ("/v1/embeddings",) +PROMPT: Final = "Say this is a test!" +UNSUPPORTED_MARKER: Final = "is not supported by model=" + + +class _OpenAiErrorDetail(BaseModel): + model_config = ConfigDict(frozen=True) + + message: str + type: str + param: str | None = None + + +def _subject(route: Route) -> Subject: + return Subject( + domain=Domain.LLM_TRANSLATION, + route=route, + providers=(Provider.OPENAI,), + models=(EMBEDDING_BACKEND,), + mode=Mode.NONSTREAM, + ) + + +def _register( + proxy: ProxyClient, resources: ResourceManager, prefix: str, supported_endpoints: tuple[str, ...] | None +) -> str: + model: Final = f"{prefix}-{unique_marker()}" + model_id: Final = proxy.register_model( + ModelNewBody( + model_name=model, + litellm_params=LiteLLMParamsBody(model=EMBEDDING_BACKEND, api_key="os.environ/OPENAI_API_KEY"), + model_info=ModelInfoBody( + mode="embedding", + supported_endpoints=None if supported_endpoints is None else list(supported_endpoints), + ), + ) + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + +def _openai_at(sdk: SdkClients, api_prefix: str, key: str) -> OpenAI: + return dataclasses.replace(sdk, base_url=f"{sdk.base_url}{api_prefix}").openai(key) + + +def _unsupported_message(endpoint: str, model: str) -> str: + return f"{endpoint} {UNSUPPORTED_MARKER}{model}. Supported endpoints: {', '.join(sorted(EMBEDDINGS_ONLY))}" + + +def _chat_error_message(client: OpenAI, model: str) -> str: + try: + _ = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE + ) + except openai.APIStatusError as error: + return f"{error.status_code} {error.message}" + return "" + + +def _assert_openai_rejection(error: openai.APIStatusError, endpoint: str, model: str) -> None: + assert error.status_code == 400, f"expected the proxy's 400, got {error.status_code}: {error.message}" + assert _OpenAiErrorDetail.model_validate(error.body) == _OpenAiErrorDetail( + message=_unsupported_message(endpoint, model), type="invalid_request_error", param="model" + ), f"unexpected error body: {error.body!r}" + + +class TestEmbeddingOnlySupportedEndpoints: + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.EMBEDDINGS)) + @pytest.mark.parametrize( + ("prefix", "supported_endpoints"), + ( + pytest.param("e2e-embed-only", EMBEDDINGS_ONLY, id="embed-only"), + pytest.param("e2e-embed-control", None, id="control"), + ), + ) + def test_embeddings_return_a_vector( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + supported_endpoints: tuple[str, ...] | None, + ) -> None: + model: Final = _register(proxy, resources, prefix, supported_endpoints) + embeddings: Final = sdk.openai(resources.key()).embeddings.create( + model=model, input=PROMPT, extra_body=NO_PROXY_CACHE + ) + assert embeddings.data, f"/embeddings returned no data for {model}: {embeddings!r}" + assert any(component != 0.0 for component in embeddings.data[0].embedding), ( + f"/embeddings returned an empty or all-zero vector for {model}" + ) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.CHAT_COMPLETIONS)) + @pytest.mark.parametrize("api_prefix", ("/v1", ""), ids=("v1-chat-completions", "chat-completions")) + def test_chat_completions_are_rejected_before_any_provider_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, api_prefix: str + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + with pytest.raises(openai.APIStatusError) as raised: + _ = _openai_at(sdk, api_prefix, resources.key()).chat.completions.create( + model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE + ) + _assert_openai_rejection(raised.value, "/v1/chat/completions", model) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.RESPONSES)) + @pytest.mark.parametrize("api_prefix", ("/v1", ""), ids=("v1-responses", "responses")) + def test_responses_are_rejected_before_any_provider_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, api_prefix: str + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + with pytest.raises(openai.APIStatusError) as raised: + _ = _openai_at(sdk, api_prefix, resources.key()).responses.create( + model=model, input=PROMPT, extra_body=NO_PROXY_CACHE + ) + _assert_openai_rejection(raised.value, "/v1/responses", model) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.MESSAGES)) + def test_messages_are_rejected_before_any_provider_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + with pytest.raises(anthropic.APIStatusError) as raised: + _ = sdk.anthropic(resources.key()).messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": PROMPT}], + extra_body=NO_PROXY_CACHE, + ) + assert raised.value.status_code == 400, ( + f"expected the proxy's 400, got {raised.value.status_code}: {raised.value.message}" + ) + assert AnthropicErrorEvent.model_validate(raised.value.body) == AnthropicErrorEvent( + type="error", + error=AnthropicErrorBody(type="invalid_request_error", message=_unsupported_message("/v1/messages", model)), + ), f"unexpected Anthropic error envelope: {raised.value.body!r}" + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.CHAT_COMPLETIONS)) + def test_chat_completions_without_supported_endpoints_reach_the_provider( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-control", None) + error_message: Final = _chat_error_message(_openai_at(sdk, "/v1", resources.key()), model) + assert UNSUPPORTED_MARKER not in error_message, ( + f"a model that declares no supported_endpoints must reach the provider, but the proxy " + f"rejected it: {error_message}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 5b1c5778c10..cc1614b048f 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1413,7 +1413,7 @@ class LiteLLMParamsBody(BaseModel): order: int | None = None -ModelMode = Literal["chat", "batch", "realtime", "image_generation"] +ModelMode = Literal["chat", "embedding", "batch", "realtime", "image_generation"] class ModelInfoBody(BaseModel): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 3eb542f223e..1d7ccf90650 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -72,6 +72,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_spend_cache_key, team_member_budget_entity_id, ) +from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment @@ -711,6 +712,67 @@ class TestProxyBaseLLMRequestProcessing: assert returned_data["metadata"]["tags"] == ["existing-tag"] tag_budget_check.assert_not_awaited() + @staticmethod + def _router_with_embedding_only_model() -> Router: + return Router( + model_list=[ + { + "model_name": "embed-only", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-test"}, + "model_info": {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, + }, + ] + ) + + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_rejects_chat_on_embedding_only_model_before_pre_call_hooks( + self, monkeypatch + ): + processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( + monkeypatch, + request_body={"model": "embed-only", "messages": [{"role": "user", "content": "hello"}]}, + guardrail_tags=[], + ) + + with pytest.raises(ProxyUnsupportedEndpointError) as exc_info: + await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="acompletion", + llm_router=self._router_with_embedding_only_model(), + ) + + assert exc_info.value.code == "400" + assert exc_info.value.message == ( + "/v1/chat/completions is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + ) + mock_proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_common_processing_pre_call_logic_lets_embeddings_through_on_embedding_only_model(self, monkeypatch): + processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( + monkeypatch, + request_body={"model": "embed-only", "input": ["hello"]}, + guardrail_tags=[], + route="/v1/embeddings", + ) + + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=mock_proxy_config, + route_type="aembedding", + llm_router=self._router_with_embedding_only_model(), + ) + + assert returned_data["model"] == "embed-only" + mock_proxy_logging_obj.pre_call_hook.assert_awaited_once() + @pytest.mark.asyncio async def test_common_processing_pre_call_logic_skips_guardrail_tag_budget_check_for_zero_cost_model( self, monkeypatch diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index a4b34825d95..a82067ef310 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -7,7 +7,13 @@ from unittest.mock import MagicMock from fastapi import HTTPException -from litellm.proxy.route_llm_request import ProxyModelNotFoundError, route_request +from litellm import Router +from litellm.proxy.route_llm_request import ( + ProxyModelNotFoundError, + ProxyUnsupportedEndpointError, + endpoint_unsupported_by_model, + route_request, +) @pytest.mark.parametrize( @@ -1760,3 +1766,153 @@ async def test_route_request_router_settings_override_skips_null_fields(): assert "context_window_fallbacks" not in call_kwargs assert "num_retries" not in call_kwargs assert "model_group_retry_policy" not in call_kwargs + + +CHAT_SHAPED_ROUTES: Final = ( + ("acompletion", "/v1/chat/completions"), + ("anthropic_messages", "/v1/messages"), + ("aresponses", "/v1/responses"), +) + +CHAT_SHAPED_ROUTE_TYPES: Final = tuple(route_type for route_type, _ in CHAT_SHAPED_ROUTES) + +EMBEDDING_ONLY: Final = {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]} + + +def _embedding_deployment(model_name: str, model_info: dict[str, object]) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": "sagemaker/openai/embed-endpoint", "aws_region_name": "us-east-1"}, + "model_info": model_info, + } + + +def _router_with(*deployments: dict[str, object], **router_kwargs: object) -> Router: + return Router(model_list=list(deployments), **router_kwargs) + + +def _unsupported(route_type: str, model: object, router: Router | None) -> ProxyUnsupportedEndpointError | None: + return endpoint_unsupported_by_model(route_type=route_type, model=model, llm_router=router, team_id=None) + + +@pytest.mark.parametrize("route_type, endpoint", CHAT_SHAPED_ROUTES) +def test_endpoint_unsupported_rejects_chat_routes_on_embedding_only_model(route_type, endpoint): + error: Final = _unsupported( + route_type, "embed-only", _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) + ) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert error.code == "400" + assert error.param == "model" + assert error.type == "invalid_request_error" + assert error.message == f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + + +@pytest.mark.parametrize("route_type", ["aembedding", "atext_completion", "aimage_generation"]) +def test_endpoint_unsupported_leaves_non_chat_routes_alone(route_type): + router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) + + assert _unsupported(route_type, "embed-only", router) is None + + +@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) +@pytest.mark.parametrize( + "model_info", + [ + {"mode": "embedding"}, + {}, + {"mode": "embedding", "supported_endpoints": None}, + {"mode": "embedding", "supported_endpoints": "/v1/embeddings"}, + {"mode": 1, "supported_endpoints": ["/v1/embeddings"]}, + ], +) +def test_endpoint_unsupported_keeps_todays_behavior_without_a_valid_declaration(route_type, model_info): + router: Final = _router_with(_embedding_deployment("embed-model", model_info)) + + assert _unsupported(route_type, "embed-model", router) is None + + +@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) +@pytest.mark.parametrize( + "model_info", + [ + {"supported_endpoints": ["/v1/embeddings"]}, + {"mode": "chat", "supported_endpoints": ["/v1/batch"]}, + {"mode": "completion", "supported_endpoints": ["/fine_tuning"]}, + {"mode": "responses", "supported_endpoints": ["/v1/realtime"]}, + ], +) +def test_endpoint_unsupported_needs_a_non_chat_mode_to_reject(route_type, model_info): + router: Final = _router_with(_embedding_deployment("advisory", model_info)) + + assert _unsupported(route_type, "advisory", router) is None + + +@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) +@pytest.mark.parametrize( + "declared", + [["/v1/messages"], ["/v1/responses"], ["/v1/chat/completions"], ["/chat/completions"], ["/responses"]], +) +def test_endpoint_unsupported_allows_every_chat_route_when_any_chat_endpoint_is_declared(route_type, declared): + router: Final = _router_with( + _embedding_deployment("chat-model", {"mode": "embedding", "supported_endpoints": declared}) + ) + + assert _unsupported(route_type, "chat-model", router) is None + + +def test_endpoint_unsupported_allows_a_group_with_one_deployment_that_declares_nothing(): + router: Final = _router_with( + _embedding_deployment("mixed", EMBEDDING_ONLY), + _embedding_deployment("mixed", {"mode": "embedding"}), + ) + + assert _unsupported("acompletion", "mixed", router) is None + + +def test_endpoint_unsupported_names_every_endpoint_the_group_declares(): + router: Final = _router_with( + _embedding_deployment("embed-group", EMBEDDING_ONLY), + _embedding_deployment( + "embed-group", {"mode": "embedding", "supported_endpoints": ["/v1/rerank", "/v1/embeddings"]} + ), + ) + + error: Final = _unsupported("acompletion", "embed-group", router) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert error.message == ( + "/v1/chat/completions is not supported by model=embed-group. Supported endpoints: /v1/embeddings, /v1/rerank" + ) + + +def test_endpoint_unsupported_says_none_for_an_empty_declared_list(): + router: Final = _router_with(_embedding_deployment("nothing", {"mode": "embedding", "supported_endpoints": []})) + + error: Final = _unsupported("aresponses", "nothing", router) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert error.message == "/v1/responses is not supported by model=nothing. Supported endpoints: none" + + +def test_endpoint_unsupported_follows_a_model_group_alias(): + router: Final = _router_with( + _embedding_deployment("embed-only", EMBEDDING_ONLY), + model_group_alias={"embed-alias": "embed-only"}, + ) + + error: Final = _unsupported("anthropic_messages", "embed-alias", router) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert error.message == "/v1/messages is not supported by model=embed-alias. Supported endpoints: /v1/embeddings" + + +@pytest.mark.parametrize("model", ["unknown-model", "", None]) +def test_endpoint_unsupported_ignores_models_without_deployments(model): + router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) + + assert _unsupported("acompletion", model, router) is None + + +def test_endpoint_unsupported_skips_the_check_without_a_router(): + assert _unsupported("acompletion", "embed-only", None) is None From 0264476d7ff071d0be1bb5987e46981d58bc9daf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 18:37:55 -0700 Subject: [PATCH 02/21] fix(proxy): let fallbacks answer and cover more chat routes in the non-chat model check The check now steps aside when a request, key, team or router fallback would answer, skips such models in the rate-limit fallback loop, ignores mode and supported_endpoints copied from the cost map by old Admin UI saves, treats a blank mode as unset, and also covers /v1/completions, the Responses WebSocket, Gemini generateContent and /v1beta/interactions --- litellm/proxy/common_request_processing.py | 5 +- .../proxy/response_api_endpoints/endpoints.py | 10 +- litellm/proxy/route_llm_request.py | 114 ++- ..._embedding_only_supported_endpoints_e2e.py | 42 +- .../test_non_chat_model_endpoint_gate.py | 883 ++++++++++++++++++ .../response_api_endpoints/test_endpoints.py | 83 ++ .../proxy/test_common_request_processing.py | 51 + tests/unit/proxy/test_route_llm_request.py | 97 +- 8 files changed, 1267 insertions(+), 18 deletions(-) create mode 100644 tests/integration/routing/test_non_chat_model_endpoint_gate.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index f7540f5d687..b2393a39ad5 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -119,6 +119,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( + ProxyUnsupportedEndpointError, endpoint_unsupported_by_model, route_request, ) @@ -2292,7 +2293,7 @@ class ProxyBaseLLMRequestProcessing: unsupported_endpoint: Final = endpoint_unsupported_by_model( route_type=route_type, - model=self.data.get("model"), + data=self.data, llm_router=llm_router, team_id=user_api_key_dict.team_id, ) @@ -2477,6 +2478,8 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, rate_limited_model=original_model, ) + except ProxyUnsupportedEndpointError: + continue except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): raise diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7f9fcca369b..c6d5221c275 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,7 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) -from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.route_llm_request import endpoint_unsupported_by_model, raise_if_required_body_param_missing from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1546,6 +1546,14 @@ async def _responses_websocket_session( user_api_key_dict=user_api_key_dict, llm_router=llm_router, ) + unsupported_endpoint: Final = endpoint_unsupported_by_model( + route_type="aresponses", + data=data, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if unsupported_endpoint is not None: + raise unsupported_endpoint ( data, litellm_logging_obj, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 23b20f017cf..087a7f56d7a 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,16 +1,25 @@ import asyncio -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol import httpx from fastapi import HTTPException, status from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request +from litellm.router_utils.fallback_event_handlers import ( + check_non_standard_fallback_format, + fallback_lookup_groups, + fallbacks_disabled_for_request, + get_fallback_model_group_for_lookup_groups, + has_unattempted_fallback_target, +) +from litellm.types.utils import COST_MAP_LOOKUP_KEY, echoed_cost_map_fields # Client-supplied params that make the router or the call path fabricate a # failure or a delay instead of calling the provider. The ``mock_testing_*`` @@ -359,15 +368,30 @@ def raise_if_required_body_param_missing( CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( { "acompletion": "/v1/chat/completions", + "atext_completion": "/v1/completions", "anthropic_messages": "/v1/messages", "aresponses": "/v1/responses", + "agenerate_content": "generateContent", + "agenerate_content_stream": "streamGenerateContent", + "acreate_interaction": "/v1beta/interactions", } ) CHAT_SHAPED_ENDPOINTS: Final[frozenset[str]] = frozenset( - {"/v1/chat/completions", "/chat/completions", "/v1/messages", "/v1/responses", "/responses"} + { + "/v1/chat/completions", + "/chat/completions", + "/v1/completions", + "/completions", + "/v1/messages", + "/v1/responses", + "/responses", + "/v1beta/interactions", + } ) +DECLARATION_FIELDS: Final[frozenset[str]] = frozenset({"mode", "supported_endpoints"}) + CHAT_SHAPED_MODES: Final[frozenset[str]] = frozenset({"chat", "completion", "responses"}) @@ -382,31 +406,107 @@ class ProxyUnsupportedEndpointError(ProxyException): pass +_FALLBACK_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter( # mutable-ok: router fallbacks are lists + list[object], config=ConfigDict(strict=True) +) + + +class _LiveCostMap(Protocol): + @property + def model_cost(self) -> Mapping[str, object]: ... + + +class _RouterFallbacks(Protocol): + @property + def fallbacks(self) -> Sequence[object] | None: ... + + +def _live_cost_map_entry(catalog: _LiveCostMap, catalog_key: str) -> Mapping[str, object]: + try: + return JSON_OBJECT_ADAPTER.validate_python(catalog.model_cost.get(catalog_key) or {}) + except ValidationError: + return MappingProxyType({}) + + +def _echoed_from_cost_map(model_info: Mapping[str, object]) -> frozenset[str]: + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + if not isinstance(catalog_key, str): + return frozenset() + return frozenset( + echoed_cost_map_fields( + model_info, + _live_cost_map_entry(litellm, catalog_key), + GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}), + ) + ) + + def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None: try: - declared: Final = _DeploymentEndpointDeclaration.model_validate(deployment.get("model_info") or {}) + model_info: Final = JSON_OBJECT_ADAPTER.validate_python(deployment.get("model_info")) + declared: Final = _DeploymentEndpointDeclaration.model_validate(model_info) except ValidationError: return None - if declared.mode is None or declared.mode in CHAT_SHAPED_MODES or declared.supported_endpoints is None: + if not DECLARATION_FIELDS.isdisjoint(_echoed_from_cost_map(model_info)): + return None + mode: Final = (declared.mode or "").strip().lower() + if not mode or mode in CHAT_SHAPED_MODES or declared.supported_endpoints is None: return None endpoints: Final = frozenset(declared.supported_endpoints) return endpoints if endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS) else None +def _override_fallbacks(data: Mapping[str, object]) -> object: + try: + return JSON_OBJECT_ADAPTER.validate_python(data.get("router_settings_override")).get("fallbacks") + except ValidationError: + return None + + +def _fallbacks_for_request(data: Mapping[str, object], llm_router: LitellmRouter) -> object: + if "fallbacks" in data: + return data["fallbacks"] + override_fallbacks: Final = _override_fallbacks(data) + return override_fallbacks if override_fallbacks is not None else _router_fallbacks(llm_router) + + +def _router_fallbacks(llm_router: _RouterFallbacks) -> Sequence[object] | None: + return llm_router.fallbacks + + +def _fallback_would_answer(model: str, data: Mapping[str, object], llm_router: LitellmRouter) -> bool: + if fallbacks_disabled_for_request(data): + return False + try: + fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(_fallbacks_for_request(data, llm_router)) + except ValidationError: + return False + if not fallbacks: + return False + if check_non_standard_fallback_format(fallbacks=fallbacks): + return True + resolved, _ = get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, + lookup_groups=fallback_lookup_groups(data, model), + ) + return has_unattempted_fallback_target(resolved, data) + + def endpoint_unsupported_by_model( route_type: str, - model: object, + data: Mapping[str, object], llm_router: LitellmRouter | None, team_id: str | None, ) -> ProxyUnsupportedEndpointError | None: endpoint: Final = CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE.get(route_type) + model: Final = data.get("model") if endpoint is None or not isinstance(model, str) or not model or llm_router is None: return None deployments: Final = tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) non_chat: Final = tuple( endpoints for deployment in deployments if (endpoints := _declared_non_chat_endpoints(deployment)) is not None ) - if not deployments or len(non_chat) < len(deployments): + if not deployments or len(non_chat) < len(deployments) or _fallback_would_answer(model, data, llm_router): return None supported: Final = ", ".join(sorted(frozenset[str]().union(*non_chat))) or "none" return ProxyUnsupportedEndpointError( diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py index 2c496dca822..9a89344cebd 100644 --- a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -18,7 +18,7 @@ import pytest from e2e_config import unique_marker from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta from lifecycle import ResourceManager -from models import AnthropicErrorBody, AnthropicErrorEvent, LiteLLMParamsBody, ModelInfoBody, ModelNewBody +from models import AnthropicErrorBody, AnthropicErrorEvent, LiteLLMParamsBody, ModelInfoBody, ModelMode, ModelNewBody from openai import OpenAI from proxy_client import ProxyClient from pydantic import BaseModel, ConfigDict @@ -27,6 +27,7 @@ from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e EMBEDDING_BACKEND: Final = "openai/text-embedding-3-small" +CHAT_FALLBACK_BACKEND: Final = "openai/gpt-5.4-nano" EMBEDDINGS_ONLY: Final = ("/v1/embeddings",) PROMPT: Final = "Say this is a test!" UNSUPPORTED_MARKER: Final = "is not supported by model=" @@ -51,15 +52,20 @@ def _subject(route: Route) -> Subject: def _register( - proxy: ProxyClient, resources: ResourceManager, prefix: str, supported_endpoints: tuple[str, ...] | None + proxy: ProxyClient, + resources: ResourceManager, + prefix: str, + supported_endpoints: tuple[str, ...] | None, + backend: str = EMBEDDING_BACKEND, + mode: ModelMode = "embedding", ) -> str: model: Final = f"{prefix}-{unique_marker()}" model_id: Final = proxy.register_model( ModelNewBody( model_name=model, - litellm_params=LiteLLMParamsBody(model=EMBEDDING_BACKEND, api_key="os.environ/OPENAI_API_KEY"), + litellm_params=LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY"), model_info=ModelInfoBody( - mode="embedding", + mode=mode, supported_endpoints=None if supported_endpoints is None else list(supported_endpoints), ), ) @@ -167,6 +173,34 @@ class TestEmbeddingOnlySupportedEndpoints: error=AnthropicErrorBody(type="invalid_request_error", message=_unsupported_message("/v1/messages", model)), ), f"unexpected Anthropic error envelope: {raised.value.body!r}" + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.COMPLETIONS)) + def test_text_completions_are_rejected_before_any_provider_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + with pytest.raises(openai.APIStatusError) as raised: + _ = _openai_at(sdk, "/v1", resources.key()).completions.create( + model=model, prompt=PROMPT, max_tokens=16, extra_body=NO_PROXY_CACHE + ) + _assert_openai_rejection(raised.value, "/v1/completions", model) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.CHAT_COMPLETIONS)) + def test_chat_completions_with_a_fallback_are_answered_by_the_fallback( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + fallback: Final = _register(proxy, resources, "e2e-chat-fallback", None, CHAT_FALLBACK_BACKEND, "chat") + completion: Final = _openai_at(sdk, "/v1", resources.key()).chat.completions.create( + model=model, + messages=[{"role": "user", "content": PROMPT}], + extra_body={**NO_PROXY_CACHE, "fallbacks": [fallback]}, + ) + assert completion.choices and completion.choices[0].message.content, ( + f"the fallback {fallback} should answer a chat call to {model}, got {completion!r}" + ) + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") @meta(_subject(Route.CHAT_COMPLETIONS)) def test_chat_completions_without_supported_endpoints_reach_the_provider( diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py new file mode 100644 index 00000000000..b12b42ec638 --- /dev/null +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -0,0 +1,883 @@ +import asyncio +import json +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + list_value, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse + +_PROVIDER_KEY: Final = "integration-provider-key" +_EMBEDDING_MODEL: Final = "text-embedding-3-small" +_CHAT_MODEL: Final = "gpt-4o-mini" +_SERVED_TEXT: Final = "served by the chat upstream" +_VECTOR: Final = (0.25, 0.5, 0.75) +_CLOSE: Final = MappingProxyType({"connection": "close"}) +_RELOAD_SECONDS: Final = 5 +_LONG_ENDPOINT: Final = "/v1/embeddings/" + "x" * 5000 +_EMBED_ONLY: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]} +) + + +def _embedding_reply() -> Reply: + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(_VECTOR)}], + "model": _EMBEDDING_MODEL, + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } + ).encode() + ) + + +def _not_a_chat_model() -> Reply: + return Reply( + status=404, + body=json.dumps( + {"error": {"message": "scripted: not a chat model", "type": "invalid_request_error", "code": None}} + ).encode(), + ) + + +def _text_completion_reply() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"cmpl-{uuid.uuid4().hex}", + "object": "text_completion", + "created": 1, + "model": _CHAT_MODEL, + "choices": [{"index": 0, "text": _SERVED_TEXT, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, + } + ).encode() + ) + + +def _respond(request: Request) -> Reply: + path: Final = request.target.split("?")[0] + body: Final = JSON_OBJECT.validate_json(request.body or b"{}") + streamed: Final = body.get("stream") is True + if path == "/v1/embeddings": + return _embedding_reply() + if body.get("model") != _CHAT_MODEL: + return _not_a_chat_model() + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path == "/v1/responses": + return responses_reply(f"resp_{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path == "/v1/completions": + return _text_completion_reply() + return _not_a_chat_model() + + +class _Upstream: + def __init__(self, wire: Wire) -> None: + self.wire = wire + self.seen: tuple[Request, ...] = () + + def calls(self, marker: str) -> tuple[Request, ...]: + self.seen = (*self.seen, *self.wire.drain()) + return tuple(request for request in self.seen if marker.encode() in request.body) + + def targets(self, marker: str) -> tuple[str, ...]: + return tuple(request.target for request in self.calls(marker)) + + +@dataclass(frozen=True, slots=True) +class _Deployment: + litellm_model: str + declarations: tuple[Mapping[str, JsonValue], ...] + + +def _single(declaration: Mapping[str, JsonValue], litellm_model: str = f"openai/{_CHAT_MODEL}") -> _Deployment: + return _Deployment(litellm_model, (declaration,)) + + +_GATED: Final[Mapping[str, tuple[_Deployment, str]]] = MappingProxyType( + { + "image_generation": ( + _single({"mode": "image_generation", "supported_endpoints": ["/v1/images/generations"]}), + "/v1/images/generations", + ), + "rerank_two_spellings": ( + _single({"mode": "rerank", "supported_endpoints": ["/v1/rerank", "/rerank"]}), + "/rerank, /v1/rerank", + ), + "empty_endpoints": (_single({"mode": "embedding", "supported_endpoints": []}), "none"), + "duplicated_endpoints": ( + _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/embeddings", "/embeddings"]}), + "/embeddings, /v1/embeddings", + ), + "trailing_slash_chat_spelling": ( + _single({"mode": "embedding", "supported_endpoints": ["/v1/chat/completions/"]}), + "/v1/chat/completions/", + ), + "five_kb_endpoint": (_single({"mode": "embedding", "supported_endpoints": [_LONG_ENDPOINT]}), _LONG_ENDPOINT), + "two_embedding_deployments": ( + _Deployment( + f"openai/{_CHAT_MODEL}", + (_EMBED_ONLY, {"mode": "embedding", "supported_endpoints": ["/embeddings"]}), + ), + "/embeddings, /v1/embeddings", + ), + } +) + +_UNGATED: Final[Mapping[str, _Deployment]] = MappingProxyType( + { + "mode_only": _single({"mode": "embedding"}), + "endpoints_only": _single({"supported_endpoints": ["/v1/embeddings"]}), + "endpoints_null": _single({"mode": "embedding", "supported_endpoints": None}), + "mode_null": _single({"mode": None, "supported_endpoints": ["/v1/embeddings"]}), + "mode_blank": _single({"mode": "", "supported_endpoints": ["/v1/embeddings"]}), + "endpoints_a_string": _single({"mode": "embedding", "supported_endpoints": "/v1/embeddings"}), + "endpoints_not_strings": _single({"mode": "embedding", "supported_endpoints": [1, 2]}), + "chat_mode": _single({"mode": "chat", "supported_endpoints": ["/v1/embeddings"]}), + "completion_mode": _single({"mode": "completion", "supported_endpoints": ["/v1/embeddings"]}), + "responses_mode": _single({"mode": "responses", "supported_endpoints": ["/v1/embeddings"]}), + "lists_v1_chat": _single( + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/chat/completions"]} + ), + "lists_bare_chat": _single( + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/chat/completions"]} + ), + "lists_messages": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/messages"]}), + "lists_bare_responses": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/responses"]}), + "mixed_group": _Deployment(f"openai/{_CHAT_MODEL}", (_EMBED_ONLY, {"mode": "chat"})), + } +) + +_SHARED: Final[Mapping[str, _Deployment]] = MappingProxyType( + { + "embedding_only": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "chat_target": _single({"mode": "chat"}), + **{name: deployment for name, (deployment, _supported) in _GATED.items()}, + **_UNGATED, + "mode_not_a_string": _single( + {"mode": 7, "supported_endpoints": ["/v1/embeddings"]}, "openai/gate-mode-not-a-string" + ), + } +) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + api_base: str + + +def _create(gateway: Gateway, name: str, litellm_params: Mapping[str, JsonValue], info: Mapping[str, JsonValue]) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"api_key": _PROVIDER_KEY, **litellm_params}, + "model_info": dict(info), + }, + ) + return string_value(object_value(created["model_info"])["id"]) + + +def _deploy(scenario: Scenario, api_base: str, label: str, deployment: _Deployment) -> tuple[str, tuple[str, ...]]: + name: Final = f"gate-{label}-{uuid.uuid4().hex[:12]}" + identities: Final = tuple( + _create(scenario.gateway, name, {"model": deployment.litellm_model, "api_base": api_base}, declaration) + for declaration in deployment.declarations + ) + tuple(scenario.cleanups.callback(scenario.delete_model, identity) for identity in identities) + return name, identities + + +def _listed_ids(gateway: Gateway) -> frozenset[str]: + listed: Final = gateway.request("GET", "/model/info", headers=_CLOSE) + assert listed.status_code == 200, listed.text + return frozenset( + string_value(object_value(object_value(entry)["model_info"])["id"]) + for entry in list_value(JSON_OBJECT.validate_json(listed.content)["data"]) + ) + + +def _served_by_every_worker(gateway: Gateway, identities: frozenset[str]) -> None: + held( + lambda: identities <= _listed_ids(gateway), + lambda served: served, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + api_base: Final = wire.url + "/v1" + deployed: Final = {label: _deploy(scenario, api_base, label, spec) for label, spec in _SHARED.items()} + _served_by_every_worker( + gateway, frozenset(identity for _name, identities in deployed.values() for identity in identities) + ) + yield _Rig( + gateway, + _Upstream(wire), + MappingProxyType({label: name for label, (name, _ids) in deployed.items()}), + api_base, + ) + + +def _marker() -> str: + return f"gate-{uuid.uuid4().hex}" + + +def _openai_error(endpoint: str, model: str, supported: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"{endpoint} is not supported by model={model}. Supported endpoints: {supported}", + "type": "invalid_request_error", + "param": "model", + "code": "400", + } + } + + +def _anthropic_error(model: str, supported: str) -> dict[str, JsonValue]: + return { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"/v1/messages is not supported by model={model}. Supported endpoints: {supported}", + }, + } + + +def _chat_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], **({"stream": True} if stream else {})} + + +def _failure_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + return rows[0] + + +def _assert_rejected(rig: _Rig, response: httpx.Response, marker: str, expected: dict[str, JsonValue]) -> None: + assert response.status_code == 400, response.text + assert JSON_OBJECT.validate_json(response.content) == expected, response.text + assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) + assert _failure_row(response.headers["x-litellm-call-id"])["status"] == "failure" + + +@dataclass(frozen=True, slots=True) +class _Route: + path: str + endpoint: str + anthropic_shape: bool + stream: bool + model_in_path: bool + shape: str + + +_ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( + { + "chat_v1": _Route("/v1/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_v1_stream": _Route("/v1/chat/completions", "/v1/chat/completions", False, True, False, "chat"), + "chat_bare": _Route("/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_engines": _Route("/engines/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat"), + "chat_azure_deployments": _Route( + "/openai/deployments/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat" + ), + "messages": _Route("/v1/messages", "/v1/messages", True, False, False, "messages"), + "messages_stream": _Route("/v1/messages", "/v1/messages", True, True, False, "messages"), + "responses_v1": _Route("/v1/responses", "/v1/responses", False, False, False, "responses"), + "responses_v1_stream": _Route("/v1/responses", "/v1/responses", False, True, False, "responses"), + "responses_bare": _Route("/responses", "/v1/responses", False, False, False, "responses"), + "cursor_responses_body": _Route("/cursor/chat/completions", "/v1/responses", False, False, False, "cursor"), + "cursor_chat_body": _Route("/cursor/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + } +) + + +def _route_body(route: _Route, model: str, marker: str) -> dict[str, JsonValue]: + shaped: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {"messages": [{"role": "user", "content": marker}]}, + "messages": {"max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + "responses": {"input": marker}, + "cursor": {"input": [{"role": "user", "content": marker}]}, + } + return { + **({} if route.model_in_path else {"model": model}), + **shaped[route.shape], + **({"stream": True} if route.stream else {}), + } + + +def _send(gateway: Gateway, route: _Route, model: str, marker: str) -> httpx.Response: + path: Final = route.path.format(model=model) + body: Final = _route_body(route, model, marker) + if not route.stream: + return gateway.request("POST", path, body) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as streamed: + streamed.read() + return streamed + + +@pytest.mark.parametrize("route", tuple(_ROUTES)) +def test_chat_shaped_route_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, route: str) -> None: + spec: Final = _ROUTES[route] + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _send(rig.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected(rig, response, marker, expected) + + +def _openai_sdk_error(rig: _Rig, client: str, model: str, marker: str) -> openai.APIStatusError: + base_url: Final = f"{rig.gateway.client.base_url}/v1" + + def sync_call() -> None: + with openai.OpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: + if client == "openai_sync_chat": + sdk.chat.completions.create(model=model, messages=[{"role": "user", "content": marker}]) + else: + sdk.responses.create(model=model, input=marker) + + async def async_call() -> None: + async with openai.AsyncOpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: + if client == "openai_async_chat_stream": + await sdk.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + else: + await sdk.responses.create(model=model, input=marker, stream=True) + + def call() -> None: + return sync_call() if client.startswith("openai_sync") else asyncio.run(async_call()) + + with pytest.raises(openai.BadRequestError) as raised: + call() + return raised.value + + +def _anthropic_sdk_error(rig: _Rig, client: str, model: str, marker: str) -> anthropic.APIStatusError: + base_url: Final = str(rig.gateway.client.base_url) + + def sync_call() -> None: + with anthropic.Anthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: + sdk.messages.create(model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]) + + async def async_call() -> None: + async with anthropic.AsyncAnthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: + await sdk.messages.create( + model=model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True + ) + + def call() -> None: + return sync_call() if client == "anthropic_sync" else asyncio.run(async_call()) + + with pytest.raises(anthropic.BadRequestError) as raised: + call() + return raised.value + + +_SDK_CLIENTS: Final[Mapping[str, str]] = MappingProxyType( + { + "openai_sync_chat": "/v1/chat/completions", + "openai_async_chat_stream": "/v1/chat/completions", + "openai_sync_responses": "/v1/responses", + "openai_async_responses_stream": "/v1/responses", + "anthropic_sync": "/v1/messages", + "anthropic_async_stream": "/v1/messages", + } +) + + +@pytest.mark.parametrize("client", tuple(_SDK_CLIENTS)) +def test_sdk_clients_see_the_rejection_as_a_bad_request(rig: _Rig, client: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + anthropic_client: Final = client.startswith("anthropic") + error: Final = ( + _anthropic_sdk_error(rig, client, model, marker) + if anthropic_client + else _openai_sdk_error(rig, client, model, marker) + ) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if anthropic_client + else _openai_error(_SDK_CLIENTS[client], model, "/v1/embeddings") + ) + assert error.status_code == 400, error + _assert_rejected(rig, error.response, marker, expected) + + +def _sdk_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + async def call() -> tuple[tuple[float, ...], str]: + async with openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", api_key=rig.gateway.key, max_retries=0 + ) as sdk: + raw: Final = await sdk.embeddings.with_raw_response.create( + model=model, input=marker, encoding_format="float" + ) + parsed: Final = raw.parse() + return tuple(parsed.data[0].embedding), raw.headers["x-litellm-call-id"] + + return asyncio.run(call()) + + +def _httpx_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + response: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"} + ) + assert response.status_code == 200, response.text + data: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + vector: Final = list_value(object_value(data[0])["embedding"]) + return tuple(float(str(value)) for value in vector), response.headers["x-litellm-call-id"] + + +@pytest.mark.parametrize("client", ("httpx_sync", "openai_async_sdk")) +def test_embeddings_still_reach_the_embedding_only_model(rig: _Rig, client: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + vector, call_id = ( + _httpx_embedding(rig, model, marker) if client == "httpx_sync" else _sdk_embedding(rig, model, marker) + ) + assert vector == _VECTOR + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + row: Final = eventually( + lambda: read_rows( + 'SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id = %s OR litellm_call_id = %s', + (call_id, call_id), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert [(entry["status"], entry["call_type"]) for entry in row] == [("success", "aembedding")], row + + +@pytest.mark.parametrize("declaration", tuple(_GATED)) +def test_non_chat_declarations_reject_chat_completions(rig: _Rig, declaration: str) -> None: + model: Final = rig.names[declaration] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, _GATED[declaration][1])) + + +def _assert_served(rig: _Rig, response: httpx.Response, marker: str, target: str) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == (target,) + + +@pytest.mark.parametrize("declaration", tuple(_UNGATED)) +def test_chat_capable_or_unclear_declarations_keep_reaching_the_upstream(rig: _Rig, declaration: str) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(rig.names[declaration], marker)) + _assert_served(rig, response, marker, "/v1/chat/completions") + + +def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_name(rig: _Rig) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(alias, marker), key=key) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", target, "/v1/embeddings")) + + +def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None: + public: Final = f"gate-team-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team() + identity: Final = _create( + rig.gateway, + public, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base}, + {**_EMBED_ONLY, "team_id": team}, + ) + scenario.cleanups.callback(scenario.delete_model, identity) + key: Final = scenario.key(team_id=team) + _served_by_every_worker(rig.gateway, frozenset({identity})) + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": public, "input": marker, "encoding_format": "float"}, key=key + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + rejected_marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(public, rejected_marker), key=key + ) + _assert_rejected( + rig, response, rejected_marker, _openai_error("/v1/chat/completions", public, "/v1/embeddings") + ) + + +def test_wildcard_embedding_only_deployment_rejects_chat_for_any_matched_model(rig: _Rig) -> None: + prefix: Final = f"gate-wild-{uuid.uuid4().hex[:12]}" + requested: Final = f"{prefix}/{_EMBEDDING_MODEL}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"{prefix}/*", {"model": "openai/*", "api_base": rig.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(requested, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", requested, "/v1/embeddings")) + + +def test_key_without_access_gets_the_access_error_and_learns_nothing_about_endpoints(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[rig.names["chat_target"]]) + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["embedding_only"], marker), key=key + ) + assert response.status_code == 403, response.text + assert object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"] == "key_model_access_denied", ( + response.text + ) + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_unknown_model_keeps_the_invalid_model_error(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(f"gate-missing-{uuid.uuid4().hex}", marker) + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text and "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_request_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": [rig.names["chat_target"]]} + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": []} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_request_fallbacks_do_not_rescue_an_unknown_primary(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(f"gate-missing-{uuid.uuid4().hex}", marker), "fallbacks": [rig.names["chat_target"]]}, + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_text_completions_reject_an_embedding_only_model_before_the_upstream(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", "/v1/completions", {"model": model, "prompt": marker}) + _assert_rejected(rig, response, marker, _openai_error("/v1/completions", model, "/v1/embeddings")) + + +@pytest.mark.parametrize("method", ("generateContent", "streamGenerateContent")) +def test_gemini_routes_reject_an_embedding_only_model_before_the_upstream(rig: _Rig, method: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", f"/v1beta/models/{model}:{method}", {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + ) + assert response.status_code == 400, response.text + assert f"{method} is not supported by model={model}. Supported endpoints: /v1/embeddings" in response.text + assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) + + +def test_non_string_mode_keeps_reaching_the_upstream(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["mode_not_a_string"], marker) + ) + assert response.status_code == 404, response.text + assert "scripted: not a chat model" in response.text, response.text + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_repeated_rejections_log_one_failure_row_per_call(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + responses: Final = tuple( + rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for _ in range(3) + ) + expected: Final = _openai_error("/v1/chat/completions", model, "/v1/embeddings") + assert [(item.status_code, item.json()) for item in responses] == [(400, expected)] * 3 + call_ids: Final = tuple(item.headers["x-litellm-call-id"] for item in responses) + assert len(set(call_ids)) == 3, call_ids + assert [_failure_row(call_id)["status"] for call_id in call_ids] == ["failure"] * 3 + assert rig.upstream.calls(marker) == () + + +def _closed_chat(rig: _Rig, model: str, marker: str) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker), headers=_CLOSE) + + +def test_live_declaration_update_turns_the_gate_on_and_off_on_every_worker(rig: _Rig) -> None: + chat_capable: Final[dict[str, JsonValue]] = { + "mode": "chat", + "supported_endpoints": ["/v1/chat/completions", "/v1/embeddings"], + } + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy(scenario, rig.api_base, "flip", _single(chat_capable)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + served: Final = _closed_chat(rig, name, marker) + _assert_served(rig, served, marker, "/v1/chat/completions") + patched: Final = rig.gateway.request( + "PATCH", f"/model/{identities[0]}/update", {"model_info": dict(_EMBED_ONLY)} + ) + assert patched.status_code == 200, patched.text + rejected: Final = held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 400, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + assert rejected.json() == _openai_error("/v1/chat/completions", name, "/v1/embeddings"), rejected.text + restored: Final = rig.gateway.request("PATCH", f"/model/{identities[0]}/update", {"model_info": chat_capable}) + assert restored.status_code == 200, restored.text + held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 200, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +_TESTS_ROOT: Final = Path(__file__).resolve().parents[2] +_CONFIG: Final = TypeAdapter(dict[str, JsonValue]) +_BURST_ROUTES: Final = ( + "chat_v1", + "chat_v1_stream", + "messages", + "messages_stream", + "responses_v1", + "responses_v1_stream", +) + + +def _scripted(url: str, identity: str, routes: Mapping[str, JsonResponse]) -> None: + registered: Final = httpx.post( + f"{url}/__scenarios", + json={ + "scenario_id": identity, + "response": RoutedResponse(content_type="application/x-routed", routes=dict(routes)).model_dump( + mode="json" + ), + }, + timeout=10, + trust_env=False, + ) + registered.raise_for_status() + + +class _Observed: + def __init__(self, url: str) -> None: + self.url = url + self.items: tuple[dict[str, JsonValue], ...] = () + + def read(self) -> tuple[dict[str, JsonValue], ...]: + payload: Final = JSON_OBJECT.validate_json( + httpx.get(f"{self.url}/__observations?include_method=true", timeout=10, trust_env=False).content + ) + self.items = (*self.items, *(object_value(item) for item in list_value(payload["requests"]))) + return self.items + + def carrying(self, marker: str) -> tuple[str, ...]: + return tuple(str(item["path"]) for item in self.read() if marker in json.dumps(item.get("body"))) + + +def _owned_config(directory: Path, url: str, embed: tuple[str, str], chat: tuple[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + litellm_settings: Final = object_value(shared["litellm_settings"]) + config: Final = { + **shared, + "model_list": [ + { + "model_name": embed[0], + "litellm_params": { + "model": f"openai/{_EMBEDDING_MODEL}", + "api_base": f"{url}/{embed[1]}/v1", + "api_key": _PROVIDER_KEY, + }, + "model_info": dict(_EMBED_ONLY), + }, + { + "model_name": chat[0], + "litellm_params": { + "model": f"openai/{_CHAT_MODEL}", + "api_base": f"{url}/{chat[1]}/v1", + "api_key": _PROVIDER_KEY, + }, + }, + ], + "litellm_settings": {**litellm_settings, "responses": {"background_mode": {"polling_via_cache": "all"}}}, + } + path: Final = directory / "gate-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +def _owned_embedding(owned: Gateway, model: str, marker: str) -> httpx.Response: + return owned.request("POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"}) + + +def _gated_expectation(route: str, model: str) -> dict[str, JsonValue]: + spec: Final = _ROUTES[route] + return ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + + +def _assert_gated_burst(results: tuple[tuple[str, httpx.Response], ...], model: str) -> tuple[str, ...]: + assert [(response.status_code, JSON_OBJECT.validate_json(response.content)) for _route, response in results] == [ + (400, _gated_expectation(route, model)) for route, _response in results + ], [response.text for _route, response in results] + return tuple(response.headers["x-litellm-call-id"] for _route, response in results) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_yaml_embedding_only_model_stays_gated_through_an_upstream_pause_and_a_worker_kill( + gateway: Gateway, tmp_path: Path +) -> None: + suffix: Final = uuid.uuid4().hex[:12] + embed: Final = (f"gate-yaml-embed-{suffix}", f"gate-embed-{suffix}") + chat: Final = (f"gate-yaml-chat-{suffix}", f"gate-chat-{suffix}") + embedding: Final = JsonResponse( + content_type="application/json", body=JSON_OBJECT.validate_json(_embedding_reply().body) + ) + chat_completion: Final = JsonResponse( + content_type="application/json", + body={ + "id": f"chatcmpl-{suffix}", + "object": "chat.completion", + "created": 1, + "model": _CHAT_MODEL, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _SERVED_TEXT}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, + }, + ) + with owned_upstream(tmp_path) as slot: + upstream: Final = slot.process + assert upstream is not None + url: Final = slot.url + _scripted(url, embed[1], {"POST /v1/embeddings": embedding}) + _scripted(url, chat[1], {"POST /v1/chat/completions": chat_completion}) + config: Final = _owned_config(tmp_path, url, embed, chat) + with owned_proxy_process( + gateway, tmp_path, {"INTEGRATION_UPSTREAM_URL": url}, config=config, workers=2 + ) as owned: + proxy: Final = owned.gateway + eventually( + lambda: proxy.request("POST", "/v1/chat/completions", _chat_body(chat[0], "probe")).status_code, + lambda status: status == 200, + seconds=30, + ) + workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) + observed: Final = _Observed(url) + background_marker: Final = _marker() + background: Final = proxy.request( + "POST", "/v1/responses", {"model": embed[0], "input": background_marker, "background": True} + ) + assert background.status_code == 400, background.text + assert background.json() == _openai_error("/v1/responses", embed[0], "/v1/embeddings"), background.text + burst: Final = tuple((_BURST_ROUTES[index % len(_BURST_ROUTES)], _marker()) for index in range(18)) + with paused(upstream), ThreadPoolExecutor(max_workers=9) as pool: + futures: Final = tuple( + (route, pool.submit(_send, proxy, _ROUTES[route], embed[0], marker)) for route, marker in burst + ) + paused_results: Final = tuple((route, future.result(timeout=15)) for route, future in futures) + gated_ids: Final = _assert_gated_burst(paused_results, embed[0]) + served_markers: Final = tuple(_marker() for _ in range(8)) + with ThreadPoolExecutor(max_workers=8) as pool: + served: Final = tuple(pool.map(partial(_owned_embedding, proxy, embed[0]), served_markers)) + assert [response.status_code for response in served] == [200] * 8, [item.text for item in served] + embedded_paths: Final = eventually( + lambda: tuple(observed.carrying(marker) for marker in served_markers), + lambda paths: all(len(entry) >= 1 for entry in paths), + seconds=20, + ) + assert embedded_paths == ((f"/{embed[1]}/v1/embeddings",),) * 8, embedded_paths + assert [observed.carrying(marker) for _route, marker in burst] == [()] * 18 + assert observed.carrying(background_marker) == () + assert [_failure_row(call_id)["status"] for call_id in gated_ids] == ["failure"] * 18 + with owned.replacing_worker(min(workers)) as kill: + kill() + after_kill: Final = tuple( + ("chat_v1", _send(proxy, _ROUTES["chat_v1"], embed[0], _marker())) for _ in range(4) + ) + _assert_gated_burst(after_kill, embed[0]) + assert _owned_embedding(proxy, embed[0], _marker()).status_code == 200 + respawned: Final = tuple( + ( + "chat_v1", + proxy.request("POST", "/v1/chat/completions", _chat_body(embed[0], _marker()), headers=_CLOSE), + ) + for _ in range(6) + ) + _assert_gated_burst(respawned, embed[0]) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5dac1d25572..4913d855d83 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1005,6 +1005,89 @@ class TestResponsesWSFirstFrameModelAuth: ) +class TestResponsesWSDeclaredNonChatModel: + @pytest.mark.asyncio + @pytest.mark.parametrize("query_model", [None, "embed-only"]) + @pytest.mark.parametrize( + "model_info, rejected", + [ + ({"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, True), + ({"mode": "embedding"}, False), + ], + ) + async def test_endpoint_rejects_a_model_declared_non_chat_before_routing( + self, query_model: str | None, model_info: dict[str, object], rejected: bool + ): + from litellm.proxy.response_api_endpoints.endpoints import ( + responses_websocket_endpoint, + ) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "embed-only", + "litellm_params": {"model": "sagemaker/openai/embed-endpoint", "aws_region_name": "us-east-1"}, + "model_info": model_info, + } + ] + ) + ws = MagicMock() + ws.headers = {} + ws.query_params = {} + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.accept = AsyncMock() + ws.receive_text = AsyncMock( + return_value=json.dumps({"type": "response.create", "model": "embed-only", "input": []}) + ) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + + processor = MagicMock() + processor.common_processing_pre_call_logic = AsyncMock( + return_value=({"model": "embed-only", "litellm_metadata": {}}, MagicMock()) + ) + + async def fake_llm_call(): + return None + + with ( + patch( # test-quality-ok: first-frame model auth needs a live key table and has its own tests above + "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the pre-call processor needs a live proxy; whether the session reaches it is under test + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + return_value=processor, + ), + patch( # test-quality-ok: routing is the seam that would reach the provider + "litellm.proxy.route_llm_request.route_request", + new_callable=AsyncMock, + return_value=fake_llm_call(), + ) as mock_route_request, + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the session reads the proxy's router global + ): + await responses_websocket_endpoint( + websocket=ws, + model=query_model, + user_api_key_dict=MagicMock(team_id=None), + ) + + if not rejected: + mock_route_request.assert_awaited_once() + ws.close.assert_not_awaited() + return + frame = json.loads(ws.send_text.await_args.args[0]) + assert frame["type"] == "error" + assert frame["error"]["type"] == "invalid_request_error" + assert frame["error"]["message"] == ( + "/v1/responses is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + ) + ws.close.assert_awaited_once_with(code=1008, reason="Pre-call error") + processor.common_processing_pre_call_logic.assert_not_awaited() + mock_route_request.assert_not_awaited() + + class TestReadWSModelFromFirstFrameErrors: @pytest.mark.asyncio async def test_transport_error_first_frame_closes_with_internal_error(self): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 7526b6617c8..3a849dc0a9b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -6933,6 +6933,57 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == "gpt-4" + @pytest.mark.asyncio + async def test_fallback_skips_a_model_that_rejects_the_endpoint(self): + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "chat-rl"}) + attempted: list[str] = [] + + async def mock_pre_call_logic(**kwargs): + model = processor.data.get("model") + attempted.append(model) + if model == "chat-rl": + raise ProxyRateLimitError(detail="RPM limit exceeded", headers={"retry-after": "30"}) + if model == "embed-only": + raise ProxyUnsupportedEndpointError( + message="/v1/chat/completions is not supported by model=embed-only", + type="invalid_request_error", + param="model", + code=400, + ) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{"chat-rl": ["embed-only", "chat-ctl"]}] + + with patch.object( + processor, + "common_processing_pre_call_logic", + side_effect=mock_pre_call_logic, + ): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="chat-rl", + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted == ["chat-rl", "embed-only", "chat-ctl"] + assert processor.data["model"] == "chat-ctl" + @pytest.mark.asyncio async def test_fallback_uses_key_level_router_settings(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index a298f57b3e2..87dc9aee7e2 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3,10 +3,11 @@ import pytest from types import SimpleNamespace from typing import Final -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from fastapi import HTTPException +import litellm from litellm import Router from litellm.proxy.route_llm_request import ( ProxyModelNotFoundError, @@ -1770,8 +1771,12 @@ async def test_route_request_router_settings_override_skips_null_fields(): CHAT_SHAPED_ROUTES: Final = ( ("acompletion", "/v1/chat/completions"), + ("atext_completion", "/v1/completions"), ("anthropic_messages", "/v1/messages"), ("aresponses", "/v1/responses"), + ("agenerate_content", "generateContent"), + ("agenerate_content_stream", "streamGenerateContent"), + ("acreate_interaction", "/v1beta/interactions"), ) CHAT_SHAPED_ROUTE_TYPES: Final = tuple(route_type for route_type, _ in CHAT_SHAPED_ROUTES) @@ -1791,8 +1796,12 @@ def _router_with(*deployments: dict[str, object], **router_kwargs: object) -> Ro return Router(model_list=list(deployments), **router_kwargs) -def _unsupported(route_type: str, model: object, router: Router | None) -> ProxyUnsupportedEndpointError | None: - return endpoint_unsupported_by_model(route_type=route_type, model=model, llm_router=router, team_id=None) +def _unsupported( + route_type: str, model: object, router: Router | None, **request: object +) -> ProxyUnsupportedEndpointError | None: + return endpoint_unsupported_by_model( + route_type=route_type, data={"model": model, **request}, llm_router=router, team_id=None + ) @pytest.mark.parametrize("route_type, endpoint", CHAT_SHAPED_ROUTES) @@ -1808,7 +1817,7 @@ def test_endpoint_unsupported_rejects_chat_routes_on_embedding_only_model(route_ assert error.message == f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" -@pytest.mark.parametrize("route_type", ["aembedding", "atext_completion", "aimage_generation"]) +@pytest.mark.parametrize("route_type", ["aembedding", "aimage_generation", "_aresponses_websocket"]) def test_endpoint_unsupported_leaves_non_chat_routes_alone(route_type): router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) @@ -1824,6 +1833,8 @@ def test_endpoint_unsupported_leaves_non_chat_routes_alone(route_type): {"mode": "embedding", "supported_endpoints": None}, {"mode": "embedding", "supported_endpoints": "/v1/embeddings"}, {"mode": 1, "supported_endpoints": ["/v1/embeddings"]}, + {"mode": "", "supported_endpoints": ["/v1/embeddings"]}, + {"mode": " ", "supported_endpoints": ["/v1/embeddings"]}, ], ) def test_endpoint_unsupported_keeps_todays_behavior_without_a_valid_declaration(route_type, model_info): @@ -1840,6 +1851,7 @@ def test_endpoint_unsupported_keeps_todays_behavior_without_a_valid_declaration( {"mode": "chat", "supported_endpoints": ["/v1/batch"]}, {"mode": "completion", "supported_endpoints": ["/fine_tuning"]}, {"mode": "responses", "supported_endpoints": ["/v1/realtime"]}, + {"mode": " Chat ", "supported_endpoints": ["/v1/embeddings"]}, ], ) def test_endpoint_unsupported_needs_a_non_chat_mode_to_reject(route_type, model_info): @@ -1851,7 +1863,16 @@ def test_endpoint_unsupported_needs_a_non_chat_mode_to_reject(route_type, model_ @pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) @pytest.mark.parametrize( "declared", - [["/v1/messages"], ["/v1/responses"], ["/v1/chat/completions"], ["/chat/completions"], ["/responses"]], + [ + ["/v1/messages"], + ["/v1/responses"], + ["/v1/chat/completions"], + ["/chat/completions"], + ["/responses"], + ["/v1/completions"], + ["/completions"], + ["/v1beta/interactions"], + ], ) def test_endpoint_unsupported_allows_every_chat_route_when_any_chat_endpoint_is_declared(route_type, declared): router: Final = _router_with( @@ -1918,6 +1939,72 @@ def test_endpoint_unsupported_skips_the_check_without_a_router(): assert _unsupported("acompletion", "embed-only", None) is None +CHAT_FALLBACK_DEPLOYMENT: Final = { + "model_name": "chat-fallback", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, +} + + +@pytest.mark.parametrize( + "router_kwargs, request_kwargs", + [ + ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {}), + ({"default_fallbacks": ["chat-fallback"]}, {}), + ({}, {"fallbacks": ["chat-fallback"]}), + ({}, {"fallbacks": [{"embed-only": ["chat-fallback"]}]}), + ({}, {"router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), + ], +) +def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): + router: Final = _router_with( + _embedding_deployment("embed-only", EMBEDDING_ONLY), CHAT_FALLBACK_DEPLOYMENT, **router_kwargs + ) + + assert _unsupported("acompletion", "embed-only", router, **request_kwargs) is None + + +@pytest.mark.parametrize( + "router_kwargs, request_kwargs", + [ + ({"fallbacks": [{"other-group": ["chat-fallback"]}]}, {}), + ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"disable_fallbacks": True}), + ({"default_fallbacks": ["chat-fallback"]}, {"fallbacks": []}), + ( + {"fallbacks": [{"embed-only": ["chat-fallback"]}]}, + {"router_settings_override": {"fallbacks": [{"other-group": ["chat-fallback"]}]}}, + ), + ({}, {"fallbacks": [], "router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), + ], +) +def test_endpoint_unsupported_rejects_when_no_fallback_would_answer(router_kwargs, request_kwargs): + router: Final = _router_with( + _embedding_deployment("embed-only", EMBEDDING_ONLY), CHAT_FALLBACK_DEPLOYMENT, **router_kwargs + ) + + assert isinstance(_unsupported("acompletion", "embed-only", router, **request_kwargs), ProxyUnsupportedEndpointError) + + +ECHOED_CATALOG_ENTRY: Final = {"mode": "image_generation", "supported_endpoints": ["/v1/images/generations"]} + + +@pytest.mark.parametrize( + "model_info, rejected", + [ + ({"key": "echoed-image-model", **ECHOED_CATALOG_ENTRY}, False), + ({"key": "echoed-image-model", "mode": "image_generation", "supported_endpoints": ["/v1/embeddings"]}, False), + ({"key": "echoed-image-model", "mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, True), + (dict(ECHOED_CATALOG_ENTRY), True), + ], +) +def test_endpoint_unsupported_ignores_fields_copied_from_the_cost_map(model_info, rejected): + router: Final = _router_with(_embedding_deployment("image-row", model_info)) + + with patch.dict(litellm.model_cost, {"echoed-image-model": ECHOED_CATALOG_ENTRY}): + error: Final = _unsupported("acompletion", "image-row", router) + + assert isinstance(error, ProxyUnsupportedEndpointError) is rejected + + def _router_with_defaults(default_litellm_params: dict[str, object]): import litellm From 6728e9562dc2e4f1658a4eb0bf7b61c780cf816b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:04:21 -0700 Subject: [PATCH 03/21] fix(proxy): only step aside for fallbacks that can answer chat, gate the Responses WebSocket after alias resolution, and compare declarations to the loaded catalog --- .../proxy/response_api_endpoints/endpoints.py | 10 +- litellm/proxy/route_llm_request.py | 109 ++++++++++++------ .../test_non_chat_model_endpoint_gate.py | 10 ++ .../response_api_endpoints/test_endpoints.py | 16 ++- .../proxy/test_common_request_processing.py | 18 ++- tests/unit/proxy/test_route_llm_request.py | 87 +++++++++++--- 6 files changed, 178 insertions(+), 72 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c6d5221c275..7f9fcca369b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,7 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) -from litellm.proxy.route_llm_request import endpoint_unsupported_by_model, raise_if_required_body_param_missing +from litellm.proxy.route_llm_request import raise_if_required_body_param_missing from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1546,14 +1546,6 @@ async def _responses_websocket_session( user_api_key_dict=user_api_key_dict, llm_router=llm_router, ) - unsupported_endpoint: Final = endpoint_unsupported_by_model( - route_type="aresponses", - data=data, - llm_router=llm_router, - team_id=user_api_key_dict.team_id, - ) - if unsupported_endpoint is not None: - raise unsupported_endpoint ( data, litellm_logging_obj, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 087a7f56d7a..2f66327a8b9 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -374,6 +374,7 @@ CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( "agenerate_content": "generateContent", "agenerate_content_stream": "streamGenerateContent", "acreate_interaction": "/v1beta/interactions", + "_aresponses_websocket": "/v1/responses", } ) @@ -390,8 +391,6 @@ CHAT_SHAPED_ENDPOINTS: Final[frozenset[str]] = frozenset( } ) -DECLARATION_FIELDS: Final[frozenset[str]] = frozenset({"mode", "supported_endpoints"}) - CHAT_SHAPED_MODES: Final[frozenset[str]] = frozenset({"chat", "completion", "responses"}) @@ -411,32 +410,18 @@ _FALLBACK_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter( # mutab ) -class _LiveCostMap(Protocol): - @property - def model_cost(self) -> Mapping[str, object]: ... - - class _RouterFallbacks(Protocol): @property def fallbacks(self) -> Sequence[object] | None: ... -def _live_cost_map_entry(catalog: _LiveCostMap, catalog_key: str) -> Mapping[str, object]: - try: - return JSON_OBJECT_ADAPTER.validate_python(catalog.model_cost.get(catalog_key) or {}) - except ValidationError: - return MappingProxyType({}) - - def _echoed_from_cost_map(model_info: Mapping[str, object]) -> frozenset[str]: catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) if not isinstance(catalog_key, str): return frozenset() return frozenset( echoed_cost_map_fields( - model_info, - _live_cost_map_entry(litellm, catalog_key), - GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}), + model_info, GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}) ) ) @@ -447,13 +432,21 @@ def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[ declared: Final = _DeploymentEndpointDeclaration.model_validate(model_info) except ValidationError: return None - if not DECLARATION_FIELDS.isdisjoint(_echoed_from_cost_map(model_info)): - return None mode: Final = (declared.mode or "").strip().lower() if not mode or mode in CHAT_SHAPED_MODES or declared.supported_endpoints is None: return None endpoints: Final = frozenset(declared.supported_endpoints) - return endpoints if endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS) else None + if not endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS): + return None + return None if "supported_endpoints" in _echoed_from_cost_map(model_info) else endpoints + + +def _group_deployments(model: str, llm_router: LitellmRouter, team_id: str | None) -> tuple[Mapping[str, object], ...]: + return tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) + + +def _any_serves_chat(deployments: Sequence[Mapping[str, object]]) -> bool: + return any(_declared_non_chat_endpoints(deployment) is None for deployment in deployments) def _override_fallbacks(data: Mapping[str, object]) -> object: @@ -474,22 +467,65 @@ def _router_fallbacks(llm_router: _RouterFallbacks) -> Sequence[object] | None: return llm_router.fallbacks -def _fallback_would_answer(model: str, data: Mapping[str, object], llm_router: LitellmRouter) -> bool: +def _fallback_targets( + model: str, + data: Mapping[str, object], + fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract +) -> Sequence[object]: + if check_non_standard_fallback_format(fallbacks=fallbacks): + return fallbacks + resolved, _ = get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, + lookup_groups=fallback_lookup_groups(data, model), + ) + return resolved or () + + +def _fallback_target_group(target: object) -> str | None: + if isinstance(target, str): + return target + try: + entry: Final = JSON_OBJECT_ADAPTER.validate_python(target) + except ValidationError: + return None + named_model: Final = entry.get("model") + return named_model if isinstance(named_model, str) else None + + +def _fallback_target_answers( + group: str | None, + data: Mapping[str, object], + llm_router: LitellmRouter, + team_id: str | None, + walked: frozenset[str], +) -> bool: + if group is None or group in walked: + return False + deployments: Final = _group_deployments(group, llm_router, team_id) + if not deployments: + return False + return _any_serves_chat(deployments) or _fallback_would_answer(group, data, llm_router, team_id, walked) + + +def _fallback_would_answer( + model: str, + data: Mapping[str, object], + llm_router: LitellmRouter, + team_id: str | None, + walked: frozenset[str] = frozenset(), +) -> bool: if fallbacks_disabled_for_request(data): return False try: fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(_fallbacks_for_request(data, llm_router)) except ValidationError: return False - if not fallbacks: - return False - if check_non_standard_fallback_format(fallbacks=fallbacks): - return True - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=fallback_lookup_groups(data, model), + walked_here: Final = walked | {model} + return any( + has_unattempted_fallback_target((target,), data) + and _fallback_target_answers(_fallback_target_group(target), data, llm_router, team_id, walked_here) + for target in _fallback_targets(model, data, fallbacks) ) - return has_unattempted_fallback_target(resolved, data) def endpoint_unsupported_by_model( @@ -502,15 +538,16 @@ def endpoint_unsupported_by_model( model: Final = data.get("model") if endpoint is None or not isinstance(model, str) or not model or llm_router is None: return None - deployments: Final = tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) - non_chat: Final = tuple( - endpoints for deployment in deployments if (endpoints := _declared_non_chat_endpoints(deployment)) is not None - ) - if not deployments or len(non_chat) < len(deployments) or _fallback_would_answer(model, data, llm_router): + deployments: Final = _group_deployments(model, llm_router, team_id) + if not deployments or _any_serves_chat(deployments): return None - supported: Final = ", ".join(sorted(frozenset[str]().union(*non_chat))) or "none" + if _fallback_would_answer(model, data, llm_router, team_id): + return None + supported: Final = ", ".join( + sorted(frozenset[str]().union(*(_declared_non_chat_endpoints(d) or frozenset() for d in deployments))) + ) return ProxyUnsupportedEndpointError( - message=f"{endpoint} is not supported by model={model}. Supported endpoints: {supported}", + message=f"{endpoint} is not supported by model={model}. Supported endpoints: {supported or 'none'}", type="invalid_request_error", param="model", code=status.HTTP_400_BAD_REQUEST, diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index b12b42ec638..cd78722af9e 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -613,6 +613,16 @@ def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) +def test_fallbacks_that_cannot_answer_chat_keep_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + fallbacks: Final = [rig.names["image_generation"], model, f"gate-missing-{uuid.uuid4().hex}"] + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": fallbacks} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + def test_request_fallbacks_do_not_rescue_an_unknown_primary(rig: _Rig) -> None: marker: Final = _marker() response: Final = rig.gateway.request( diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 4913d855d83..4f71fedd174 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1015,12 +1015,13 @@ class TestResponsesWSDeclaredNonChatModel: ({"mode": "embedding"}, False), ], ) - async def test_endpoint_rejects_a_model_declared_non_chat_before_routing( + async def test_session_closes_on_a_model_declared_non_chat_before_routing( self, query_model: str | None, model_info: dict[str, object], rejected: bool ): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) + from litellm.proxy.route_llm_request import endpoint_unsupported_by_model router: Final = litellm.Router( model_list=[ @@ -1043,10 +1044,16 @@ class TestResponsesWSDeclaredNonChatModel: ws.send_text = AsyncMock() ws.close = AsyncMock() + async def gated_pre_call(**kwargs): + unsupported = endpoint_unsupported_by_model( + route_type=kwargs["route_type"], data={"model": kwargs["model"]}, llm_router=router, team_id=None + ) + if unsupported is not None: + raise unsupported + return {"model": kwargs["model"], "litellm_metadata": {}}, MagicMock() + processor = MagicMock() - processor.common_processing_pre_call_logic = AsyncMock( - return_value=({"model": "embed-only", "litellm_metadata": {}}, MagicMock()) - ) + processor.common_processing_pre_call_logic = AsyncMock(side_effect=gated_pre_call) async def fake_llm_call(): return None @@ -1084,7 +1091,6 @@ class TestResponsesWSDeclaredNonChatModel: "/v1/responses is not supported by model=embed-only. Supported endpoints: /v1/embeddings" ) ws.close.assert_awaited_once_with(code=1008, reason="Pre-call error") - processor.common_processing_pre_call_logic.assert_not_awaited() mock_route_request.assert_not_awaited() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 3a849dc0a9b..bfee48f8b7b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -725,12 +725,20 @@ class TestProxyBaseLLMRequestProcessing: ) @pytest.mark.asyncio + @pytest.mark.parametrize( + "route_type, endpoint", + [("acompletion", "/v1/chat/completions"), ("_aresponses_websocket", "/v1/responses")], + ) + @pytest.mark.parametrize( + "requested_model, key_aliases", + [("embed-only", {}), ("embed-alias", {"embed-alias": "embed-only"})], + ) async def test_common_processing_pre_call_logic_rejects_chat_on_embedding_only_model_before_pre_call_hooks( - self, monkeypatch + self, monkeypatch, route_type, endpoint, requested_model, key_aliases ): processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( monkeypatch, - request_body={"model": "embed-only", "messages": [{"role": "user", "content": "hello"}]}, + request_body={"model": requested_model, "messages": [{"role": "user", "content": "hello"}]}, guardrail_tags=[], ) @@ -738,16 +746,16 @@ class TestProxyBaseLLMRequestProcessing: await processing_obj.common_processing_pre_call_logic( request=mock_request, general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases=key_aliases), proxy_logging_obj=mock_proxy_logging_obj, proxy_config=mock_proxy_config, - route_type="acompletion", + route_type=route_type, llm_router=self._router_with_embedding_only_model(), ) assert exc_info.value.code == "400" assert exc_info.value.message == ( - "/v1/chat/completions is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" ) mock_proxy_logging_obj.pre_call_hook.assert_not_awaited() diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 87dc9aee7e2..475aeaa896d 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -3,13 +3,16 @@ import pytest from types import SimpleNamespace from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock from fastapi import HTTPException import litellm from litellm import Router +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy.route_llm_request import ( + CHAT_SHAPED_ENDPOINTS, + CHAT_SHAPED_MODES, ProxyModelNotFoundError, ProxyUnsupportedEndpointError, endpoint_unsupported_by_model, @@ -1777,6 +1780,7 @@ CHAT_SHAPED_ROUTES: Final = ( ("agenerate_content", "generateContent"), ("agenerate_content_stream", "streamGenerateContent"), ("acreate_interaction", "/v1beta/interactions"), + ("_aresponses_websocket", "/v1/responses"), ) CHAT_SHAPED_ROUTE_TYPES: Final = tuple(route_type for route_type, _ in CHAT_SHAPED_ROUTES) @@ -1817,7 +1821,7 @@ def test_endpoint_unsupported_rejects_chat_routes_on_embedding_only_model(route_ assert error.message == f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" -@pytest.mark.parametrize("route_type", ["aembedding", "aimage_generation", "_aresponses_websocket"]) +@pytest.mark.parametrize("route_type", ["aembedding", "aimage_generation", "aspeech", "_arealtime"]) def test_endpoint_unsupported_leaves_non_chat_routes_alone(route_type): router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) @@ -1944,6 +1948,8 @@ CHAT_FALLBACK_DEPLOYMENT: Final = { "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, } +EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment("embed-fallback", EMBEDDING_ONLY) + @pytest.mark.parametrize( "router_kwargs, request_kwargs", @@ -1953,11 +1959,18 @@ CHAT_FALLBACK_DEPLOYMENT: Final = { ({}, {"fallbacks": ["chat-fallback"]}), ({}, {"fallbacks": [{"embed-only": ["chat-fallback"]}]}), ({}, {"router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), + ({"fallbacks": [{"embed-only": ["embed-fallback", "missing-group", "chat-fallback"]}]}, {}), + ({}, {"fallbacks": [{"model": "chat-fallback", "temperature": 0.2}]}), + ({"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "chat-fallback"}}, {}), + ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}), ], ) def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), CHAT_FALLBACK_DEPLOYMENT, **router_kwargs + _embedding_deployment("embed-only", EMBEDDING_ONLY), + CHAT_FALLBACK_DEPLOYMENT, + EMBEDDING_FALLBACK_DEPLOYMENT, + **router_kwargs, ) assert _unsupported("acompletion", "embed-only", router, **request_kwargs) is None @@ -1974,35 +1987,75 @@ def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, r {"router_settings_override": {"fallbacks": [{"other-group": ["chat-fallback"]}]}}, ), ({}, {"fallbacks": [], "router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), + ({"fallbacks": [{"embed-only": ["embed-fallback"]}]}, {}), + ({"fallbacks": [{"embed-only": ["embed-only"]}]}, {}), + ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), + ({"default_fallbacks": ["embed-fallback"]}, {}), + ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), + ({}, {"fallbacks": [{"model": "embed-fallback", "temperature": 0.2}]}), + ({"fallbacks": [{"embed-only": ["embed-alias"]}], "model_group_alias": {"embed-alias": "embed-fallback"}}, {}), + ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["embed-only"]}]}, {}), + ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {"fallbacks": []}), ], ) def test_endpoint_unsupported_rejects_when_no_fallback_would_answer(router_kwargs, request_kwargs): router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), CHAT_FALLBACK_DEPLOYMENT, **router_kwargs + _embedding_deployment("embed-only", EMBEDDING_ONLY), + CHAT_FALLBACK_DEPLOYMENT, + EMBEDDING_FALLBACK_DEPLOYMENT, + **router_kwargs, ) - assert isinstance(_unsupported("acompletion", "embed-only", router, **request_kwargs), ProxyUnsupportedEndpointError) + assert isinstance( + _unsupported("acompletion", "embed-only", router, **request_kwargs), ProxyUnsupportedEndpointError + ) -ECHOED_CATALOG_ENTRY: Final = {"mode": "image_generation", "supported_endpoints": ["/v1/images/generations"]} +def _catalog_key_declaring_non_chat() -> str: + return min( + key + for key, entry in GetModelCostMap.loaded_model_cost_map().items() + if entry.get("litellm_provider") == "openai" + and entry.get("mode") not in CHAT_SHAPED_MODES + and isinstance(endpoints := entry.get("supported_endpoints"), list) + and CHAT_SHAPED_ENDPOINTS.isdisjoint(endpoints) + ) + + +def _catalog_declaration(catalog_key: str) -> dict[str, object]: + entry: Final = GetModelCostMap.loaded_model_cost_map()[catalog_key] + return {"mode": entry["mode"], "supported_endpoints": entry["supported_endpoints"]} @pytest.mark.parametrize( - "model_info, rejected", + "declaration, with_key, rejected", [ - ({"key": "echoed-image-model", **ECHOED_CATALOG_ENTRY}, False), - ({"key": "echoed-image-model", "mode": "image_generation", "supported_endpoints": ["/v1/embeddings"]}, False), - ({"key": "echoed-image-model", "mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, True), - (dict(ECHOED_CATALOG_ENTRY), True), + pytest.param("catalog", True, False, id="echo-of-the-catalog-row"), + pytest.param("catalog-mode-typed-endpoints", True, True, id="endpoints-typed-over-an-echo"), + pytest.param("typed-mode-catalog-endpoints", True, False, id="endpoints-left-as-echoed"), + pytest.param("typed", True, True, id="both-typed-over-an-echo"), + pytest.param("catalog", False, True, id="catalog-values-typed-by-hand"), ], ) -def test_endpoint_unsupported_ignores_fields_copied_from_the_cost_map(model_info, rejected): - router: Final = _router_with(_embedding_deployment("image-row", model_info)) +def test_endpoint_unsupported_ignores_endpoints_copied_from_the_loaded_catalog(declaration, with_key, rejected): + catalog_key: Final = _catalog_key_declaring_non_chat() + catalog: Final = _catalog_declaration(catalog_key) + declared: Final = { + "catalog": catalog, + "catalog-mode-typed-endpoints": {**catalog, "supported_endpoints": ["/v1/embeddings"]}, + "typed-mode-catalog-endpoints": {**catalog, "mode": "embedding"}, + "typed": EMBEDDING_ONLY, + }[declaration] + router: Final = _router_with( + { + "model_name": "catalog-row", + "litellm_params": {"model": catalog_key, "api_key": "test-key"}, + "model_info": {**({"key": catalog_key} if with_key else {}), **declared}, + } + ) - with patch.dict(litellm.model_cost, {"echoed-image-model": ECHOED_CATALOG_ENTRY}): - error: Final = _unsupported("acompletion", "image-row", router) - - assert isinstance(error, ProxyUnsupportedEndpointError) is rejected + assert litellm.model_cost[catalog_key]["supported_endpoints"] == declared["supported_endpoints"] + assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected def _router_with_defaults(default_litellm_params: dict[str, object]): From 19872dd0ecb1a72de8abce2f5344bda32b6e0a3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:31:35 -0700 Subject: [PATCH 04/21] fix(proxy): gate the Responses WebSocket with the router after aliases resolve, and ignore client-set internal fallback flags --- .../proxy/response_api_endpoints/endpoints.py | 10 +++++- litellm/proxy/route_llm_request.py | 10 +++--- .../test_non_chat_model_endpoint_gate.py | 36 +++++++++++++++++++ .../response_api_endpoints/test_endpoints.py | 22 +++++------- .../proxy/test_common_request_processing.py | 10 ++---- tests/unit/proxy/test_route_llm_request.py | 2 ++ 6 files changed, 65 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7f9fcca369b..536abfbfac7 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -36,7 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) -from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.proxy.route_llm_request import endpoint_unsupported_by_model, raise_if_required_body_param_missing from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1564,6 +1564,14 @@ async def _responses_websocket_session( model=resolved_model, route_type="_aresponses_websocket", ) + unsupported_endpoint: Final = endpoint_unsupported_by_model( + route_type="_aresponses_websocket", + data=data, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if unsupported_endpoint is not None: + raise unsupported_endpoint except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") try: diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 2f66327a8b9..b171caf3475 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -15,7 +15,6 @@ from litellm.router_utils.common_utils import is_proxy_admin_request from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, fallback_lookup_groups, - fallbacks_disabled_for_request, get_fallback_model_group_for_lookup_groups, has_unattempted_fallback_target, ) @@ -514,12 +513,15 @@ def _fallback_would_answer( team_id: str | None, walked: frozenset[str] = frozenset(), ) -> bool: - if fallbacks_disabled_for_request(data): + if data.get("disable_fallbacks") is True: + return False + raw_fallbacks: Final = _fallbacks_for_request(data, llm_router) + if raw_fallbacks is None: return False try: - fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(_fallbacks_for_request(data, llm_router)) + fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks) except ValidationError: - return False + return True walked_here: Final = walked | {model} return any( has_unattempted_fallback_target((target,), data) diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index cd78722af9e..43cf611a0c5 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import json import uuid from collections.abc import Iterator, Mapping @@ -13,6 +14,7 @@ import anthropic import httpx import openai import pytest +import websockets import yaml from integration._support.client import ( JSON_OBJECT, @@ -30,6 +32,7 @@ from integration._support.openai_wire import answering_model_discovery, chat_rep from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter +from websockets.exceptions import ConnectionClosed from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse @@ -526,6 +529,39 @@ def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_nam _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", target, "/v1/embeddings")) +async def _websocket_first_frame_and_close( + gateway: Gateway, key: str, model: str, marker: str +) -> tuple[dict[str, JsonValue], int | None, str | None]: + url: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://", 1) + "/v1/responses" + async with websockets.connect(url, additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10) as ws: + await ws.send(json.dumps({"type": "response.create", "model": model, "input": marker})) + first: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(ws.recv(), 30)) + with contextlib.suppress(ConnectionClosed): + await asyncio.wait_for(ws.recv(), 30) + return first, ws.close_code, ws.close_reason + + +@pytest.mark.parametrize("aliased", [False, True]) +def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstream(rig: _Rig, aliased: bool) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-ws-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + frame, close_code, close_reason = asyncio.run( + _websocket_first_frame_and_close(rig.gateway, key, alias if aliased else target, marker) + ) + assert frame == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"/v1/responses is not supported by model={target}. Supported endpoints: /v1/embeddings", + }, + } + assert (close_code, close_reason) == (1008, "Pre-call error") + assert rig.upstream.calls(marker) == () + + def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None: public: Final = f"gate-team-{uuid.uuid4().hex[:12]}" marker: Final = _marker() diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 4f71fedd174..2df84eae062 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1007,7 +1007,8 @@ class TestResponsesWSFirstFrameModelAuth: class TestResponsesWSDeclaredNonChatModel: @pytest.mark.asyncio - @pytest.mark.parametrize("query_model", [None, "embed-only"]) + @pytest.mark.parametrize("requested_model", ["embed-only", "embed-alias"]) + @pytest.mark.parametrize("use_query_model", [False, True]) @pytest.mark.parametrize( "model_info, rejected", [ @@ -1016,12 +1017,11 @@ class TestResponsesWSDeclaredNonChatModel: ], ) async def test_session_closes_on_a_model_declared_non_chat_before_routing( - self, query_model: str | None, model_info: dict[str, object], rejected: bool + self, requested_model: str, use_query_model: bool, model_info: dict[str, object], rejected: bool ): from litellm.proxy.response_api_endpoints.endpoints import ( responses_websocket_endpoint, ) - from litellm.proxy.route_llm_request import endpoint_unsupported_by_model router: Final = litellm.Router( model_list=[ @@ -1039,21 +1039,17 @@ class TestResponsesWSDeclaredNonChatModel: ws.url = "ws://testserver/v1/responses" ws.accept = AsyncMock() ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "response.create", "model": "embed-only", "input": []}) + return_value=json.dumps({"type": "response.create", "model": requested_model, "input": []}) ) ws.send_text = AsyncMock() ws.close = AsyncMock() + key_aliases: Final = {"embed-alias": "embed-only"} - async def gated_pre_call(**kwargs): - unsupported = endpoint_unsupported_by_model( - route_type=kwargs["route_type"], data={"model": kwargs["model"]}, llm_router=router, team_id=None - ) - if unsupported is not None: - raise unsupported - return {"model": kwargs["model"], "litellm_metadata": {}}, MagicMock() + async def alias_resolving_pre_call(**kwargs): + return {"model": key_aliases.get(kwargs["model"], kwargs["model"]), "litellm_metadata": {}}, MagicMock() processor = MagicMock() - processor.common_processing_pre_call_logic = AsyncMock(side_effect=gated_pre_call) + processor.common_processing_pre_call_logic = AsyncMock(side_effect=alias_resolving_pre_call) async def fake_llm_call(): return None @@ -1076,7 +1072,7 @@ class TestResponsesWSDeclaredNonChatModel: ): await responses_websocket_endpoint( websocket=ws, - model=query_model, + model=requested_model if use_query_model else None, user_api_key_dict=MagicMock(team_id=None), ) diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index bfee48f8b7b..b38fa820437 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -725,16 +725,12 @@ class TestProxyBaseLLMRequestProcessing: ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "route_type, endpoint", - [("acompletion", "/v1/chat/completions"), ("_aresponses_websocket", "/v1/responses")], - ) @pytest.mark.parametrize( "requested_model, key_aliases", [("embed-only", {}), ("embed-alias", {"embed-alias": "embed-only"})], ) async def test_common_processing_pre_call_logic_rejects_chat_on_embedding_only_model_before_pre_call_hooks( - self, monkeypatch, route_type, endpoint, requested_model, key_aliases + self, monkeypatch, requested_model, key_aliases ): processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( monkeypatch, @@ -749,13 +745,13 @@ class TestProxyBaseLLMRequestProcessing: user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases=key_aliases), proxy_logging_obj=mock_proxy_logging_obj, proxy_config=mock_proxy_config, - route_type=route_type, + route_type="acompletion", llm_router=self._router_with_embedding_only_model(), ) assert exc_info.value.code == "400" assert exc_info.value.message == ( - f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + "/v1/chat/completions is not supported by model=embed-only. Supported endpoints: /v1/embeddings" ) mock_proxy_logging_obj.pre_call_hook.assert_not_awaited() diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 475aeaa896d..fba3e63aec8 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1963,6 +1963,8 @@ EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment("embed-fallback", E ({}, {"fallbacks": [{"model": "chat-fallback", "temperature": 0.2}]}), ({"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "chat-fallback"}}, {}), ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}), + ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}), + ({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}), ], ) def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): From e8729db3aa137c50fcc0700f8189c1f554bc6683 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:58:28 -0700 Subject: [PATCH 05/21] fix(proxy): walk fallback groups breadth first and release reservations on a rejected WebSocket Each fallback group is looked up once however often the request repeats it, and a WebSocket closed by the endpoint gate runs the failure hook so rate-limit reservations are released --- .../proxy/response_api_endpoints/endpoints.py | 5 +++ litellm/proxy/route_llm_request.py | 41 ++++++++++++------- litellm/router.py | 2 +- .../response_api_endpoints/test_endpoints.py | 9 ++++ tests/unit/proxy/test_route_llm_request.py | 23 +++++++++++ 5 files changed, 64 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 536abfbfac7..5f17a761110 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1571,6 +1571,11 @@ async def _responses_websocket_session( team_id=user_api_key_dict.team_id, ) if unsupported_endpoint is not None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=unsupported_endpoint, + request_data=data, + ) raise unsupported_endpoint except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index b171caf3475..868115ac7b1 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,6 +1,7 @@ import asyncio from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol @@ -491,19 +492,35 @@ def _fallback_target_group(target: object) -> str | None: return named_model if isinstance(named_model, str) else None -def _fallback_target_answers( - group: str | None, +def _next_fallback_groups( + model: str, data: Mapping[str, object], + fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract +) -> tuple[str, ...]: + return tuple( + group + for target in _fallback_targets(model, data, fallbacks) + if has_unattempted_fallback_target((target,), data) and (group := _fallback_target_group(target)) is not None + ) + + +def _frontier_reaches_chat( + frontier: tuple[str, ...], + seen: frozenset[str], + data: Mapping[str, object], + fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract llm_router: LitellmRouter, team_id: str | None, - walked: frozenset[str], ) -> bool: - if group is None or group in walked: + reached: Final = dict.fromkeys(chain.from_iterable(_next_fallback_groups(m, data, fallbacks) for m in frontier)) + fresh: Final = tuple(group for group in reached if group not in seen) + if not fresh: return False - deployments: Final = _group_deployments(group, llm_router, team_id) - if not deployments: - return False - return _any_serves_chat(deployments) or _fallback_would_answer(group, data, llm_router, team_id, walked) + deployments_by_group: Final = tuple((group, _group_deployments(group, llm_router, team_id)) for group in fresh) + if any(deployments and _any_serves_chat(deployments) for _, deployments in deployments_by_group): + return True + expandable: Final = tuple(group for group, deployments in deployments_by_group if deployments) + return _frontier_reaches_chat(expandable, seen | frozenset(fresh), data, fallbacks, llm_router, team_id) def _fallback_would_answer( @@ -511,7 +528,6 @@ def _fallback_would_answer( data: Mapping[str, object], llm_router: LitellmRouter, team_id: str | None, - walked: frozenset[str] = frozenset(), ) -> bool: if data.get("disable_fallbacks") is True: return False @@ -522,12 +538,7 @@ def _fallback_would_answer( fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks) except ValidationError: return True - walked_here: Final = walked | {model} - return any( - has_unattempted_fallback_target((target,), data) - and _fallback_target_answers(_fallback_target_group(target), data, llm_router, team_id, walked_here) - for target in _fallback_targets(model, data, fallbacks) - ) + return _frontier_reaches_chat((model,), frozenset({model}), data, fallbacks, llm_router, team_id) def endpoint_unsupported_by_model( diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..47a00287d5b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1076,7 +1076,7 @@ class Router: self.assistants_config = assistants_config self.search_tools = search_tools or [] self.guardrail_list = guardrail_list or [] - self.deployment_names: list = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 + self.deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local" # default to an in-memory cache diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 2df84eae062..bc086a29109 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -15,6 +15,7 @@ from httpx import Response import litellm from litellm.proxy.proxy_server import app +from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError from litellm.types.llms.openai import ResponsesAPIResponse from tests._master_key import MASTER_KEY @@ -1069,6 +1070,10 @@ class TestResponsesWSDeclaredNonChatModel: return_value=fake_llm_call(), ) as mock_route_request, patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the session reads the proxy's router global + patch( # test-quality-ok: the session reads the proxy's logging global; the failure hook releases reservations + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + ) as mock_failure_hook, ): await responses_websocket_endpoint( websocket=ws, @@ -1079,7 +1084,11 @@ class TestResponsesWSDeclaredNonChatModel: if not rejected: mock_route_request.assert_awaited_once() ws.close.assert_not_awaited() + mock_failure_hook.assert_not_awaited() return + mock_failure_hook.assert_awaited_once() + assert isinstance(mock_failure_hook.await_args.kwargs["original_exception"], ProxyUnsupportedEndpointError) + assert mock_failure_hook.await_args.kwargs["request_data"]["model"] == "embed-only" frame = json.loads(ws.send_text.await_args.args[0]) assert frame["type"] == "error" assert frame["error"]["type"] == "invalid_request_error" diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index fba3e63aec8..13acef2f40b 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -2060,6 +2060,29 @@ def test_endpoint_unsupported_ignores_endpoints_copied_from_the_loaded_catalog(d assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected +class _LookupCountingRouter(Router): + def __init__(self, **router_kwargs: object) -> None: + self.looked_up: list[str | None] = [] + super().__init__(**router_kwargs) + self.looked_up.clear() + + def get_model_list(self, model_name: str | None = None, team_id: str | None = None): + self.looked_up.append(model_name) + return super().get_model_list(model_name=model_name, team_id=team_id) + + +def test_endpoint_unsupported_looks_up_each_repeated_fallback_group_once(): + fallback_groups: Final = ("embed-a", "embed-b", "embed-c") + router: Final = _LookupCountingRouter( + model_list=[_embedding_deployment(name, EMBEDDING_ONLY) for name in ("embed-only", *fallback_groups)] + ) + + error: Final = _unsupported("acompletion", "embed-only", router, fallbacks=list(fallback_groups) * 5) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert sorted(name for name in router.looked_up if name in fallback_groups) == sorted(fallback_groups) + + def _router_with_defaults(default_litellm_params: dict[str, object]): import litellm From b64803ab17bd10418b5f13d8cfb95f83cc626bd6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:17:30 -0700 Subject: [PATCH 06/21] fix(proxy): walk fallback layers without recursion and gate the compact Responses route The fallback walk now folds one layer at a time, so the recursion check passes, and /v1/responses/compact on a model declared non-chat gets the same 400 as the other chat-shaped routes --- litellm/proxy/route_llm_request.py | 51 ++++++++++++++++------ tests/unit/proxy/test_route_llm_request.py | 1 + 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 868115ac7b1..b274121822e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,7 +1,7 @@ import asyncio from collections.abc import Mapping, Sequence from dataclasses import dataclass -from itertools import chain +from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol @@ -375,6 +375,7 @@ CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( "agenerate_content_stream": "streamGenerateContent", "acreate_interaction": "/v1beta/interactions", "_aresponses_websocket": "/v1/responses", + "acompact_responses": "/v1/responses/compact", } ) @@ -504,23 +505,45 @@ def _next_fallback_groups( ) -def _frontier_reaches_chat( - frontier: tuple[str, ...], - seen: frozenset[str], +@dataclass(frozen=True, slots=True) +class _FallbackWalk: + frontier: tuple[str, ...] + seen: frozenset[str] + reaches_chat: bool = False + + +def _walk_one_layer( + walk: _FallbackWalk, + data: Mapping[str, object], + fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract + llm_router: LitellmRouter, + team_id: str | None, +) -> _FallbackWalk: + reached: Final = dict.fromkeys( + chain.from_iterable(_next_fallback_groups(m, data, fallbacks) for m in walk.frontier) + ) + fresh: Final = tuple(group for group in reached if group not in walk.seen) + deployments_by_group: Final = tuple((group, _group_deployments(group, llm_router, team_id)) for group in fresh) + return _FallbackWalk( + frontier=tuple(group for group, deployments in deployments_by_group if deployments), + seen=walk.seen | frozenset(fresh), + reaches_chat=any(deployments and _any_serves_chat(deployments) for _, deployments in deployments_by_group), + ) + + +def _fallbacks_reach_chat( + model: str, data: Mapping[str, object], fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract llm_router: LitellmRouter, team_id: str | None, ) -> bool: - reached: Final = dict.fromkeys(chain.from_iterable(_next_fallback_groups(m, data, fallbacks) for m in frontier)) - fresh: Final = tuple(group for group in reached if group not in seen) - if not fresh: - return False - deployments_by_group: Final = tuple((group, _group_deployments(group, llm_router, team_id)) for group in fresh) - if any(deployments and _any_serves_chat(deployments) for _, deployments in deployments_by_group): - return True - expandable: Final = tuple(group for group, deployments in deployments_by_group if deployments) - return _frontier_reaches_chat(expandable, seen | frozenset(fresh), data, fallbacks, llm_router, team_id) + walks: Final = accumulate( + repeat(None), + lambda walk, _: _walk_one_layer(walk, data, fallbacks, llm_router, team_id), + initial=_FallbackWalk(frontier=(model,), seen=frozenset({model})), + ) + return any(walk.reaches_chat for walk in takewhile(lambda walk: walk.frontier or walk.reaches_chat, walks)) def _fallback_would_answer( @@ -538,7 +561,7 @@ def _fallback_would_answer( fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks) except ValidationError: return True - return _frontier_reaches_chat((model,), frozenset({model}), data, fallbacks, llm_router, team_id) + return _fallbacks_reach_chat(model, data, fallbacks, llm_router, team_id) def endpoint_unsupported_by_model( diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 13acef2f40b..c9128eba8c2 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1781,6 +1781,7 @@ CHAT_SHAPED_ROUTES: Final = ( ("agenerate_content_stream", "streamGenerateContent"), ("acreate_interaction", "/v1beta/interactions"), ("_aresponses_websocket", "/v1/responses"), + ("acompact_responses", "/v1/responses/compact"), ) CHAT_SHAPED_ROUTE_TYPES: Final = tuple(route_type for route_type, _ in CHAT_SHAPED_ROUTES) From e2ab40bf532cfa97e3970630799e8c4ebb7e7063 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:37:09 -0700 Subject: [PATCH 07/21] fix(proxy): keep chat on rows whose cost map entry serves chat A row that carries the cost map key of an entry that serves chat counts as declaring nothing, so a catalog edit that adds a chat path never turns old Admin UI copies into 400s --- litellm/proxy/route_llm_request.py | 27 +++++++++++---- tests/unit/proxy/test_route_llm_request.py | 39 ++++++++++++++++++++++ 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index b274121822e..37f292816f7 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -416,14 +416,27 @@ class _RouterFallbacks(Protocol): def fallbacks(self) -> Sequence[object] | None: ... -def _echoed_from_cost_map(model_info: Mapping[str, object]) -> frozenset[str]: +def _loaded_catalog_entry(model_info: Mapping[str, object]) -> Mapping[str, object]: catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) if not isinstance(catalog_key, str): - return frozenset() - return frozenset( - echoed_cost_map_fields( - model_info, GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}) - ) + return MappingProxyType({}) + return GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}) + + +def _catalog_serves_chat(catalog_entry: Mapping[str, object]) -> bool: + try: + declared: Final = _DeploymentEndpointDeclaration.model_validate(dict(catalog_entry)) + except ValidationError: + return False + return (declared.mode or "").strip().lower() in CHAT_SHAPED_MODES or not CHAT_SHAPED_ENDPOINTS.isdisjoint( + declared.supported_endpoints or () + ) + + +def _declaration_overruled_by_catalog(model_info: Mapping[str, object]) -> bool: + catalog_entry: Final = _loaded_catalog_entry(model_info) + return _catalog_serves_chat(catalog_entry) or "supported_endpoints" in echoed_cost_map_fields( + model_info, catalog_entry ) @@ -439,7 +452,7 @@ def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[ endpoints: Final = frozenset(declared.supported_endpoints) if not endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS): return None - return None if "supported_endpoints" in _echoed_from_cost_map(model_info) else endpoints + return None if _declaration_overruled_by_catalog(model_info) else endpoints def _group_deployments(model: str, llm_router: LitellmRouter, team_id: str | None) -> tuple[Mapping[str, object], ...]: diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index c9128eba8c2..5e4b5452bfd 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -2061,6 +2061,45 @@ def test_endpoint_unsupported_ignores_endpoints_copied_from_the_loaded_catalog(d assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected +def _catalog_key_serving_chat_by_mode() -> str: + return min( + key + for key, entry in GetModelCostMap.loaded_model_cost_map().items() + if entry.get("mode") == "chat" and "supported_endpoints" not in entry + ) + + +def _catalog_key_serving_chat_by_endpoints() -> str: + return min( + key + for key, entry in GetModelCostMap.loaded_model_cost_map().items() + if entry.get("mode") not in CHAT_SHAPED_MODES + and isinstance(endpoints := entry.get("supported_endpoints"), list) + and not CHAT_SHAPED_ENDPOINTS.isdisjoint(endpoints) + ) + + +@pytest.mark.parametrize( + "catalog_key, rejected", + [ + pytest.param(_catalog_key_serving_chat_by_mode, False, id="key-of-a-chat-mode-catalog-entry"), + pytest.param(_catalog_key_serving_chat_by_endpoints, False, id="key-of-a-catalog-entry-listing-chat"), + pytest.param(_catalog_key_declaring_non_chat, True, id="key-of-a-non-chat-catalog-entry"), + pytest.param(lambda: "not-in-the-catalog", True, id="key-missing-from-the-catalog"), + ], +) +def test_endpoint_unsupported_ignores_non_chat_values_on_a_row_whose_catalog_entry_serves_chat(catalog_key, rejected): + router: Final = _router_with( + { + "model_name": "catalog-row", + "litellm_params": {"model": "openai/catalog-row", "api_key": "test-key"}, + "model_info": {"key": catalog_key(), **EMBEDDING_ONLY}, + } + ) + + assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected + + class _LookupCountingRouter(Router): def __init__(self, **router_kwargs: object) -> None: self.looked_up: list[str | None] = [] From 4accd0c2f27f9357e0130ae528093d10280b4f3e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:33:37 -0700 Subject: [PATCH 08/21] fix(proxy): check chat sent to a deployment id or backend model name --- litellm/proxy/route_llm_request.py | 39 +++++++++++++++++---- tests/unit/proxy/test_route_llm_request.py | 40 +++++++++++++++++++++- 2 files changed, 71 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 37f292816f7..972eb6dda46 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -381,14 +381,10 @@ CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( CHAT_SHAPED_ENDPOINTS: Final[frozenset[str]] = frozenset( { - "/v1/chat/completions", "/chat/completions", - "/v1/completions", "/completions", - "/v1/messages", - "/v1/responses", "/responses", - "/v1beta/interactions", + *CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE.values(), } ) @@ -455,8 +451,37 @@ def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[ return None if _declaration_overruled_by_catalog(model_info) else endpoints +def _deployment_by_id(model_id: str, llm_router: LitellmRouter) -> tuple[Mapping[str, object], ...]: + try: + return (JSON_OBJECT_ADAPTER.validate_python(llm_router.get_model_info(id=model_id)),) + except ValidationError: + return () + + def _group_deployments(model: str, llm_router: LitellmRouter, team_id: str | None) -> tuple[Mapping[str, object], ...]: - return tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) + return tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) or _deployment_by_id( + model, llm_router + ) + + +def _deployments_by_backend_name(model: str, llm_router: LitellmRouter) -> tuple[Mapping[str, object], ...]: + if ( + model not in llm_router.deployment_names + or llm_router.router_general_settings.pass_through_all_models + or llm_router.pattern_router.patterns + ): + return () + return tuple( + deployment + for deployment in llm_router.get_model_list() or () + if deployment["litellm_params"].get("model") == model + ) + + +def _requested_deployments( + model: str, llm_router: LitellmRouter, team_id: str | None +) -> tuple[Mapping[str, object], ...]: + return _group_deployments(model, llm_router, team_id) or _deployments_by_backend_name(model, llm_router) def _any_serves_chat(deployments: Sequence[Mapping[str, object]]) -> bool: @@ -587,7 +612,7 @@ def endpoint_unsupported_by_model( model: Final = data.get("model") if endpoint is None or not isinstance(model, str) or not model or llm_router is None: return None - deployments: Final = _group_deployments(model, llm_router, team_id) + deployments: Final = _requested_deployments(model, llm_router, team_id) if not deployments or _any_serves_chat(deployments): return None if _fallback_would_answer(model, data, llm_router, team_id): diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 5e4b5452bfd..709b4569bbe 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -10,6 +10,7 @@ from fastapi import HTTPException import litellm from litellm import Router from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.types.router import RouterGeneralSettings from litellm.proxy.route_llm_request import ( CHAT_SHAPED_ENDPOINTS, CHAT_SHAPED_MODES, @@ -1877,6 +1878,9 @@ def test_endpoint_unsupported_needs_a_non_chat_mode_to_reject(route_type, model_ ["/v1/completions"], ["/completions"], ["/v1beta/interactions"], + ["generateContent"], + ["streamGenerateContent"], + ["/v1/responses/compact"], ], ) def test_endpoint_unsupported_allows_every_chat_route_when_any_chat_endpoint_is_declared(route_type, declared): @@ -1933,6 +1937,35 @@ def test_endpoint_unsupported_follows_a_model_group_alias(): assert error.message == "/v1/messages is not supported by model=embed-alias. Supported endpoints: /v1/embeddings" +@pytest.mark.parametrize("model", ["embed-1", "sagemaker/openai/embed-endpoint"]) +def test_endpoint_unsupported_rejects_chat_sent_to_the_deployment_id_or_backend_model(model): + router: Final = _router_with( + _embedding_deployment("embed-only", {**EMBEDDING_ONLY, "id": "embed-1"}), CHAT_FALLBACK_DEPLOYMENT + ) + + error: Final = _unsupported("acompletion", model, router) + + assert isinstance(error, ProxyUnsupportedEndpointError) + assert error.message == ( + f"/v1/chat/completions is not supported by model={model}. Supported endpoints: /v1/embeddings" + ) + + +@pytest.mark.parametrize( + "extra_deployments, router_kwargs", + [ + ((), {"router_general_settings": RouterGeneralSettings(pass_through_all_models=True)}), + (({"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": "test-key"}},), {}), + ], +) +def test_endpoint_unsupported_leaves_a_backend_model_to_the_route_that_serves_it(extra_deployments, router_kwargs): + router: Final = _router_with( + _embedding_deployment("embed-only", EMBEDDING_ONLY), *extra_deployments, **router_kwargs + ) + + assert _unsupported("acompletion", "sagemaker/openai/embed-endpoint", router) is None + + @pytest.mark.parametrize("model", ["unknown-model", "", None]) def test_endpoint_unsupported_ignores_models_without_deployments(model): router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) @@ -1947,9 +1980,12 @@ def test_endpoint_unsupported_skips_the_check_without_a_router(): CHAT_FALLBACK_DEPLOYMENT: Final = { "model_name": "chat-fallback", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, + "model_info": {"id": "chat-fallback-id"}, } -EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment("embed-fallback", EMBEDDING_ONLY) +EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment( + "embed-fallback", {**EMBEDDING_ONLY, "id": "embed-fallback-id"} +) @pytest.mark.parametrize( @@ -1966,6 +2002,7 @@ EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment("embed-fallback", E ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}), ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}), ({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}), + ({}, {"fallbacks": ["chat-fallback-id"]}), ], ) def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): @@ -1995,6 +2032,7 @@ def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, r ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), ({"default_fallbacks": ["embed-fallback"]}, {}), ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), + ({}, {"fallbacks": ["embed-fallback-id"]}), ({}, {"fallbacks": [{"model": "embed-fallback", "temperature": 0.2}]}), ({"fallbacks": [{"embed-only": ["embed-alias"]}], "model_group_alias": {"embed-alias": "embed-fallback"}}, {}), ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["embed-only"]}]}, {}), From 65c34c9d982d09fc96d19117793f1155e8315d02 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 01:50:36 -0700 Subject: [PATCH 09/21] fix(proxy): check a model the read-through loads from the database --- litellm/proxy/route_llm_request.py | 8 +++++ tests/unit/proxy/test_route_llm_request.py | 34 ++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 972eb6dda46..2d0e70f6c7c 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -901,6 +901,14 @@ async def route_request( if not await model_registry_read_through.attempt(requested_model): raise + unsupported_endpoint: Final = endpoint_unsupported_by_model( + route_type=route_type, + data=JSON_OBJECT_ADAPTER.validate_python(data), + llm_router=proxy_server.llm_router, + team_id=user_api_key_dict.team_id if user_api_key_dict is not None else None, + ) + if unsupported_endpoint is not None: + raise unsupported_endpoint return await _route_request_single_attempt( data=data, llm_router=proxy_server.llm_router, diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 709b4569bbe..4a5d0b60943 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1532,6 +1532,40 @@ async def test_route_request_read_through_recovers_model_created_on_sibling_repl assert table.find_many_wheres[0] == {"model_name": model_name} +@pytest.mark.asyncio +async def test_route_request_checks_the_endpoints_of_a_model_loaded_by_read_through(monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + model_name: Final = "embed-only-on-sibling-replica" + row: Final = SimpleNamespace( + model_id=f"{model_name}-id", + model_name=model_name, + litellm_params={"model": "openai/text-embedding-3-small", "api_key": "fake"}, + model_info={"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, + blocked=False, + ) + router: Final = litellm.Router( + model_list=[{"model_name": "some-other-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] + ) + fake_prisma, table = _fake_prisma_client_with_models([row]) + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with pytest.raises(ProxyUnsupportedEndpointError) as exc_info: + await route_request( + data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, + llm_router=router, + user_model=None, + route_type="acompletion", + ) + + assert exc_info.value.message == ( + f"/v1/chat/completions is not supported by model={model_name}. Supported endpoints: /v1/embeddings" + ) + assert table.find_many_wheres == [{"model_name": model_name}] + + @pytest.mark.asyncio async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch): import litellm From 19dabb6c265f334c2a324ab149ddec2587c8395f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 02:47:53 -0700 Subject: [PATCH 10/21] fix(proxy): step aside for fallbacks the check cannot follow and wildcards that serve a backend name A fallback target with no deployments of its own, a client routing stamp, and a wildcard that does not match the requested backend name no longer turn a working request into a 400 --- litellm/proxy/route_llm_request.py | 6 +- .../test_non_chat_model_endpoint_gate.py | 742 +++++++++++++++++- tests/unit/proxy/test_route_llm_request.py | 33 +- 3 files changed, 760 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 2d0e70f6c7c..f45171ec34e 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -468,7 +468,7 @@ def _deployments_by_backend_name(model: str, llm_router: LitellmRouter) -> tuple if ( model not in llm_router.deployment_names or llm_router.router_general_settings.pass_through_all_models - or llm_router.pattern_router.patterns + or llm_router.pattern_router.get_deployments_by_pattern(model=model) ): return () return tuple( @@ -565,7 +565,7 @@ def _walk_one_layer( return _FallbackWalk( frontier=tuple(group for group, deployments in deployments_by_group if deployments), seen=walk.seen | frozenset(fresh), - reaches_chat=any(deployments and _any_serves_chat(deployments) for _, deployments in deployments_by_group), + reaches_chat=any(not deployments or _any_serves_chat(deployments) for _, deployments in deployments_by_group), ) @@ -599,6 +599,8 @@ def _fallback_would_answer( fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks) except ValidationError: return True + if fallback_lookup_groups(data, model) != (model,): + return True return _fallbacks_reach_chat(model, data, fallbacks, llm_router, team_id) diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 43cf611a0c5..fdef4f579f2 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -2,7 +2,7 @@ import asyncio import contextlib import json import uuid -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -41,12 +41,76 @@ _EMBEDDING_MODEL: Final = "text-embedding-3-small" _CHAT_MODEL: Final = "gpt-4o-mini" _SERVED_TEXT: Final = "served by the chat upstream" _VECTOR: Final = (0.25, 0.5, 0.75) +_COUNTED_TOKENS: Final = 17 _CLOSE: Final = MappingProxyType({"connection": "close"}) _RELOAD_SECONDS: Final = 5 _LONG_ENDPOINT: Final = "/v1/embeddings/" + "x" * 5000 _EMBED_ONLY: Final[Mapping[str, JsonValue]] = MappingProxyType( {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]} ) +_CHAT_ENDPOINTS: Final = frozenset( + { + "/v1/chat/completions", + "/chat/completions", + "/v1/completions", + "/completions", + "/v1/messages", + "/v1/responses", + "/responses", + "/v1beta/interactions", + } +) +_CHAT_MODES: Final = frozenset({"chat", "completion", "responses"}) +_CATALOG: Final = TypeAdapter(dict[str, JsonValue]).validate_json( + (Path(__file__).resolve().parents[3] / "litellm/model_prices_and_context_window_backup.json").read_bytes() +) +_CATALOG_ENTRIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + key: object_value(entry) + for key, entry in sorted(_CATALOG.items()) + if isinstance(entry, dict) and key != "sample_spec" + } +) + + +def _listed_endpoints(entry: Mapping[str, JsonValue]) -> frozenset[str] | None: + listed: Final = entry.get("supported_endpoints") + if not isinstance(listed, list) or not listed or not all(isinstance(item, str) for item in listed): + return None + return frozenset(string_value(item) for item in listed) + + +def _non_chat_mode(entry: Mapping[str, JsonValue]) -> bool: + mode: Final = entry.get("mode") + return isinstance(mode, str) and mode not in _CHAT_MODES + + +def _lists_only_other_endpoints(entry: Mapping[str, JsonValue]) -> bool: + listed: Final = _listed_endpoints(entry) + return ( + _non_chat_mode(entry) + and listed is not None + and listed.isdisjoint(_CHAT_ENDPOINTS) + and listed != {"/v1/embeddings"} + ) + + +def _lists_a_chat_endpoint(entry: Mapping[str, JsonValue]) -> bool: + listed: Final = _listed_endpoints(entry) + return _non_chat_mode(entry) and listed is not None and not listed.isdisjoint(_CHAT_ENDPOINTS) + + +def _catalog_key(matches: Callable[[Mapping[str, JsonValue]], bool]) -> str: + return next(key for key, entry in _CATALOG_ENTRIES.items() if matches(entry)) + + +_CATALOG_CHAT_KEY: Final = _catalog_key(lambda entry: entry.get("mode") == "chat") +_CATALOG_EMBEDDING_KEY: Final = _catalog_key( + lambda entry: entry.get("mode") == "embedding" and "supported_endpoints" not in entry +) +_CATALOG_OTHER_ENDPOINTS_KEY: Final = _catalog_key(_lists_only_other_endpoints) +_CATALOG_CHAT_LISTED_KEY: Final = _catalog_key(_lists_a_chat_endpoint) +_CATALOG_OTHER_ENDPOINTS_ENTRY: Final = _CATALOG_ENTRIES[_CATALOG_OTHER_ENDPOINTS_KEY] def _embedding_reply() -> Reply: @@ -92,11 +156,13 @@ def _respond(request: Request) -> Reply: streamed: Final = body.get("stream") is True if path == "/v1/embeddings": return _embedding_reply() + if path == "/v1/responses/input_tokens": + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS}).encode()) if body.get("model") != _CHAT_MODEL: return _not_a_chat_model() if path == "/v1/chat/completions": return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) - if path == "/v1/responses": + if path in ("/v1/responses", "/v1/responses/compact"): return responses_reply(f"resp_{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) if path == "/v1/completions": return _text_completion_reply() @@ -153,6 +219,19 @@ _GATED: Final[Mapping[str, tuple[_Deployment, str]]] = MappingProxyType( ), "/embeddings, /v1/embeddings", ), + "catalog_key_lists_other_endpoints": ( + _single({**_EMBED_ONLY, "key": _CATALOG_OTHER_ENDPOINTS_KEY}), + "/v1/embeddings", + ), + "catalog_key_names_an_embedding_entry": ( + _single({**_EMBED_ONLY, "key": _CATALOG_EMBEDDING_KEY}), + "/v1/embeddings", + ), + "catalog_key_not_in_the_catalog": ( + _single({**_EMBED_ONLY, "key": "gate-uncatalogued-entry"}), + "/v1/embeddings", + ), + "catalog_key_not_a_string": (_single({**_EMBED_ONLY, "key": 7}), "/v1/embeddings"), } ) @@ -176,14 +255,27 @@ _UNGATED: Final[Mapping[str, _Deployment]] = MappingProxyType( ), "lists_messages": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/messages"]}), "lists_bare_responses": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/responses"]}), + "lists_generate_content": _single({"mode": "embedding", "supported_endpoints": ["generateContent"]}), + "lists_responses_compact": _single({"mode": "embedding", "supported_endpoints": ["/v1/responses/compact"]}), "mixed_group": _Deployment(f"openai/{_CHAT_MODEL}", (_EMBED_ONLY, {"mode": "chat"})), + "catalog_key_serves_chat": _single({**_EMBED_ONLY, "key": _CATALOG_CHAT_KEY}), + "catalog_key_lists_a_chat_endpoint": _single({**_EMBED_ONLY, "key": _CATALOG_CHAT_LISTED_KEY}), + "catalog_endpoints_echoed": _single( + { + "key": _CATALOG_OTHER_ENDPOINTS_KEY, + "mode": _CATALOG_OTHER_ENDPOINTS_ENTRY["mode"], + "supported_endpoints": _CATALOG_OTHER_ENDPOINTS_ENTRY["supported_endpoints"], + } + ), } ) _SHARED: Final[Mapping[str, _Deployment]] = MappingProxyType( { "embedding_only": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "embedding_only_second": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), "chat_target": _single({"mode": "chat"}), + "chat_rate_limited": _single({"mode": "chat"}), **{name: deployment for name, (deployment, _supported) in _GATED.items()}, **_UNGATED, "mode_not_a_string": _single( @@ -300,13 +392,19 @@ def _failure_row(call_id: str) -> dict[str, JsonValue]: return rows[0] -def _assert_rejected(rig: _Rig, response: httpx.Response, marker: str, expected: dict[str, JsonValue]) -> None: +def _assert_rejected_before( + upstream: _Upstream, response: httpx.Response, marker: str, expected: dict[str, JsonValue] +) -> None: assert response.status_code == 400, response.text assert JSON_OBJECT.validate_json(response.content) == expected, response.text - assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) + assert upstream.calls(marker) == (), upstream.targets(marker) assert _failure_row(response.headers["x-litellm-call-id"])["status"] == "failure" +def _assert_rejected(rig: _Rig, response: httpx.Response, marker: str, expected: dict[str, JsonValue]) -> None: + _assert_rejected_before(rig.upstream, response, marker, expected) + + @dataclass(frozen=True, slots=True) class _Route: path: str @@ -541,6 +639,20 @@ async def _websocket_first_frame_and_close( return first, ws.close_code, ws.close_reason +def _websocket_rejection(model: str) -> tuple[dict[str, JsonValue], int | None, str | None]: + return ( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"/v1/responses is not supported by model={model}. Supported endpoints: /v1/embeddings", + }, + }, + 1008, + "Pre-call error", + ) + + @pytest.mark.parametrize("aliased", [False, True]) def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstream(rig: _Rig, aliased: bool) -> None: target: Final = rig.names["embedding_only"] @@ -548,20 +660,39 @@ def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstre marker: Final = _marker() with rig.gateway.scenario() as scenario: key: Final = scenario.key(aliases={alias: target}) - frame, close_code, close_reason = asyncio.run( + closed: Final = asyncio.run( _websocket_first_frame_and_close(rig.gateway, key, alias if aliased else target, marker) ) - assert frame == { - "type": "error", - "error": { - "type": "invalid_request_error", - "message": f"/v1/responses is not supported by model={target}. Supported endpoints: /v1/embeddings", - }, - } - assert (close_code, close_reason) == (1008, "Pre-call error") + assert closed == _websocket_rejection(target) assert rig.upstream.calls(marker) == () +_RESERVING_KEYS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "max_parallel_requests": {"max_parallel_requests": 1}, + "tpm_limit": {"tpm_limit": 100, "default_estimated_output_tokens": 80}, + } +) + + +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_responses_websocket_rejections_leave_the_key_reservation_free(rig: _Rig, limit: str) -> None: + target: Final = rig.names["embedding_only"] + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rig.gateway, key, target, marker)) for marker in markers + ) + assert closed == (_websocket_rejection(target),) * 3 + served: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["chat_target"], served_marker), key=key + ) + _assert_served(rig, served, served_marker, "/v1/chat/completions") + assert [rig.upstream.calls(marker) for marker in markers] == [()] * 3 + + def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None: public: Final = f"gate-team-{uuid.uuid4().hex[:12]}" marker: Final = _marker() @@ -652,7 +783,7 @@ def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: def test_fallbacks_that_cannot_answer_chat_keep_the_rejection(rig: _Rig) -> None: model: Final = rig.names["embedding_only"] marker: Final = _marker() - fallbacks: Final = [rig.names["image_generation"], model, f"gate-missing-{uuid.uuid4().hex}"] + fallbacks: Final = [rig.names["image_generation"], model] response: Final = rig.gateway.request( "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": fallbacks} ) @@ -671,6 +802,203 @@ def test_request_fallbacks_do_not_rescue_an_unknown_primary(rig: _Rig) -> None: assert rig.upstream.calls(marker) == () +def _chat(rig: _Rig, model: str, marker: str, *, key: str | None = None, **extra: JsonValue) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", {**_chat_body(model, marker), **extra}, key=key) + + +def _assert_rejected_chat(rig: _Rig, response: httpx.Response, marker: str, model: str) -> None: + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_a_fallback_chain_that_reaches_chat_two_hops_away_answers(rig: _Rig) -> None: + first, second, chat = rig.names["embedding_only"], rig.names["embedding_only_second"], rig.names["chat_target"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [chat]}]) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) * 3 + + +def test_a_fallback_chain_that_loops_between_embedding_only_groups_keeps_the_rejection(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [first]}]) + _assert_rejected_chat(rig, response, marker, first) + + +def test_a_specific_chain_that_loops_beats_a_catch_all_that_could_answer(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat( + rig, first, marker, fallbacks=[{"*": [rig.names["chat_target"]]}, {first: [second]}, {second: [first]}] + ) + _assert_rejected_chat(rig, response, marker, first) + + +def test_disable_fallbacks_keeps_the_rejection_over_request_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, fallbacks=[rig.names["chat_target"]], disable_fallbacks=True) + _assert_rejected_chat(rig, response, marker, model) + + +def test_client_sent_metadata_disable_fallbacks_does_not_block_the_fallback(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat( + rig, model, marker, fallbacks=[rig.names["chat_target"]], metadata={"_disable_fallbacks": True} + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +_MALFORMED_FALLBACKS: Final[Mapping[str, Callable[[_Rig], JsonValue]]] = MappingProxyType( + { + "a_string": lambda rig: rig.names["chat_target"], + "an_object": lambda rig: {rig.names["embedding_only"]: [rig.names["chat_target"]]}, + } +) + + +@pytest.mark.parametrize("shape", tuple(_MALFORMED_FALLBACKS)) +def test_malformed_request_fallbacks_step_aside_for_the_router(rig: _Rig, shape: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, fallbacks=_MALFORMED_FALLBACKS[shape](rig)) + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.targets(marker)[:1] == ("/v1/chat/completions",), response.text + + +def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: + ghost: Final = f"gate-ghost-{uuid.uuid4().hex}" + return {"fallbacks": [{rig.names["embedding_only"]: [ghost]}, {ghost: [rig.names["chat_target"]]}]} + + +def _with_a_client_routing_stamp(rig: _Rig) -> Mapping[str, JsonValue]: + model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + return { + "fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], + "metadata": {"model_group": second}, + } + + +_FALLBACKS_THE_CHECK_CANNOT_FOLLOW: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( + { + "a_group_with_no_deployments": _through_a_group_with_no_deployments, + "a_client_routing_stamp": _with_a_client_routing_stamp, + } +) + + +@pytest.mark.parametrize("shape", tuple(_FALLBACKS_THE_CHECK_CANNOT_FOLLOW)) +def test_fallbacks_the_check_cannot_follow_step_aside_for_the_router(rig: _Rig, shape: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, **_FALLBACKS_THE_CHECK_CANNOT_FOLLOW[shape](rig)) + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.targets(marker)[:1] == ("/v1/chat/completions",), response.text + + +def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=scenario.key(team_id=team)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +def test_key_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=key) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +def test_key_router_settings_that_cannot_answer_override_team_fallbacks_that_could(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + key: Final = scenario.key( + team_id=team, router_settings={"fallbacks": [{model: [rig.names["embedding_only_second"]]}]} + ) + response: Final = _chat(rig, model, marker, key=key) + _assert_rejected_chat(rig, response, marker, model) + + +def test_an_empty_request_fallback_list_overrides_key_router_settings_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + response: Final = _chat(rig, model, marker, key=key, fallbacks=[]) + _assert_rejected_chat(rig, response, marker, model) + + +def test_local_rate_limit_fallback_skips_an_embedding_only_hop(rig: _Rig) -> None: + limited, chat, model = rig.names["chat_rate_limited"], rig.names["chat_target"], rig.names["embedding_only"] + first, second = _marker(), _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: [model, chat]}]} + ) + served_first: Final = _chat(rig, limited, first, key=key) + served_second: Final = _chat(rig, limited, second, key=key) + assert (served_first.status_code, served_first.headers.get("x-litellm-model-group")) == (200, limited), ( + served_first.text + ) + assert (served_second.status_code, served_second.headers.get("x-litellm-model-group")) == (200, chat), ( + served_second.text + ) + assert (rig.upstream.targets(first), rig.upstream.targets(second)) == ( + ("/v1/chat/completions",), + ("/v1/chat/completions",), + ) + + +_COMPACT_PATHS: Final = ("/v1/responses/compact", "/responses/compact", "/openai/v1/responses/compact") + + +def _compact_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": [{"role": "user", "content": marker}]} + + +@pytest.mark.parametrize("path", _COMPACT_PATHS) +def test_responses_compact_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, path: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", path, _compact_body(model, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/responses/compact", model, "/v1/embeddings")) + + +def test_responses_compact_keeps_reaching_the_upstream_for_a_chat_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/compact", _compact_body(rig.names["chat_target"], marker) + ) + _assert_served(rig, response, marker, "/v1/responses/compact") + + +def test_count_tokens_keeps_counting_for_an_embedding_only_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": rig.names["embedding_only"], "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNTED_TOKENS, response.text + assert rig.upstream.targets(marker) == ("/v1/responses/input_tokens",) + + def test_text_completions_reject_an_embedding_only_model_before_the_upstream(rig: _Rig) -> None: model: Final = rig.names["embedding_only"] marker: Final = _marker() @@ -927,3 +1255,389 @@ def test_yaml_embedding_only_model_stays_gated_through_an_upstream_pause_and_a_w for _ in range(6) ) _assert_gated_burst(respawned, embed[0]) + + +@dataclass(frozen=True, slots=True) +class _Routed: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + + +def _routed_entry(name: str, api_base: str, litellm_model: str, info: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": litellm_model, "api_base": api_base, "api_key": _PROVIDER_KEY}, + "model_info": dict(info), + } + + +def _routed_config(directory: Path, api_base: str, names: Mapping[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["embed"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "default_fallbacks": [names["chat"]], + "model_group_alias": {names["alias_embed"]: names["embed"], names["alias_chat"]: names["chat"]}, + }, + } + path: Final = directory / "gate-routed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def routed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Routed]: + directory: Final = tmp_path_factory.mktemp("gate-routed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + {label: f"gate-routed-{label}-{suffix}" for label in ("embed", "chat", "alias_embed", "alias_chat")} + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, directory, {}, config=_routed_config(directory, wire.url + "/v1", names), workers=2 + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Routed(owned.gateway, _Upstream(wire), names) + + +_ROUTED_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +_DEFAULT_FALLBACK_ROUTES: Final[Mapping[str, str]] = MappingProxyType( + {"chat_v1": "/v1/chat/completions", "messages": "/v1/responses", "responses_v1": "/v1/responses"} +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", tuple(_DEFAULT_FALLBACK_ROUTES)) +def test_router_default_fallbacks_answer_for_an_embedding_only_group(routed: _Routed, route: str) -> None: + marker: Final = _marker() + response: Final = _send(routed.gateway, _ROUTES[route], routed.names["embed"], marker) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == (_DEFAULT_FALLBACK_ROUTES[route],) * 2 + + +_OPT_OUTS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "empty_fallbacks": {"fallbacks": []}, + "null_fallbacks": {"fallbacks": None}, + "disable_fallbacks": {"disable_fallbacks": True}, + } +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("opt_out", tuple(_OPT_OUTS)) +def test_request_opt_outs_keep_the_rejection_over_router_default_fallbacks(routed: _Routed, opt_out: str) -> None: + model: Final = routed.names["embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), **_OPT_OUTS[opt_out]} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_falls_back_to_the_default(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_embed"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_is_rejected_without_fallbacks(routed: _Routed) -> None: + alias: Final = routed.names["alias_embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(alias, marker), "disable_fallbacks": True} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", alias, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_a_chat_group_keeps_reaching_the_upstream(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_chat"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_keeps_serving_embeddings(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(routed.gateway, routed.names["alias_embed"], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert routed.upstream.targets(marker) == ("/v1/embeddings",) + + +_REWRITER: Final = """from litellm.integrations.custom_guardrail import CustomGuardrail + + +class WebSocketModelRewriter(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if call_type == "_aresponses_websocket" and data.get("model") == "__SOURCE__": + data["model"] = "__TARGET__" + return data +""" + + +@dataclass(frozen=True, slots=True) +class _Rewritten: + gateway: Gateway + upstream: _Upstream + source: str + target: str + + +def _rewritten_config(directory: Path, api_base: str, source: str, target: str, module: str) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(source, api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(target, api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + ], + "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, + "guardrails": [ + { + "guardrail_name": f"{module}-guardrail", + "litellm_params": { + "guardrail": f"{module}.WebSocketModelRewriter", + "mode": "pre_call", + "default_on": True, + }, + } + ], + } + path: Final = directory / "gate-rewritten-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def rewritten(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rewritten]: + directory: Final = tmp_path_factory.mktemp("gate-rewritten") + suffix: Final = uuid.uuid4().hex[:12] + source: Final = f"gate-rewrite-chat-{suffix}" + target: Final = f"gate-rewrite-embed-{suffix}" + module: Final = f"gate_ws_rewriter_{suffix}" + (directory / f"{module}.py").write_text(_REWRITER.replace("__SOURCE__", source).replace("__TARGET__", target)) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_rewritten_config(directory, wire.url + "/v1", source, target, module), + workers=2, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(source, "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Rewritten(owned.gateway, _Upstream(wire), source, target) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_a_websocket_model_a_hook_rewrites_to_an_embedding_only_group_releases_the_key_reservation( + rewritten: _Rewritten, limit: str +) -> None: + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rewritten.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rewritten.gateway, key, rewritten.source, marker)) + for marker in markers + ) + assert closed == (_websocket_rejection(rewritten.target),) * 3 + served: Final = rewritten.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rewritten.source, served_marker), key=key + ) + assert served.status_code == 200, served.text + assert _SERVED_TEXT in served.text, served.text + assert [rewritten.upstream.calls(marker) for marker in markers] == [()] * 3 + assert rewritten.upstream.targets(served_marker) == ("/v1/chat/completions",) + + +@dataclass(frozen=True, slots=True) +class _Addressed: + gateway: Gateway + rig: Gateway + upstream: _Upstream + names: Mapping[str, str] + addresses: Mapping[str, str] + api_base: str + + +def _addressed_config(directory: Path, api_base: str, names: Mapping[str, str], suffix: str) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry( + names["embed"], api_base, f"openai/gate-addr-backend-{suffix}", {**_EMBED_ONLY, "id": names["embed_id"]} + ), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat", "id": names["chat_id"]}), + _routed_entry(names["shared_embed"], api_base, f"openai/{_CHAT_MODEL}", _EMBED_ONLY), + ], + "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, + "general_settings": { + **object_value(shared["general_settings"]), + "proxy_config_reload_interval_seconds": 3600, + }, + } + path: Final = directory / "gate-addressed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + directory: Final = tmp_path_factory.mktemp("gate-addressed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + {label: f"gate-addr-{label}-{suffix}" for label in ("embed", "embed_id", "chat", "chat_id", "shared_embed")} + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_addressed_config(directory, wire.url + "/v1", names, suffix), + workers=2, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Addressed( + owned.gateway, + rig_gateway, + _Upstream(wire), + names, + MappingProxyType( + {"deployment_id": names["embed_id"], "backend_name": f"openai/gate-addr-backend-{suffix}"} + ), + wire.url + "/v1", + ) + + +_ADDRESSED_ROUTES: Final = ("chat_v1", "messages", "responses_v1") + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", _ADDRESSED_ROUTES) +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_chat_sent_to_an_embedding_only_deployment_id_or_backend_name_is_rejected_before_the_upstream( + addressed: _Addressed, address: str, route: str +) -> None: + spec: Final = _ROUTES[route] + model: Final = addressed.addresses[address] + marker: Final = _marker() + response: Final = _send(addressed.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected_before(addressed.upstream, response, marker, expected) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_embeddings_sent_to_an_embedding_only_deployment_id_or_backend_name_keep_working( + addressed: _Addressed, address: str +) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(addressed.gateway, addressed.addresses[address], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert addressed.upstream.targets(marker) == ("/v1/embeddings",) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("chat_deployment_id", "backend_name_a_chat_deployment_shares")) +def test_chat_sent_to_a_chat_capable_deployment_id_or_backend_name_keeps_reaching_the_upstream( + addressed: _Addressed, address: str +) -> None: + model: Final = addressed.names["chat_id"] if address == "chat_deployment_id" else f"openai/{_CHAT_MODEL}" + marker: Final = _marker() + response: Final = addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_a_chat_deployment_id_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["chat_id"]]}, + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + + +@_ROUTED_TIMEOUT +def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rejected_on_every_call( + addressed: _Addressed, +) -> None: + name: Final = f"gate-addr-loaded-{uuid.uuid4().hex[:12]}" + markers: Final = tuple(_marker() for _ in range(6)) + with addressed.rig.scenario() as scenario: + identity: Final = _create( + addressed.rig, name, {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": addressed.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + responses: Final = tuple( + addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(name, marker), headers=_CLOSE) + for marker in markers + ) + expected: Final = _openai_error("/v1/chat/completions", name, "/v1/embeddings") + assert [(response.status_code, JSON_OBJECT.validate_json(response.content)) for response in responses] == [ + (400, expected) + ] * len(markers), [response.text for response in responses] + assert [addressed.upstream.calls(marker) for marker in markers] == [()] * len(markers) diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 4a5d0b60943..b7189f3990d 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1971,10 +1971,19 @@ def test_endpoint_unsupported_follows_a_model_group_alias(): assert error.message == "/v1/messages is not supported by model=embed-alias. Supported endpoints: /v1/embeddings" +UNRELATED_WILDCARD_DEPLOYMENT: Final = { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "test-key"}, +} + + @pytest.mark.parametrize("model", ["embed-1", "sagemaker/openai/embed-endpoint"]) -def test_endpoint_unsupported_rejects_chat_sent_to_the_deployment_id_or_backend_model(model): +@pytest.mark.parametrize("extra_deployments", [(), (UNRELATED_WILDCARD_DEPLOYMENT,)]) +def test_endpoint_unsupported_rejects_chat_sent_to_the_deployment_id_or_backend_model(model, extra_deployments): router: Final = _router_with( - _embedding_deployment("embed-only", {**EMBEDDING_ONLY, "id": "embed-1"}), CHAT_FALLBACK_DEPLOYMENT + _embedding_deployment("embed-only", {**EMBEDDING_ONLY, "id": "embed-1"}), + CHAT_FALLBACK_DEPLOYMENT, + *extra_deployments, ) error: Final = _unsupported("acompletion", model, router) @@ -1989,7 +1998,7 @@ def test_endpoint_unsupported_rejects_chat_sent_to_the_deployment_id_or_backend_ "extra_deployments, router_kwargs", [ ((), {"router_general_settings": RouterGeneralSettings(pass_through_all_models=True)}), - (({"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": "test-key"}},), {}), + (({"model_name": "sagemaker/*", "litellm_params": {"model": "openai/*", "api_key": "test-key"}},), {}), ], ) def test_endpoint_unsupported_leaves_a_backend_model_to_the_route_that_serves_it(extra_deployments, router_kwargs): @@ -2037,6 +2046,22 @@ EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment( ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}), ({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}), ({}, {"fallbacks": ["chat-fallback-id"]}), + ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), + ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), + ({"fallbacks": [{"embed-only": ["missing-group"]}, {"missing-group": ["chat-fallback"]}]}, {}), + ({}, {"fallbacks": ["openai/gpt-4o-mini"]}), + ( + {"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "openai/gpt-4o-mini"}}, + {}, + ), + ( + {"fallbacks": [{"embed-only": ["chat-fallback"]}, {"embed-fallback": ["embed-only"]}]}, + {"metadata": {"model_group": "embed-fallback"}}, + ), + ( + {"fallbacks": [{"embed-only": ["chat-fallback"]}, {"embed-fallback": ["embed-only"]}]}, + {"litellm_metadata": {"pre_routing_selected_model": "embed-fallback"}}, + ), ], ) def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): @@ -2063,9 +2088,7 @@ def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, r ({}, {"fallbacks": [], "router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), ({"fallbacks": [{"embed-only": ["embed-fallback"]}]}, {}), ({"fallbacks": [{"embed-only": ["embed-only"]}]}, {}), - ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), ({"default_fallbacks": ["embed-fallback"]}, {}), - ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), ({}, {"fallbacks": ["embed-fallback-id"]}), ({}, {"fallbacks": [{"model": "embed-fallback", "temperature": 0.2}]}), ({"fallbacks": [{"embed-only": ["embed-alias"]}], "model_group_alias": {"embed-alias": "embed-fallback"}}, {}), From d010fd16a1482d6f9fb9b6181224e34104202fd7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 04:18:30 -0700 Subject: [PATCH 11/21] fix(proxy): ignore a client model_group that is not a string when building rate-limit headers A client metadata.model_group sent as an int or a list crashed the error handler with a 500 once the embedding-only check rejected the call before routing. The header helper now treats a non-string stamp as no model group. Drops the pattern clause in the backend-name lookup, since the group lookup already returns wildcard matches. Adds the third audit round's routing-stamp cells. --- litellm/proxy/common_utils/callback_utils.py | 5 +- litellm/proxy/route_llm_request.py | 6 +- .../test_non_chat_model_endpoint_gate.py | 271 ++++++++++++++++-- .../proxy/common_utils/test_callback_utils.py | 13 + 4 files changed, 269 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 4fa7709e5c7..b95633f25fe 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -429,10 +429,7 @@ def get_model_group_from_litellm_kwargs(kwargs: dict) -> str | None: def get_model_group_from_request_data(data: dict) -> str | None: _metadata: Final = data.get("metadata", None) or {} _model_group: Final = _metadata.get("model_group", None) - if _model_group is not None: - return _model_group - - return None + return _model_group if isinstance(_model_group, str) else None def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, str]: diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index f45171ec34e..9d9cf11002d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -465,11 +465,7 @@ def _group_deployments(model: str, llm_router: LitellmRouter, team_id: str | Non def _deployments_by_backend_name(model: str, llm_router: LitellmRouter) -> tuple[Mapping[str, object], ...]: - if ( - model not in llm_router.deployment_names - or llm_router.router_general_settings.pass_through_all_models - or llm_router.pattern_router.get_deployments_by_pattern(model=model) - ): + if model not in llm_router.deployment_names or llm_router.router_general_settings.pass_through_all_models: return () return tuple( deployment diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index fdef4f579f2..3adc8ba5939 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -39,6 +39,7 @@ from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, _PROVIDER_KEY: Final = "integration-provider-key" _EMBEDDING_MODEL: Final = "text-embedding-3-small" _CHAT_MODEL: Final = "gpt-4o-mini" +_CHAT_BACKEND_PREFIX: Final = "gate-chat-backend-" _SERVED_TEXT: Final = "served by the chat upstream" _VECTOR: Final = (0.25, 0.5, 0.75) _COUNTED_TOKENS: Final = 17 @@ -150,6 +151,10 @@ def _text_completion_reply() -> Reply: ) +def _serves_chat(model: JsonValue) -> bool: + return model == _CHAT_MODEL or (isinstance(model, str) and model.startswith(_CHAT_BACKEND_PREFIX)) + + def _respond(request: Request) -> Reply: path: Final = request.target.split("?")[0] body: Final = JSON_OBJECT.validate_json(request.body or b"{}") @@ -158,7 +163,7 @@ def _respond(request: Request) -> Reply: return _embedding_reply() if path == "/v1/responses/input_tokens": return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS}).encode()) - if body.get("model") != _CHAT_MODEL: + if not _serves_chat(body.get("model")): return _not_a_chat_model() if path == "/v1/chat/completions": return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) @@ -870,23 +875,50 @@ def test_malformed_request_fallbacks_step_aside_for_the_router(rig: _Rig, shape: assert rig.upstream.targets(marker)[:1] == ("/v1/chat/completions",), response.text +def _assert_answered_after_one_miss( + upstream: _Upstream, response: httpx.Response, marker: str, target: str = "/v1/chat/completions" +) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert upstream.targets(marker) == (target, target), response.text + + +def _assert_the_provider_answered(rig: _Rig, response: httpx.Response, marker: str) -> None: + assert response.status_code == 404, response.text + assert "scripted: not a chat model" in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",), response.text + + def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: ghost: Final = f"gate-ghost-{uuid.uuid4().hex}" return {"fallbacks": [{rig.names["embedding_only"]: [ghost]}, {ghost: [rig.names["chat_target"]]}]} -def _with_a_client_routing_stamp(rig: _Rig) -> Mapping[str, JsonValue]: - model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] - return { - "fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], - "metadata": {"model_group": second}, +_ROUTING_STAMPS: Final[Mapping[str, Callable[[JsonValue], Mapping[str, JsonValue]]]] = MappingProxyType( + { + "a_client_routing_stamp": lambda group: {"metadata": {"model_group": group}}, + "a_pre_routing_stamp": lambda group: {"metadata": {"pre_routing_selected_model": group}}, + "a_litellm_metadata_routing_stamp": lambda group: {"litellm_metadata": {"model_group": group}}, + "a_litellm_metadata_pre_routing_stamp": lambda group: { + "litellm_metadata": {"pre_routing_selected_model": group} + }, + "a_routing_stamp_in_both_buckets": lambda group: { + "metadata": {"model_group": group}, + "litellm_metadata": {"model_group": group}, + }, } +) + + +def _with_a_routing_stamp(stamp: Callable[[JsonValue], Mapping[str, JsonValue]], rig: _Rig) -> Mapping[str, JsonValue]: + model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + return {"fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], **stamp(second)} _FALLBACKS_THE_CHECK_CANNOT_FOLLOW: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( { "a_group_with_no_deployments": _through_a_group_with_no_deployments, - "a_client_routing_stamp": _with_a_client_routing_stamp, + **{name: partial(_with_a_routing_stamp, stamp) for name, stamp in _ROUTING_STAMPS.items()}, } ) @@ -896,8 +928,122 @@ def test_fallbacks_the_check_cannot_follow_step_aside_for_the_router(rig: _Rig, model: Final = rig.names["embedding_only"] marker: Final = _marker() response: Final = _chat(rig, model, marker, **_FALLBACKS_THE_CHECK_CANNOT_FOLLOW[shape](rig)) - assert "Supported endpoints" not in response.text, response.text - assert rig.upstream.targets(marker)[:1] == ("/v1/chat/completions",), response.text + _assert_answered_after_one_miss(rig.upstream, response, marker) + + +_STAMPED_ROUTES: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( + { + "chat_v1": ("metadata", "/v1/chat/completions"), + "messages": ("litellm_metadata", "/v1/responses"), + "responses_v1": ("litellm_metadata", "/v1/responses"), + } +) + + +@pytest.mark.parametrize("route", tuple(_STAMPED_ROUTES)) +def test_a_client_routing_stamp_steps_aside_on_every_chat_shaped_route(rig: _Rig, route: str) -> None: + model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + bucket, target = _STAMPED_ROUTES[route] + marker: Final = _marker() + spec: Final = _ROUTES[route] + response: Final = rig.gateway.request( + "POST", + spec.path, + { + **_route_body(spec, model, marker), + "fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], + bucket: {"model_group": second}, + }, + ) + _assert_answered_after_one_miss(rig.upstream, response, marker, target) + + +_STAMPS_THE_CHECK_IGNORES: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( + { + "an_int_routing_stamp": lambda rig: {"metadata": {"model_group": 7}}, + "a_list_routing_stamp": lambda rig: {"metadata": {"model_group": [rig.names["embedding_only_second"]]}}, + "an_empty_routing_stamp": lambda rig: {"metadata": {"model_group": ""}}, + "a_routing_stamp_naming_the_requested_group": lambda rig: { + "metadata": {"model_group": rig.names["embedding_only"]} + }, + "an_int_pre_routing_stamp": lambda rig: {"litellm_metadata": {"pre_routing_selected_model": 7}}, + "an_int_original_group_stamp": lambda rig: {"metadata": {"original_model_group": 7}}, + "an_original_group_stamp_the_proxy_strips": lambda rig: { + "metadata": {"original_model_group": rig.names["embedding_only_second"]} + }, + } +) + + +@pytest.mark.parametrize("stamp", tuple(_STAMPS_THE_CHECK_IGNORES)) +def test_routing_stamps_that_name_no_other_group_keep_the_rejection(rig: _Rig, stamp: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat( + rig, + model, + marker, + fallbacks=[{rig.names["embedding_only_second"]: [rig.names["chat_target"]]}], + **_STAMPS_THE_CHECK_IGNORES[stamp](rig), + ) + _assert_rejected_chat(rig, response, marker, model) + + +def test_a_routing_stamp_without_any_fallbacks_keeps_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, metadata={"model_group": rig.names["embedding_only_second"]}) + _assert_rejected_chat(rig, response, marker, model) + + +_STAMPS_THE_ROUTER_RESETS: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( + { + "a_routing_stamp_naming_a_group_with_a_chat_chain": lambda rig: { + "metadata": {"model_group": rig.names["embedding_only_second"]} + }, + "a_pre_routing_stamp_naming_a_group_with_a_chat_chain": lambda rig: { + "metadata": {"pre_routing_selected_model": rig.names["embedding_only_second"]} + }, + "a_five_kb_routing_stamp": lambda rig: {"metadata": {"model_group": "g" * 5000}}, + } +) + + +@pytest.mark.parametrize("stamp", tuple(_STAMPS_THE_ROUTER_RESETS)) +def test_routing_stamps_the_router_resets_step_aside_and_the_provider_answers(rig: _Rig, stamp: str) -> None: + marker: Final = _marker() + response: Final = _chat( + rig, + rig.names["embedding_only"], + marker, + fallbacks=[{rig.names["embedding_only_second"]: [rig.names["chat_target"]]}], + **_STAMPS_THE_ROUTER_RESETS[stamp](rig), + ) + _assert_the_provider_answered(rig, response, marker) + + +def test_request_fallbacks_naming_a_chat_backend_name_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + backend: Final = f"openai/{_CHAT_BACKEND_PREFIX}{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"gate-backend-{uuid.uuid4().hex[:12]}", {"model": backend, "api_base": rig.api_base}, {} + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = _chat(rig, model, marker, fallbacks=[{model: [backend]}]) + _assert_answered_after_one_miss(rig.upstream, response, marker) + + +@pytest.mark.parametrize("level", ("key", "team")) +def test_router_settings_fallbacks_through_a_group_with_no_deployments_answer(rig: _Rig, level: str) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + settings: Final = {"router_settings": dict(_through_a_group_with_no_deployments(rig))} + key: Final = scenario.key(**settings) if level == "key" else scenario.key(team_id=scenario.team(**settings)) + response: Final = _chat(rig, rig.names["embedding_only"], marker, key=key) + _assert_answered_after_one_miss(rig.upstream, response, marker) def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: @@ -1501,8 +1647,25 @@ class _Addressed: api_base: str -def _addressed_config(directory: Path, api_base: str, names: Mapping[str, str], suffix: str) -> Path: +def _wildcard_backend_entry(name: str, api_base: str, backend: str) -> dict[str, JsonValue]: + entry: Final = _routed_entry(name, api_base, backend, _EMBED_ONLY) + return {**entry, "litellm_params": {**object_value(entry["litellm_params"]), "custom_llm_provider": "openai"}} + + +def _addressed_config( + directory: Path, api_base: str, names: Mapping[str, str], suffix: str, with_wildcards: bool +) -> Path: shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + wildcard: Final = f"gate-addr-wild-{suffix}" + wildcard_entries: Final = ( + ( + _routed_entry(f"{wildcard}/*", api_base, "openai/*", {}), + _wildcard_backend_entry(names["wildcard_backend"], api_base, f"{wildcard}/{_CHAT_MODEL}"), + ) + if with_wildcards + else () + ) + aliases: Final = {names["wildcard_alias"]: f"{wildcard}/{_CHAT_BACKEND_PREFIX}{suffix}"} if with_wildcards else {} config: Final = { **shared, "model_list": [ @@ -1511,8 +1674,15 @@ def _addressed_config(directory: Path, api_base: str, names: Mapping[str, str], ), _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat", "id": names["chat_id"]}), _routed_entry(names["shared_embed"], api_base, f"openai/{_CHAT_MODEL}", _EMBED_ONLY), + *wildcard_entries, + _routed_entry(names["ghost_primary"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), ], - "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "model_group_alias": aliases, + "fallbacks": [{names["ghost_primary"]: [names["ghost"]]}, {names["ghost"]: [names["chat"]]}], + }, "general_settings": { **object_value(shared["general_settings"]), "proxy_config_reload_interval_seconds": 3600, @@ -1523,12 +1693,25 @@ def _addressed_config(directory: Path, api_base: str, names: Mapping[str, str], return path -@pytest.fixture(scope="module") -def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: +@contextlib.contextmanager +def _addressed_proxy(tmp_path_factory: pytest.TempPathFactory, with_wildcards: bool) -> Iterator[_Addressed]: directory: Final = tmp_path_factory.mktemp("gate-addressed") suffix: Final = uuid.uuid4().hex[:12] names: Final = MappingProxyType( - {label: f"gate-addr-{label}-{suffix}" for label in ("embed", "embed_id", "chat", "chat_id", "shared_embed")} + { + label: f"gate-addr-{label}-{suffix}" + for label in ( + "embed", + "embed_id", + "chat", + "chat_id", + "shared_embed", + "wildcard_backend", + "wildcard_alias", + "ghost_primary", + "ghost", + ) + } ) with ( wire_server(answering_model_discovery(_respond)) as wire, @@ -1537,7 +1720,7 @@ def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: rig_gateway, directory, {}, - config=_addressed_config(directory, wire.url + "/v1", names, suffix), + config=_addressed_config(directory, wire.url + "/v1", names, suffix, with_wildcards), workers=2, ) as owned, ): @@ -1552,12 +1735,28 @@ def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: _Upstream(wire), names, MappingProxyType( - {"deployment_id": names["embed_id"], "backend_name": f"openai/gate-addr-backend-{suffix}"} + { + "deployment_id": names["embed_id"], + "backend_name": f"openai/gate-addr-backend-{suffix}", + "wildcard_backend_name": f"gate-addr-wild-{suffix}/{_CHAT_MODEL}", + } ), wire.url + "/v1", ) +@pytest.fixture(scope="module") +def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=True) as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def addressed_without_wildcards(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=False) as proxy: + yield proxy + + _ADDRESSED_ROUTES: Final = ("chat_v1", "messages", "responses_v1") @@ -1622,9 +1821,47 @@ def test_request_fallbacks_naming_a_chat_deployment_id_answer_for_an_embedding_o @_ROUTED_TIMEOUT -def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rejected_on_every_call( +def test_chat_sent_to_an_embedding_only_backend_name_a_wildcard_also_matches_is_served_by_the_wildcard( addressed: _Addressed, ) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.addresses["wildcard_backend_name"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in addressed.upstream.calls(marker)] == [_CHAT_MODEL] + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_an_alias_only_a_wildcard_serves_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["wildcard_alias"]]}, + ) + _assert_answered_after_one_miss(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_yaml_fallbacks_through_a_group_with_no_deployments_answer_for_an_embedding_only_group( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.names["ghost_primary"], marker) + ) + _assert_answered_after_one_miss(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rejected_on_every_call( + addressed_without_wildcards: _Addressed, +) -> None: + addressed: Final = addressed_without_wildcards name: Final = f"gate-addr-loaded-{uuid.uuid4().hex[:12]}" markers: Final = tuple(_marker() for _ in range(6)) with addressed.rig.scenario() as scenario: diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index ecb2375d495..c8d909ac2eb 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -49,6 +49,19 @@ def test_get_remaining_tokens_and_requests_from_request_data(): } +@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}]) +def test_remaining_limit_headers_ignore_a_model_group_that_is_not_a_string(client_model_group: object): + casedata: Final = { + "metadata": { + "model_group": client_model_group, + f"litellm-key-remaining-requests-{client_model_group}": 100, + f"litellm-key-remaining-tokens-{client_model_group}": 200, + } + } + + assert get_remaining_tokens_and_requests_from_request_data(casedata) == {} + + @patch( "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", return_value=["API_KEY", "MISSING_VAR"], From cafe24f0c8a88b4ceb61af0025f6f7bd36b6245a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 04:41:32 -0700 Subject: [PATCH 12/21] fix(proxy): keep a non-string client model_group out of spend rows --- .../spend_tracking/spend_tracking_utils.py | 8 +++++-- .../test_spend_tracking_utils.py | 22 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 99bd2fbc480..4446dcff066 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -401,7 +401,11 @@ def _sl_attribution_fallback( ) -> str: if standard_logging_payload is None: return "" - return standard_logging_payload.get(field) or "" + return _string_or_empty(standard_logging_payload.get(field)) + + +def _string_or_empty(value: object) -> str: + return value if isinstance(value, str) else "" def _deployment_provider(deployment: DeploymentTypedDict) -> str | None: @@ -600,7 +604,7 @@ def get_logging_payload( _model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback( standard_logging_payload, "model_id" ) - _model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback( + _model_group: Final = _string_or_empty(metadata.get("model_group")) or _sl_attribution_fallback( standard_logging_payload, "model_group" ) _api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback( diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 4d299f05fbe..94fa6c3767e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4900,6 +4900,28 @@ def test_get_logging_payload_router_rejected_request_takes_provider_from_model_g assert payload["custom_llm_provider"] == expected_provider +@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}]) +def test_get_logging_payload_drops_a_client_model_group_that_is_not_a_string(client_model_group: object): + payload: Final = get_logging_payload( + kwargs={ + "model": "embed-only", + "litellm_params": { + "metadata": {"user_api_key": "test-key", "model_group": client_model_group, "status": "failure"} + }, + "standard_logging_object": { + **_make_failed_request_standard_logging_payload(), + "model_group": client_model_group, + }, + }, + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + llm_router=None, + ) + + assert payload["model_group"] == "" + + def test_get_logging_payload_router_rejected_request_without_router_leaves_provider_empty(): assert _router_rejected_failure_payload("openai-group", None)["custom_llm_provider"] == "" From 4307b2b75bef0a80e64f36246a52e2cb8d422190 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:52:45 -0700 Subject: [PATCH 13/21] fix(proxy): validate standard fallbacks before the typed fallback-group lookup --- litellm/proxy/route_llm_request.py | 14 ++++++++------ tests/unit/proxy/test_route_llm_request.py | 1 + 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 9d9cf11002d..833f73f3d69 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -405,6 +405,7 @@ class ProxyUnsupportedEndpointError(ProxyException): _FALLBACK_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter( # mutable-ok: router fallbacks are lists list[object], config=ConfigDict(strict=True) ) +_STANDARD_FALLBACKS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str]], ...]) class _RouterFallbacks(Protocol): @@ -510,7 +511,7 @@ def _fallback_targets( if check_non_standard_fallback_format(fallbacks=fallbacks): return fallbacks resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, + fallbacks=_STANDARD_FALLBACKS_ADAPTER.validate_python(fallbacks), lookup_groups=fallback_lookup_groups(data, model), ) return resolved or () @@ -591,13 +592,14 @@ def _fallback_would_answer( raw_fallbacks: Final = _fallbacks_for_request(data, llm_router) if raw_fallbacks is None: return False - try: - fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks) - except ValidationError: - return True if fallback_lookup_groups(data, model) != (model,): return True - return _fallbacks_reach_chat(model, data, fallbacks, llm_router, team_id) + try: + return _fallbacks_reach_chat( + model, data, _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks), llm_router, team_id + ) + except ValidationError: + return True def endpoint_unsupported_by_model( diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index b7189f3990d..626de363e33 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -2045,6 +2045,7 @@ EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment( ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}), ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}), ({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}), + ({}, {"fallbacks": [{"embed-only": [{"model": "embed-fallback"}]}]}), ({}, {"fallbacks": ["chat-fallback-id"]}), ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), From e08adbed6158583254f8e07b8faf5797c2295f3e Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 17:43:23 -0700 Subject: [PATCH 14/21] feat(router): opt in to non-chat deployment filtering --- litellm/proxy/common_request_processing.py | 18 +- .../proxy/response_api_endpoints/endpoints.py | 27 +- litellm/proxy/route_llm_request.py | 289 +------ litellm/router.py | 34 +- litellm/router_utils/endpoint_eligibility.py | 95 +++ litellm/types/router.py | 1 + .../llm_nonconversational.yaml | 2 +- ..._embedding_only_supported_endpoints_e2e.py | 116 +-- tests/e2e/models.py | 1 + .../test_non_chat_model_endpoint_gate.py | 742 ++++-------------- .../response_api_endpoints/test_endpoints.py | 92 --- .../proxy/test_common_request_processing.py | 120 +-- tests/unit/proxy/test_route_llm_request.py | 450 ----------- .../router_utils/test_endpoint_eligibility.py | 76 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 15 files changed, 386 insertions(+), 1682 deletions(-) create mode 100644 litellm/router_utils/endpoint_eligibility.py create mode 100644 tests/unit/router_utils/test_endpoint_eligibility.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 0766299d91d..95d1123bfc9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -119,7 +119,6 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( - ProxyUnsupportedEndpointError, endpoint_unsupported_by_model, route_request, ) @@ -132,6 +131,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model +from litellm.router_utils.endpoint_eligibility import UnsupportedEndpointError from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError from litellm.types.router_weights import validate_router_weights @@ -2301,15 +2301,6 @@ class ProxyBaseLLMRequestProcessing: self.data["litellm_logging_obj"] = logging_obj - unsupported_endpoint: Final = endpoint_unsupported_by_model( - route_type=route_type, - data=self.data, - llm_router=llm_router, - team_id=user_api_key_dict.team_id, - ) - if unsupported_endpoint is not None: - raise unsupported_endpoint - # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set # via litellm_params.guardrails are only honored on post_call paths @@ -2479,6 +2470,11 @@ class ProxyBaseLLMRequestProcessing: self.data = independent_snapshot(pristine) self.data["model"] = fallback_model try: + if ( + endpoint_unsupported_by_model(route_type, self.data, llm_router, user_api_key_dict.team_id) + is not None + ): + continue return await self.common_processing_pre_call_logic( request=request, general_settings=general_settings, @@ -2496,7 +2492,7 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, rate_limited_model=original_model, ) - except ProxyUnsupportedEndpointError: + except UnsupportedEndpointError: continue except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 937a51faf07..ad82a1ed724 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -56,7 +56,8 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, ) -from litellm.proxy.route_llm_request import endpoint_unsupported_by_model, raise_if_required_body_param_missing +from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.router_utils.endpoint_eligibility import UnsupportedEndpointError from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( REASONING_EFFORT, @@ -1731,6 +1732,13 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec def _responses_ws_failure_frame(failure: Exception) -> str: + if isinstance(failure, UnsupportedEndpointError): + return json.dumps( + { + "type": "error", + "error": {"type": "invalid_request_error", "message": failure.message.split("\n\n", 1)[0]}, + } + ) raw_status: Final = getattr(failure, "status_code", None) status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500 error_type: Final = ( @@ -1857,20 +1865,8 @@ async def _responses_websocket_session( user_api_base=user_api_base, model=resolved_model, route_type="_aresponses_websocket", - ) - unsupported_endpoint: Final = endpoint_unsupported_by_model( - route_type="_aresponses_websocket", - data=data, llm_router=llm_router, - team_id=user_api_key_dict.team_id, ) - if unsupported_endpoint is not None: - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=unsupported_endpoint, - request_data=data, - ) - raise unsupported_endpoint except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") try: @@ -1917,7 +1913,10 @@ async def _responses_websocket_session( original_exception=e, request_data=routed_data, ) - await websocket.close(code=1011, reason="Internal server error") + await websocket.close( + code=1008 if isinstance(e, UnsupportedEndpointError) else 1011, + reason="Pre-call error" if isinstance(e, UnsupportedEndpointError) else "Internal server error", + ) finally: await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 833f73f3d69..b1d0c72f97c 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,25 +1,17 @@ import asyncio -from collections.abc import Mapping, Sequence +from collections.abc import Mapping from dataclasses import dataclass -from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal import httpx from fastapi import HTTPException, status -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import TypeAdapter, ValidationError import litellm -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request -from litellm.router_utils.fallback_event_handlers import ( - check_non_standard_fallback_format, - fallback_lookup_groups, - get_fallback_model_group_for_lookup_groups, - has_unattempted_fallback_target, -) -from litellm.types.utils import COST_MAP_LOOKUP_KEY, echoed_cost_map_fields +from litellm.router_utils.endpoint_eligibility import UnsupportedEndpointError # Client-supplied params that make the router or the call path fabricate a # failure or a delay instead of calling the provider. The ``mock_testing_*`` @@ -365,267 +357,24 @@ def raise_if_required_body_param_missing( ) -CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE: Final[Mapping[str, str]] = MappingProxyType( - { - "acompletion": "/v1/chat/completions", - "atext_completion": "/v1/completions", - "anthropic_messages": "/v1/messages", - "aresponses": "/v1/responses", - "agenerate_content": "generateContent", - "agenerate_content_stream": "streamGenerateContent", - "acreate_interaction": "/v1beta/interactions", - "_aresponses_websocket": "/v1/responses", - "acompact_responses": "/v1/responses/compact", - } -) - -CHAT_SHAPED_ENDPOINTS: Final[frozenset[str]] = frozenset( - { - "/chat/completions", - "/completions", - "/responses", - *CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE.values(), - } -) - -CHAT_SHAPED_MODES: Final[frozenset[str]] = frozenset({"chat", "completion", "responses"}) - - -class _DeploymentEndpointDeclaration(BaseModel): - model_config = ConfigDict(frozen=True) - - mode: str | None = None - supported_endpoints: tuple[str, ...] | None = None - - -class ProxyUnsupportedEndpointError(ProxyException): - pass - - -_FALLBACK_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter( # mutable-ok: router fallbacks are lists - list[object], config=ConfigDict(strict=True) -) -_STANDARD_FALLBACKS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str]], ...]) - - -class _RouterFallbacks(Protocol): - @property - def fallbacks(self) -> Sequence[object] | None: ... - - -def _loaded_catalog_entry(model_info: Mapping[str, object]) -> Mapping[str, object]: - catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) - if not isinstance(catalog_key, str): - return MappingProxyType({}) - return GetModelCostMap.loaded_model_cost_map().get(catalog_key) or MappingProxyType({}) - - -def _catalog_serves_chat(catalog_entry: Mapping[str, object]) -> bool: - try: - declared: Final = _DeploymentEndpointDeclaration.model_validate(dict(catalog_entry)) - except ValidationError: - return False - return (declared.mode or "").strip().lower() in CHAT_SHAPED_MODES or not CHAT_SHAPED_ENDPOINTS.isdisjoint( - declared.supported_endpoints or () - ) - - -def _declaration_overruled_by_catalog(model_info: Mapping[str, object]) -> bool: - catalog_entry: Final = _loaded_catalog_entry(model_info) - return _catalog_serves_chat(catalog_entry) or "supported_endpoints" in echoed_cost_map_fields( - model_info, catalog_entry - ) - - -def _declared_non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None: - try: - model_info: Final = JSON_OBJECT_ADAPTER.validate_python(deployment.get("model_info")) - declared: Final = _DeploymentEndpointDeclaration.model_validate(model_info) - except ValidationError: - return None - mode: Final = (declared.mode or "").strip().lower() - if not mode or mode in CHAT_SHAPED_MODES or declared.supported_endpoints is None: - return None - endpoints: Final = frozenset(declared.supported_endpoints) - if not endpoints.isdisjoint(CHAT_SHAPED_ENDPOINTS): - return None - return None if _declaration_overruled_by_catalog(model_info) else endpoints - - -def _deployment_by_id(model_id: str, llm_router: LitellmRouter) -> tuple[Mapping[str, object], ...]: - try: - return (JSON_OBJECT_ADAPTER.validate_python(llm_router.get_model_info(id=model_id)),) - except ValidationError: - return () - - -def _group_deployments(model: str, llm_router: LitellmRouter, team_id: str | None) -> tuple[Mapping[str, object], ...]: - return tuple(llm_router.get_model_list(model_name=model, team_id=team_id) or ()) or _deployment_by_id( - model, llm_router - ) - - -def _deployments_by_backend_name(model: str, llm_router: LitellmRouter) -> tuple[Mapping[str, object], ...]: - if model not in llm_router.deployment_names or llm_router.router_general_settings.pass_through_all_models: - return () - return tuple( - deployment - for deployment in llm_router.get_model_list() or () - if deployment["litellm_params"].get("model") == model - ) - - -def _requested_deployments( - model: str, llm_router: LitellmRouter, team_id: str | None -) -> tuple[Mapping[str, object], ...]: - return _group_deployments(model, llm_router, team_id) or _deployments_by_backend_name(model, llm_router) - - -def _any_serves_chat(deployments: Sequence[Mapping[str, object]]) -> bool: - return any(_declared_non_chat_endpoints(deployment) is None for deployment in deployments) - - -def _override_fallbacks(data: Mapping[str, object]) -> object: - try: - return JSON_OBJECT_ADAPTER.validate_python(data.get("router_settings_override")).get("fallbacks") - except ValidationError: - return None - - -def _fallbacks_for_request(data: Mapping[str, object], llm_router: LitellmRouter) -> object: - if "fallbacks" in data: - return data["fallbacks"] - override_fallbacks: Final = _override_fallbacks(data) - return override_fallbacks if override_fallbacks is not None else _router_fallbacks(llm_router) - - -def _router_fallbacks(llm_router: _RouterFallbacks) -> Sequence[object] | None: - return llm_router.fallbacks - - -def _fallback_targets( - model: str, - data: Mapping[str, object], - fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract -) -> Sequence[object]: - if check_non_standard_fallback_format(fallbacks=fallbacks): - return fallbacks - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=_STANDARD_FALLBACKS_ADAPTER.validate_python(fallbacks), - lookup_groups=fallback_lookup_groups(data, model), - ) - return resolved or () - - -def _fallback_target_group(target: object) -> str | None: - if isinstance(target, str): - return target - try: - entry: Final = JSON_OBJECT_ADAPTER.validate_python(target) - except ValidationError: - return None - named_model: Final = entry.get("model") - return named_model if isinstance(named_model, str) else None - - -def _next_fallback_groups( - model: str, - data: Mapping[str, object], - fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract -) -> tuple[str, ...]: - return tuple( - group - for target in _fallback_targets(model, data, fallbacks) - if has_unattempted_fallback_target((target,), data) and (group := _fallback_target_group(target)) is not None - ) - - -@dataclass(frozen=True, slots=True) -class _FallbackWalk: - frontier: tuple[str, ...] - seen: frozenset[str] - reaches_chat: bool = False - - -def _walk_one_layer( - walk: _FallbackWalk, - data: Mapping[str, object], - fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract - llm_router: LitellmRouter, - team_id: str | None, -) -> _FallbackWalk: - reached: Final = dict.fromkeys( - chain.from_iterable(_next_fallback_groups(m, data, fallbacks) for m in walk.frontier) - ) - fresh: Final = tuple(group for group in reached if group not in walk.seen) - deployments_by_group: Final = tuple((group, _group_deployments(group, llm_router, team_id)) for group in fresh) - return _FallbackWalk( - frontier=tuple(group for group, deployments in deployments_by_group if deployments), - seen=walk.seen | frozenset(fresh), - reaches_chat=any(not deployments or _any_serves_chat(deployments) for _, deployments in deployments_by_group), - ) - - -def _fallbacks_reach_chat( - model: str, - data: Mapping[str, object], - fallbacks: list[object], # mutable-ok: mirrors get_fallback_model_group's contract - llm_router: LitellmRouter, - team_id: str | None, -) -> bool: - walks: Final = accumulate( - repeat(None), - lambda walk, _: _walk_one_layer(walk, data, fallbacks, llm_router, team_id), - initial=_FallbackWalk(frontier=(model,), seen=frozenset({model})), - ) - return any(walk.reaches_chat for walk in takewhile(lambda walk: walk.frontier or walk.reaches_chat, walks)) - - -def _fallback_would_answer( - model: str, - data: Mapping[str, object], - llm_router: LitellmRouter, - team_id: str | None, -) -> bool: - if data.get("disable_fallbacks") is True: - return False - raw_fallbacks: Final = _fallbacks_for_request(data, llm_router) - if raw_fallbacks is None: - return False - if fallback_lookup_groups(data, model) != (model,): - return True - try: - return _fallbacks_reach_chat( - model, data, _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks), llm_router, team_id - ) - except ValidationError: - return True - - def endpoint_unsupported_by_model( route_type: str, data: Mapping[str, object], llm_router: LitellmRouter | None, team_id: str | None, -) -> ProxyUnsupportedEndpointError | None: - endpoint: Final = CHAT_SHAPED_ENDPOINT_BY_ROUTE_TYPE.get(route_type) +) -> UnsupportedEndpointError | None: model: Final = data.get("model") - if endpoint is None or not isinstance(model, str) or not model or llm_router is None: + if not isinstance(model, str) or llm_router is None: return None - deployments: Final = _requested_deployments(model, llm_router, team_id) - if not deployments or _any_serves_chat(deployments): - return None - if _fallback_would_answer(model, data, llm_router, team_id): - return None - supported: Final = ", ".join( - sorted(frozenset[str]().union(*(_declared_non_chat_endpoints(d) or frozenset() for d in deployments))) - ) - return ProxyUnsupportedEndpointError( - message=f"{endpoint} is not supported by model={model}. Supported endpoints: {supported or 'none'}", - type="invalid_request_error", - param="model", - code=status.HTTP_400_BAD_REQUEST, - ) + metadata: Final = data.get("metadata") + request_kwargs: Final = { + **data, + "metadata": { + **(metadata if isinstance(metadata, Mapping) else {}), + "user_api_key_team_id": team_id, + }, + } + return llm_router.endpoint_error_for_request(model, route_type, request_kwargs) class MockTestingParamsDisabledError(HTTPException): @@ -901,14 +650,6 @@ async def route_request( if not await model_registry_read_through.attempt(requested_model): raise - unsupported_endpoint: Final = endpoint_unsupported_by_model( - route_type=route_type, - data=JSON_OBJECT_ADAPTER.validate_python(data), - llm_router=proxy_server.llm_router, - team_id=user_api_key_dict.team_id if user_api_key_dict is not None else None, - ) - if unsupported_endpoint is not None: - raise unsupported_endpoint return await _route_request_single_attempt( data=data, llm_router=proxy_server.llm_router, diff --git a/litellm/router.py b/litellm/router.py index f6b60c8953f..d73ec1b3180 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -222,6 +222,11 @@ from litellm.router_utils.cooldown_handlers import ( is_caller_timeout_408, set_cooldown_deployments, ) +from litellm.router_utils.endpoint_eligibility import ( + UnsupportedEndpointError, + eligible_deployments, + endpoint_selection, +) from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, @@ -8391,9 +8396,17 @@ class Router: Handler for making a call to the .completion()/.embeddings()/etc. functions. """ model_group: Final = kwargs.get("model") - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response + endpoint_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python( + kwargs.get("original_generic_function") or original_function + ) + with endpoint_selection( + getattr(endpoint_function, "__name__", None), + getattr(endpoint_function, "__module__", None), + model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), + ): + response = original_function(*args, **kwargs) + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + response = await response await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) @@ -13459,6 +13472,16 @@ class Router: deployment_model: Final = litellm_params.get("model") return isinstance(deployment_model, str) and classify_strategy_router_model(deployment_model) is not None + def endpoint_error_for_request( + self, model: str, call_type: str, request_kwargs: Mapping[str, object] + ) -> UnsupportedEndpointError | None: + try: + with endpoint_selection(call_type, model=model): + self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + except UnsupportedEndpointError as error: + return error + return None + def _common_checks_available_deployment( self, model: str, @@ -13643,7 +13666,10 @@ class Router: model=model, llm_provider="", ) - return selectable + eligible: Final = eligible_deployments(model, selectable) + if isinstance(eligible, UnsupportedEndpointError): + raise eligible + return list(eligible) def _filter_reserved_deployments( self, diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py new file mode 100644 index 00000000000..85db045d6a5 --- /dev/null +++ b/litellm/router_utils/endpoint_eligibility.py @@ -0,0 +1,95 @@ +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from types import MappingProxyType +from typing import Final, TypeVar + +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.exceptions import BadRequestError + +CHAT_ENDPOINT_BY_CALL: Final[Mapping[str, str]] = MappingProxyType( + { + "completion": "/v1/chat/completions", + "acompletion": "/v1/chat/completions", + "_completion": "/v1/chat/completions", + "_acompletion": "/v1/chat/completions", + "text_completion": "/v1/completions", + "atext_completion": "/v1/completions", + "_atext_completion": "/v1/completions", + "anthropic_messages": "/v1/messages", + "aresponses": "/v1/responses", + "agenerate_content": "generateContent", + "agenerate_content_stream": "streamGenerateContent", + "acreate_interaction": "/v1beta/interactions", + "litellm.interactions.main.acreate": "/v1beta/interactions", + "litellm.interactions.main.create": "/v1beta/interactions", + "_aresponses_websocket": "/v1/responses", + "acompact_responses": "/v1/responses/compact", + } +) +CHAT_ENDPOINTS: Final = frozenset((*CHAT_ENDPOINT_BY_CALL.values(), "/chat/completions", "/completions", "/responses")) +_Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object]) +_REQUEST_SELECTION: Final[ContextVar[tuple[str | None, str | None]]] = ContextVar( + "router_request_endpoint", default=(None, None) +) + + +class _EndpointDeclaration(BaseModel): + model_config = ConfigDict(frozen=True) + + enforce_supported_endpoints: bool = False + supported_endpoints: tuple[str, ...] | None = None + + +class UnsupportedEndpointError(BadRequestError): + def __init__(self, message: str, model: str, llm_provider: str = "") -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # inherited public exception constructor is untyped + message=message, model=model, llm_provider=llm_provider + ) + self.message = message + self.type = "invalid_request_error" + self.param = "model" + self.code = "400" + + +@contextmanager +def endpoint_selection(call_type: object, module: object = None, model: object = None) -> Generator[None]: + qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type + endpoint: Final = CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None + token: Final = _REQUEST_SELECTION.set((endpoint, model if isinstance(model, str) else None)) + try: + yield + finally: + _REQUEST_SELECTION.reset(token) + + +def non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None: + try: + declaration: Final = _EndpointDeclaration.model_validate(deployment.get("model_info")) + except ValidationError: + return None + if not declaration.enforce_supported_endpoints or declaration.supported_endpoints is None: + return None + endpoints: Final = frozenset(declaration.supported_endpoints) + return endpoints if endpoints.isdisjoint(CHAT_ENDPOINTS) else None + + +def eligible_deployments( + model: str, deployments: Sequence[_Deployment] +) -> tuple[_Deployment, ...] | UnsupportedEndpointError: + endpoint, requested_model = _REQUEST_SELECTION.get() + if endpoint is None or not deployments: + return tuple(deployments) + eligible: Final = tuple(deployment for deployment in deployments if non_chat_endpoints(deployment) is None) + if eligible: + return eligible + requested: Final = requested_model or model + supported: Final = ", ".join( + sorted(frozenset[str]().union(*(non_chat_endpoints(deployment) or frozenset() for deployment in deployments))) + ) + return UnsupportedEndpointError( + message=f"{endpoint} is not supported by model={requested}. Supported endpoints: {supported or 'none'}", + model=model, + llm_provider="", + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 0ca39e2400d..6ae7a77d571 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -255,6 +255,7 @@ def _whole_float_as_int(value: object) -> object: class ModelInfo(MirroredPricingParams): + enforce_supported_endpoints: bool = False id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: datetime.datetime | None = None diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index bbd03b445f1..a88ccf06c44 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -2,7 +2,7 @@ - {id: llm.completions.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: completions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_completions_endpoint_e2e.py", rationale: "Legacy text /completions endpoint, second-highest production request volume"} - {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"} - {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy ยง9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"} -- {id: llm.embeddings.openai.supported_endpoints_gate.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], fail_before_fix: proven, source: "test_embedding_only_supported_endpoints_e2e.py", rationale: "A model whose every deployment lists only non-chat supported_endpoints answers each chat-shaped route 400 at the proxy, naming the endpoint and the declared list, while /embeddings keeps serving"} +- {id: llm.embeddings.openai.supported_endpoints_gate.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], fail_before_fix: proven, source: "test_embedding_only_supported_endpoints_e2e.py", rationale: "Explicit endpoint restrictions preserve real embeddings and route chat through a live fallback; rejection and zero-provider-call contracts live in integration/routing/test_non_chat_model_endpoint_gate.py"} - {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"} - {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (GitHub issue #36646)"} - {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"} diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py index 9a89344cebd..1f3ba3d05db 100644 --- a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -1,27 +1,17 @@ -"""Live e2e: an embedding model whose deployments declare supported_endpoints refuses chat-shaped calls. - -Each test registers an OpenAI embedding deployment whose model_info.supported_endpoints lists only -/v1/embeddings, or a control deployment on the same backend that declares nothing. /embeddings serves -real vectors on both. Every chat-shaped route to the declared model answers 400 with the proxy's own -message, while the same chat call to the control goes to OpenAI and comes back with whatever OpenAI -answers, so getting the proxy's message is the proof that no provider call was made. -""" +"""Live provider checks for embeddings and chat fallback from an explicitly restricted model.""" from __future__ import annotations import dataclasses from typing import Final -import anthropic -import openai import pytest from e2e_config import unique_marker from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta from lifecycle import ResourceManager -from models import AnthropicErrorBody, AnthropicErrorEvent, LiteLLMParamsBody, ModelInfoBody, ModelMode, ModelNewBody +from models import LiteLLMParamsBody, ModelInfoBody, ModelMode, ModelNewBody from openai import OpenAI from proxy_client import ProxyClient -from pydantic import BaseModel, ConfigDict from sdk_clients import NO_PROXY_CACHE, SdkClients pytestmark = pytest.mark.e2e @@ -30,15 +20,6 @@ EMBEDDING_BACKEND: Final = "openai/text-embedding-3-small" CHAT_FALLBACK_BACKEND: Final = "openai/gpt-5.4-nano" EMBEDDINGS_ONLY: Final = ("/v1/embeddings",) PROMPT: Final = "Say this is a test!" -UNSUPPORTED_MARKER: Final = "is not supported by model=" - - -class _OpenAiErrorDetail(BaseModel): - model_config = ConfigDict(frozen=True) - - message: str - type: str - param: str | None = None def _subject(route: Route) -> Subject: @@ -65,6 +46,7 @@ def _register( model_name=model, litellm_params=LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY"), model_info=ModelInfoBody( + enforce_supported_endpoints=supported_endpoints is not None, mode=mode, supported_endpoints=None if supported_endpoints is None else list(supported_endpoints), ), @@ -78,27 +60,6 @@ def _openai_at(sdk: SdkClients, api_prefix: str, key: str) -> OpenAI: return dataclasses.replace(sdk, base_url=f"{sdk.base_url}{api_prefix}").openai(key) -def _unsupported_message(endpoint: str, model: str) -> str: - return f"{endpoint} {UNSUPPORTED_MARKER}{model}. Supported endpoints: {', '.join(sorted(EMBEDDINGS_ONLY))}" - - -def _chat_error_message(client: OpenAI, model: str) -> str: - try: - _ = client.chat.completions.create( - model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE - ) - except openai.APIStatusError as error: - return f"{error.status_code} {error.message}" - return "" - - -def _assert_openai_rejection(error: openai.APIStatusError, endpoint: str, model: str) -> None: - assert error.status_code == 400, f"expected the proxy's 400, got {error.status_code}: {error.message}" - assert _OpenAiErrorDetail.model_validate(error.body) == _OpenAiErrorDetail( - message=_unsupported_message(endpoint, model), type="invalid_request_error", param="model" - ), f"unexpected error body: {error.body!r}" - - class TestEmbeddingOnlySupportedEndpoints: @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") @meta(_subject(Route.EMBEDDINGS)) @@ -126,65 +87,6 @@ class TestEmbeddingOnlySupportedEndpoints: f"/embeddings returned an empty or all-zero vector for {model}" ) - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.CHAT_COMPLETIONS)) - @pytest.mark.parametrize("api_prefix", ("/v1", ""), ids=("v1-chat-completions", "chat-completions")) - def test_chat_completions_are_rejected_before_any_provider_call( - self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, api_prefix: str - ) -> None: - model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) - with pytest.raises(openai.APIStatusError) as raised: - _ = _openai_at(sdk, api_prefix, resources.key()).chat.completions.create( - model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE - ) - _assert_openai_rejection(raised.value, "/v1/chat/completions", model) - - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.RESPONSES)) - @pytest.mark.parametrize("api_prefix", ("/v1", ""), ids=("v1-responses", "responses")) - def test_responses_are_rejected_before_any_provider_call( - self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, api_prefix: str - ) -> None: - model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) - with pytest.raises(openai.APIStatusError) as raised: - _ = _openai_at(sdk, api_prefix, resources.key()).responses.create( - model=model, input=PROMPT, extra_body=NO_PROXY_CACHE - ) - _assert_openai_rejection(raised.value, "/v1/responses", model) - - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.MESSAGES)) - def test_messages_are_rejected_before_any_provider_call( - self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients - ) -> None: - model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) - with pytest.raises(anthropic.APIStatusError) as raised: - _ = sdk.anthropic(resources.key()).messages.create( - model=model, - max_tokens=16, - messages=[{"role": "user", "content": PROMPT}], - extra_body=NO_PROXY_CACHE, - ) - assert raised.value.status_code == 400, ( - f"expected the proxy's 400, got {raised.value.status_code}: {raised.value.message}" - ) - assert AnthropicErrorEvent.model_validate(raised.value.body) == AnthropicErrorEvent( - type="error", - error=AnthropicErrorBody(type="invalid_request_error", message=_unsupported_message("/v1/messages", model)), - ), f"unexpected Anthropic error envelope: {raised.value.body!r}" - - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.COMPLETIONS)) - def test_text_completions_are_rejected_before_any_provider_call( - self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients - ) -> None: - model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) - with pytest.raises(openai.APIStatusError) as raised: - _ = _openai_at(sdk, "/v1", resources.key()).completions.create( - model=model, prompt=PROMPT, max_tokens=16, extra_body=NO_PROXY_CACHE - ) - _assert_openai_rejection(raised.value, "/v1/completions", model) - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") @meta(_subject(Route.CHAT_COMPLETIONS)) def test_chat_completions_with_a_fallback_are_answered_by_the_fallback( @@ -200,15 +102,3 @@ class TestEmbeddingOnlySupportedEndpoints: assert completion.choices and completion.choices[0].message.content, ( f"the fallback {fallback} should answer a chat call to {model}, got {completion!r}" ) - - @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.CHAT_COMPLETIONS)) - def test_chat_completions_without_supported_endpoints_reach_the_provider( - self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients - ) -> None: - model: Final = _register(proxy, resources, "e2e-embed-control", None) - error_message: Final = _chat_error_message(_openai_at(sdk, "/v1", resources.key()), model) - assert UNSUPPORTED_MARKER not in error_message, ( - f"a model that declares no supported_endpoints must reach the provider, but the proxy " - f"rejected it: {error_message}" - ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 42a0b6efac7..25b48201bef 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1426,6 +1426,7 @@ ModelMode = Literal["chat", "embedding", "batch", "realtime", "image_generation" class ModelInfoBody(BaseModel): + enforce_supported_endpoints: bool | None = None # id is left unset so the proxy assigns a unique model_id per deployment. # Pinning it to the model_name made re-registrations of a fixed-name model # (e.g. the batch suite's openai-batch) collide on the model_id unique diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 3adc8ba5939..163f5983c3b 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -2,15 +2,12 @@ import asyncio import contextlib import json import uuid -from collections.abc import Callable, Iterator, Mapping -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Iterator, Mapping from dataclasses import dataclass -from functools import partial from pathlib import Path from types import MappingProxyType from typing import Final -import anthropic import httpx import openai import pytest @@ -27,15 +24,15 @@ from integration._support.client import ( object_value, string_value, ) +from integration._support import responses_ws_peer from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply -from integration._support.process import graceful_stop_seconds, owned_proxy_process, owned_upstream, paused +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.tls import write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter from websockets.exceptions import ConnectionClosed -from tests.integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse - _PROVIDER_KEY: Final = "integration-provider-key" _EMBEDDING_MODEL: Final = "text-embedding-3-small" _CHAT_MODEL: Final = "gpt-4o-mini" @@ -45,73 +42,9 @@ _VECTOR: Final = (0.25, 0.5, 0.75) _COUNTED_TOKENS: Final = 17 _CLOSE: Final = MappingProxyType({"connection": "close"}) _RELOAD_SECONDS: Final = 5 -_LONG_ENDPOINT: Final = "/v1/embeddings/" + "x" * 5000 _EMBED_ONLY: Final[Mapping[str, JsonValue]] = MappingProxyType( - {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]} + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"], "enforce_supported_endpoints": True} ) -_CHAT_ENDPOINTS: Final = frozenset( - { - "/v1/chat/completions", - "/chat/completions", - "/v1/completions", - "/completions", - "/v1/messages", - "/v1/responses", - "/responses", - "/v1beta/interactions", - } -) -_CHAT_MODES: Final = frozenset({"chat", "completion", "responses"}) -_CATALOG: Final = TypeAdapter(dict[str, JsonValue]).validate_json( - (Path(__file__).resolve().parents[3] / "litellm/model_prices_and_context_window_backup.json").read_bytes() -) -_CATALOG_ENTRIES: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( - { - key: object_value(entry) - for key, entry in sorted(_CATALOG.items()) - if isinstance(entry, dict) and key != "sample_spec" - } -) - - -def _listed_endpoints(entry: Mapping[str, JsonValue]) -> frozenset[str] | None: - listed: Final = entry.get("supported_endpoints") - if not isinstance(listed, list) or not listed or not all(isinstance(item, str) for item in listed): - return None - return frozenset(string_value(item) for item in listed) - - -def _non_chat_mode(entry: Mapping[str, JsonValue]) -> bool: - mode: Final = entry.get("mode") - return isinstance(mode, str) and mode not in _CHAT_MODES - - -def _lists_only_other_endpoints(entry: Mapping[str, JsonValue]) -> bool: - listed: Final = _listed_endpoints(entry) - return ( - _non_chat_mode(entry) - and listed is not None - and listed.isdisjoint(_CHAT_ENDPOINTS) - and listed != {"/v1/embeddings"} - ) - - -def _lists_a_chat_endpoint(entry: Mapping[str, JsonValue]) -> bool: - listed: Final = _listed_endpoints(entry) - return _non_chat_mode(entry) and listed is not None and not listed.isdisjoint(_CHAT_ENDPOINTS) - - -def _catalog_key(matches: Callable[[Mapping[str, JsonValue]], bool]) -> str: - return next(key for key, entry in _CATALOG_ENTRIES.items() if matches(entry)) - - -_CATALOG_CHAT_KEY: Final = _catalog_key(lambda entry: entry.get("mode") == "chat") -_CATALOG_EMBEDDING_KEY: Final = _catalog_key( - lambda entry: entry.get("mode") == "embedding" and "supported_endpoints" not in entry -) -_CATALOG_OTHER_ENDPOINTS_KEY: Final = _catalog_key(_lists_only_other_endpoints) -_CATALOG_CHAT_LISTED_KEY: Final = _catalog_key(_lists_a_chat_endpoint) -_CATALOG_OTHER_ENDPOINTS_ENTRY: Final = _CATALOG_ENTRIES[_CATALOG_OTHER_ENDPOINTS_KEY] def _embedding_reply() -> Reply: @@ -181,7 +114,7 @@ class _Upstream: def calls(self, marker: str) -> tuple[Request, ...]: self.seen = (*self.seen, *self.wire.drain()) - return tuple(request for request in self.seen if marker.encode() in request.body) + return tuple(request for request in self.seen if marker.encode() in request.body or marker in request.target) def targets(self, marker: str) -> tuple[str, ...]: return tuple(request.target for request in self.calls(marker)) @@ -197,92 +130,15 @@ def _single(declaration: Mapping[str, JsonValue], litellm_model: str = f"openai/ return _Deployment(litellm_model, (declaration,)) -_GATED: Final[Mapping[str, tuple[_Deployment, str]]] = MappingProxyType( - { - "image_generation": ( - _single({"mode": "image_generation", "supported_endpoints": ["/v1/images/generations"]}), - "/v1/images/generations", - ), - "rerank_two_spellings": ( - _single({"mode": "rerank", "supported_endpoints": ["/v1/rerank", "/rerank"]}), - "/rerank, /v1/rerank", - ), - "empty_endpoints": (_single({"mode": "embedding", "supported_endpoints": []}), "none"), - "duplicated_endpoints": ( - _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/embeddings", "/embeddings"]}), - "/embeddings, /v1/embeddings", - ), - "trailing_slash_chat_spelling": ( - _single({"mode": "embedding", "supported_endpoints": ["/v1/chat/completions/"]}), - "/v1/chat/completions/", - ), - "five_kb_endpoint": (_single({"mode": "embedding", "supported_endpoints": [_LONG_ENDPOINT]}), _LONG_ENDPOINT), - "two_embedding_deployments": ( - _Deployment( - f"openai/{_CHAT_MODEL}", - (_EMBED_ONLY, {"mode": "embedding", "supported_endpoints": ["/embeddings"]}), - ), - "/embeddings, /v1/embeddings", - ), - "catalog_key_lists_other_endpoints": ( - _single({**_EMBED_ONLY, "key": _CATALOG_OTHER_ENDPOINTS_KEY}), - "/v1/embeddings", - ), - "catalog_key_names_an_embedding_entry": ( - _single({**_EMBED_ONLY, "key": _CATALOG_EMBEDDING_KEY}), - "/v1/embeddings", - ), - "catalog_key_not_in_the_catalog": ( - _single({**_EMBED_ONLY, "key": "gate-uncatalogued-entry"}), - "/v1/embeddings", - ), - "catalog_key_not_a_string": (_single({**_EMBED_ONLY, "key": 7}), "/v1/embeddings"), - } -) - -_UNGATED: Final[Mapping[str, _Deployment]] = MappingProxyType( - { - "mode_only": _single({"mode": "embedding"}), - "endpoints_only": _single({"supported_endpoints": ["/v1/embeddings"]}), - "endpoints_null": _single({"mode": "embedding", "supported_endpoints": None}), - "mode_null": _single({"mode": None, "supported_endpoints": ["/v1/embeddings"]}), - "mode_blank": _single({"mode": "", "supported_endpoints": ["/v1/embeddings"]}), - "endpoints_a_string": _single({"mode": "embedding", "supported_endpoints": "/v1/embeddings"}), - "endpoints_not_strings": _single({"mode": "embedding", "supported_endpoints": [1, 2]}), - "chat_mode": _single({"mode": "chat", "supported_endpoints": ["/v1/embeddings"]}), - "completion_mode": _single({"mode": "completion", "supported_endpoints": ["/v1/embeddings"]}), - "responses_mode": _single({"mode": "responses", "supported_endpoints": ["/v1/embeddings"]}), - "lists_v1_chat": _single( - {"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/chat/completions"]} - ), - "lists_bare_chat": _single( - {"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/chat/completions"]} - ), - "lists_messages": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/v1/messages"]}), - "lists_bare_responses": _single({"mode": "embedding", "supported_endpoints": ["/v1/embeddings", "/responses"]}), - "lists_generate_content": _single({"mode": "embedding", "supported_endpoints": ["generateContent"]}), - "lists_responses_compact": _single({"mode": "embedding", "supported_endpoints": ["/v1/responses/compact"]}), - "mixed_group": _Deployment(f"openai/{_CHAT_MODEL}", (_EMBED_ONLY, {"mode": "chat"})), - "catalog_key_serves_chat": _single({**_EMBED_ONLY, "key": _CATALOG_CHAT_KEY}), - "catalog_key_lists_a_chat_endpoint": _single({**_EMBED_ONLY, "key": _CATALOG_CHAT_LISTED_KEY}), - "catalog_endpoints_echoed": _single( - { - "key": _CATALOG_OTHER_ENDPOINTS_KEY, - "mode": _CATALOG_OTHER_ENDPOINTS_ENTRY["mode"], - "supported_endpoints": _CATALOG_OTHER_ENDPOINTS_ENTRY["supported_endpoints"], - } - ), - } -) - _SHARED: Final[Mapping[str, _Deployment]] = MappingProxyType( { "embedding_only": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), "embedding_only_second": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "image_generation": _single( + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/images/generations"]} + ), "chat_target": _single({"mode": "chat"}), "chat_rate_limited": _single({"mode": "chat"}), - **{name: deployment for name, (deployment, _supported) in _GATED.items()}, - **_UNGATED, "mode_not_a_string": _single( {"mode": 7, "supported_endpoints": ["/v1/embeddings"]}, "openai/gate-mode-not-a-string" ), @@ -397,11 +253,18 @@ def _failure_row(call_id: str) -> dict[str, JsonValue]: return rows[0] +def _assert_error(response: httpx.Response, expected: dict[str, JsonValue]) -> None: + assert response.status_code == 400, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + message: Final = string_value(error["message"]).split("\n\n", 1)[0] + assert {**body, "error": {**error, "message": message}} == expected, response.text + + def _assert_rejected_before( upstream: _Upstream, response: httpx.Response, marker: str, expected: dict[str, JsonValue] ) -> None: - assert response.status_code == 400, response.text - assert JSON_OBJECT.validate_json(response.content) == expected, response.text + _assert_error(response, expected) assert upstream.calls(marker) == (), upstream.targets(marker) assert _failure_row(response.headers["x-litellm-call-id"])["status"] == "failure" @@ -422,6 +285,7 @@ class _Route: _ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( { + "interactions": _Route("/v1beta/interactions", "/v1beta/interactions", False, False, False, "interactions"), "chat_v1": _Route("/v1/chat/completions", "/v1/chat/completions", False, False, False, "chat"), "chat_v1_stream": _Route("/v1/chat/completions", "/v1/chat/completions", False, True, False, "chat"), "chat_bare": _Route("/chat/completions", "/v1/chat/completions", False, False, False, "chat"), @@ -445,6 +309,7 @@ def _route_body(route: _Route, model: str, marker: str) -> dict[str, JsonValue]: "chat": {"messages": [{"role": "user", "content": marker}]}, "messages": {"max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, "responses": {"input": marker}, + "interactions": {"input": marker}, "cursor": {"input": [{"role": "user", "content": marker}]}, } return { @@ -478,85 +343,6 @@ def test_chat_shaped_route_rejects_an_embedding_only_model_before_the_upstream(r _assert_rejected(rig, response, marker, expected) -def _openai_sdk_error(rig: _Rig, client: str, model: str, marker: str) -> openai.APIStatusError: - base_url: Final = f"{rig.gateway.client.base_url}/v1" - - def sync_call() -> None: - with openai.OpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: - if client == "openai_sync_chat": - sdk.chat.completions.create(model=model, messages=[{"role": "user", "content": marker}]) - else: - sdk.responses.create(model=model, input=marker) - - async def async_call() -> None: - async with openai.AsyncOpenAI(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: - if client == "openai_async_chat_stream": - await sdk.chat.completions.create( - model=model, messages=[{"role": "user", "content": marker}], stream=True - ) - else: - await sdk.responses.create(model=model, input=marker, stream=True) - - def call() -> None: - return sync_call() if client.startswith("openai_sync") else asyncio.run(async_call()) - - with pytest.raises(openai.BadRequestError) as raised: - call() - return raised.value - - -def _anthropic_sdk_error(rig: _Rig, client: str, model: str, marker: str) -> anthropic.APIStatusError: - base_url: Final = str(rig.gateway.client.base_url) - - def sync_call() -> None: - with anthropic.Anthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: - sdk.messages.create(model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]) - - async def async_call() -> None: - async with anthropic.AsyncAnthropic(base_url=base_url, api_key=rig.gateway.key, max_retries=0) as sdk: - await sdk.messages.create( - model=model, max_tokens=16, messages=[{"role": "user", "content": marker}], stream=True - ) - - def call() -> None: - return sync_call() if client == "anthropic_sync" else asyncio.run(async_call()) - - with pytest.raises(anthropic.BadRequestError) as raised: - call() - return raised.value - - -_SDK_CLIENTS: Final[Mapping[str, str]] = MappingProxyType( - { - "openai_sync_chat": "/v1/chat/completions", - "openai_async_chat_stream": "/v1/chat/completions", - "openai_sync_responses": "/v1/responses", - "openai_async_responses_stream": "/v1/responses", - "anthropic_sync": "/v1/messages", - "anthropic_async_stream": "/v1/messages", - } -) - - -@pytest.mark.parametrize("client", tuple(_SDK_CLIENTS)) -def test_sdk_clients_see_the_rejection_as_a_bad_request(rig: _Rig, client: str) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - anthropic_client: Final = client.startswith("anthropic") - error: Final = ( - _anthropic_sdk_error(rig, client, model, marker) - if anthropic_client - else _openai_sdk_error(rig, client, model, marker) - ) - expected: Final = ( - _anthropic_error(model, "/v1/embeddings") - if anthropic_client - else _openai_error(_SDK_CLIENTS[client], model, "/v1/embeddings") - ) - assert error.status_code == 400, error - _assert_rejected(rig, error.response, marker, expected) - - def _sdk_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: async def call() -> tuple[tuple[float, ...], str]: async with openai.AsyncOpenAI( @@ -601,27 +387,12 @@ def test_embeddings_still_reach_the_embedding_only_model(rig: _Rig, client: str) assert [(entry["status"], entry["call_type"]) for entry in row] == [("success", "aembedding")], row -@pytest.mark.parametrize("declaration", tuple(_GATED)) -def test_non_chat_declarations_reject_chat_completions(rig: _Rig, declaration: str) -> None: - model: Final = rig.names[declaration] - marker: Final = _marker() - response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) - _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, _GATED[declaration][1])) - - def _assert_served(rig: _Rig, response: httpx.Response, marker: str, target: str) -> None: assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text assert rig.upstream.targets(marker) == (target,) -@pytest.mark.parametrize("declaration", tuple(_UNGATED)) -def test_chat_capable_or_unclear_declarations_keep_reaching_the_upstream(rig: _Rig, declaration: str) -> None: - marker: Final = _marker() - response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(rig.names[declaration], marker)) - _assert_served(rig, response, marker, "/v1/chat/completions") - - def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_name(rig: _Rig) -> None: target: Final = rig.names["embedding_only"] alias: Final = f"gate-alias-{uuid.uuid4().hex[:12]}" @@ -773,7 +544,7 @@ def test_request_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> No ) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: @@ -821,7 +592,7 @@ def test_a_fallback_chain_that_reaches_chat_two_hops_away_answers(rig: _Rig) -> response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [chat]}]) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions",) * 3 + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) def test_a_fallback_chain_that_loops_between_embedding_only_groups_keeps_the_rejection(rig: _Rig) -> None: @@ -847,46 +618,12 @@ def test_disable_fallbacks_keeps_the_rejection_over_request_fallbacks(rig: _Rig) _assert_rejected_chat(rig, response, marker, model) -def test_client_sent_metadata_disable_fallbacks_does_not_block_the_fallback(rig: _Rig) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - response: Final = _chat( - rig, model, marker, fallbacks=[rig.names["chat_target"]], metadata={"_disable_fallbacks": True} - ) - assert response.status_code == 200, response.text - assert _SERVED_TEXT in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") - - -_MALFORMED_FALLBACKS: Final[Mapping[str, Callable[[_Rig], JsonValue]]] = MappingProxyType( - { - "a_string": lambda rig: rig.names["chat_target"], - "an_object": lambda rig: {rig.names["embedding_only"]: [rig.names["chat_target"]]}, - } -) - - -@pytest.mark.parametrize("shape", tuple(_MALFORMED_FALLBACKS)) -def test_malformed_request_fallbacks_step_aside_for_the_router(rig: _Rig, shape: str) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - response: Final = _chat(rig, model, marker, fallbacks=_MALFORMED_FALLBACKS[shape](rig)) - assert "Supported endpoints" not in response.text, response.text - assert rig.upstream.targets(marker)[:1] == ("/v1/chat/completions",), response.text - - -def _assert_answered_after_one_miss( +def _assert_fallback_answered( upstream: _Upstream, response: httpx.Response, marker: str, target: str = "/v1/chat/completions" ) -> None: assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert upstream.targets(marker) == (target, target), response.text - - -def _assert_the_provider_answered(rig: _Rig, response: httpx.Response, marker: str) -> None: - assert response.status_code == 404, response.text - assert "scripted: not a chat model" in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions",), response.text + assert upstream.targets(marker) == (target,), response.text def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: @@ -894,134 +631,6 @@ def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: return {"fallbacks": [{rig.names["embedding_only"]: [ghost]}, {ghost: [rig.names["chat_target"]]}]} -_ROUTING_STAMPS: Final[Mapping[str, Callable[[JsonValue], Mapping[str, JsonValue]]]] = MappingProxyType( - { - "a_client_routing_stamp": lambda group: {"metadata": {"model_group": group}}, - "a_pre_routing_stamp": lambda group: {"metadata": {"pre_routing_selected_model": group}}, - "a_litellm_metadata_routing_stamp": lambda group: {"litellm_metadata": {"model_group": group}}, - "a_litellm_metadata_pre_routing_stamp": lambda group: { - "litellm_metadata": {"pre_routing_selected_model": group} - }, - "a_routing_stamp_in_both_buckets": lambda group: { - "metadata": {"model_group": group}, - "litellm_metadata": {"model_group": group}, - }, - } -) - - -def _with_a_routing_stamp(stamp: Callable[[JsonValue], Mapping[str, JsonValue]], rig: _Rig) -> Mapping[str, JsonValue]: - model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] - return {"fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], **stamp(second)} - - -_FALLBACKS_THE_CHECK_CANNOT_FOLLOW: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( - { - "a_group_with_no_deployments": _through_a_group_with_no_deployments, - **{name: partial(_with_a_routing_stamp, stamp) for name, stamp in _ROUTING_STAMPS.items()}, - } -) - - -@pytest.mark.parametrize("shape", tuple(_FALLBACKS_THE_CHECK_CANNOT_FOLLOW)) -def test_fallbacks_the_check_cannot_follow_step_aside_for_the_router(rig: _Rig, shape: str) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - response: Final = _chat(rig, model, marker, **_FALLBACKS_THE_CHECK_CANNOT_FOLLOW[shape](rig)) - _assert_answered_after_one_miss(rig.upstream, response, marker) - - -_STAMPED_ROUTES: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( - { - "chat_v1": ("metadata", "/v1/chat/completions"), - "messages": ("litellm_metadata", "/v1/responses"), - "responses_v1": ("litellm_metadata", "/v1/responses"), - } -) - - -@pytest.mark.parametrize("route", tuple(_STAMPED_ROUTES)) -def test_a_client_routing_stamp_steps_aside_on_every_chat_shaped_route(rig: _Rig, route: str) -> None: - model, second = rig.names["embedding_only"], rig.names["embedding_only_second"] - bucket, target = _STAMPED_ROUTES[route] - marker: Final = _marker() - spec: Final = _ROUTES[route] - response: Final = rig.gateway.request( - "POST", - spec.path, - { - **_route_body(spec, model, marker), - "fallbacks": [{model: [rig.names["chat_target"]]}, {second: [model]}], - bucket: {"model_group": second}, - }, - ) - _assert_answered_after_one_miss(rig.upstream, response, marker, target) - - -_STAMPS_THE_CHECK_IGNORES: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( - { - "an_int_routing_stamp": lambda rig: {"metadata": {"model_group": 7}}, - "a_list_routing_stamp": lambda rig: {"metadata": {"model_group": [rig.names["embedding_only_second"]]}}, - "an_empty_routing_stamp": lambda rig: {"metadata": {"model_group": ""}}, - "a_routing_stamp_naming_the_requested_group": lambda rig: { - "metadata": {"model_group": rig.names["embedding_only"]} - }, - "an_int_pre_routing_stamp": lambda rig: {"litellm_metadata": {"pre_routing_selected_model": 7}}, - "an_int_original_group_stamp": lambda rig: {"metadata": {"original_model_group": 7}}, - "an_original_group_stamp_the_proxy_strips": lambda rig: { - "metadata": {"original_model_group": rig.names["embedding_only_second"]} - }, - } -) - - -@pytest.mark.parametrize("stamp", tuple(_STAMPS_THE_CHECK_IGNORES)) -def test_routing_stamps_that_name_no_other_group_keep_the_rejection(rig: _Rig, stamp: str) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - response: Final = _chat( - rig, - model, - marker, - fallbacks=[{rig.names["embedding_only_second"]: [rig.names["chat_target"]]}], - **_STAMPS_THE_CHECK_IGNORES[stamp](rig), - ) - _assert_rejected_chat(rig, response, marker, model) - - -def test_a_routing_stamp_without_any_fallbacks_keeps_the_rejection(rig: _Rig) -> None: - model: Final = rig.names["embedding_only"] - marker: Final = _marker() - response: Final = _chat(rig, model, marker, metadata={"model_group": rig.names["embedding_only_second"]}) - _assert_rejected_chat(rig, response, marker, model) - - -_STAMPS_THE_ROUTER_RESETS: Final[Mapping[str, Callable[[_Rig], Mapping[str, JsonValue]]]] = MappingProxyType( - { - "a_routing_stamp_naming_a_group_with_a_chat_chain": lambda rig: { - "metadata": {"model_group": rig.names["embedding_only_second"]} - }, - "a_pre_routing_stamp_naming_a_group_with_a_chat_chain": lambda rig: { - "metadata": {"pre_routing_selected_model": rig.names["embedding_only_second"]} - }, - "a_five_kb_routing_stamp": lambda rig: {"metadata": {"model_group": "g" * 5000}}, - } -) - - -@pytest.mark.parametrize("stamp", tuple(_STAMPS_THE_ROUTER_RESETS)) -def test_routing_stamps_the_router_resets_step_aside_and_the_provider_answers(rig: _Rig, stamp: str) -> None: - marker: Final = _marker() - response: Final = _chat( - rig, - rig.names["embedding_only"], - marker, - fallbacks=[{rig.names["embedding_only_second"]: [rig.names["chat_target"]]}], - **_STAMPS_THE_ROUTER_RESETS[stamp](rig), - ) - _assert_the_provider_answered(rig, response, marker) - - def test_request_fallbacks_naming_a_chat_backend_name_answer_for_an_embedding_only_primary(rig: _Rig) -> None: model: Final = rig.names["embedding_only"] backend: Final = f"openai/{_CHAT_BACKEND_PREFIX}{uuid.uuid4().hex[:12]}" @@ -1033,7 +642,7 @@ def test_request_fallbacks_naming_a_chat_backend_name_answer_for_an_embedding_on scenario.cleanups.callback(scenario.delete_model, identity) _served_by_every_worker(rig.gateway, frozenset({identity})) response: Final = _chat(rig, model, marker, fallbacks=[{model: [backend]}]) - _assert_answered_after_one_miss(rig.upstream, response, marker) + _assert_fallback_answered(rig.upstream, response, marker) @pytest.mark.parametrize("level", ("key", "team")) @@ -1043,7 +652,7 @@ def test_router_settings_fallbacks_through_a_group_with_no_deployments_answer(ri settings: Final = {"router_settings": dict(_through_a_group_with_no_deployments(rig))} key: Final = scenario.key(**settings) if level == "key" else scenario.key(team_id=scenario.team(**settings)) response: Final = _chat(rig, rig.names["embedding_only"], marker, key=key) - _assert_answered_after_one_miss(rig.upstream, response, marker) + _assert_fallback_answered(rig.upstream, response, marker) def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: @@ -1054,7 +663,7 @@ def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig response: Final = _chat(rig, model, marker, key=scenario.key(team_id=team)) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) def test_key_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: @@ -1065,7 +674,7 @@ def test_key_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: response: Final = _chat(rig, model, marker, key=key) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) def test_key_router_settings_that_cannot_answer_override_team_fallbacks_that_could(rig: _Rig) -> None: @@ -1164,17 +773,6 @@ def test_gemini_routes_reject_an_embedding_only_model_before_the_upstream(rig: _ assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) -def test_non_string_mode_keeps_reaching_the_upstream(rig: _Rig) -> None: - marker: Final = _marker() - response: Final = rig.gateway.request( - "POST", "/v1/chat/completions", _chat_body(rig.names["mode_not_a_string"], marker) - ) - assert response.status_code == 404, response.text - assert "scripted: not a chat model" in response.text, response.text - assert "Supported endpoints" not in response.text, response.text - assert rig.upstream.targets(marker) == ("/v1/chat/completions",) - - def test_repeated_rejections_log_one_failure_row_per_call(rig: _Rig) -> None: model: Final = rig.names["embedding_only"] marker: Final = _marker() @@ -1182,7 +780,8 @@ def test_repeated_rejections_log_one_failure_row_per_call(rig: _Rig) -> None: rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for _ in range(3) ) expected: Final = _openai_error("/v1/chat/completions", model, "/v1/embeddings") - assert [(item.status_code, item.json()) for item in responses] == [(400, expected)] * 3 + for response in responses: + _assert_error(response, expected) call_ids: Final = tuple(item.headers["x-litellm-call-id"] for item in responses) assert len(set(call_ids)) == 3, call_ids assert [_failure_row(call_id)["status"] for call_id in call_ids] == ["failure"] * 3 @@ -1214,7 +813,7 @@ def test_live_declaration_update_turns_the_gate_on_and_off_on_every_worker(rig: holding=_RELOAD_SECONDS * 2, seconds=_RELOAD_SECONDS * 12, ) - assert rejected.json() == _openai_error("/v1/chat/completions", name, "/v1/embeddings"), rejected.text + _assert_error(rejected, _openai_error("/v1/chat/completions", name, "/v1/embeddings")) restored: Final = rig.gateway.request("PATCH", f"/model/{identities[0]}/update", {"model_info": chat_capable}) assert restored.status_code == 200, restored.text held( @@ -1227,182 +826,12 @@ def test_live_declaration_update_turns_the_gate_on_and_off_on_every_worker(rig: _TESTS_ROOT: Final = Path(__file__).resolve().parents[2] _CONFIG: Final = TypeAdapter(dict[str, JsonValue]) -_BURST_ROUTES: Final = ( - "chat_v1", - "chat_v1_stream", - "messages", - "messages_stream", - "responses_v1", - "responses_v1_stream", -) - - -def _scripted(url: str, identity: str, routes: Mapping[str, JsonResponse]) -> None: - registered: Final = httpx.post( - f"{url}/__scenarios", - json={ - "scenario_id": identity, - "response": RoutedResponse(content_type="application/x-routed", routes=dict(routes)).model_dump( - mode="json" - ), - }, - timeout=10, - trust_env=False, - ) - registered.raise_for_status() - - -class _Observed: - def __init__(self, url: str) -> None: - self.url = url - self.items: tuple[dict[str, JsonValue], ...] = () - - def read(self) -> tuple[dict[str, JsonValue], ...]: - payload: Final = JSON_OBJECT.validate_json( - httpx.get(f"{self.url}/__observations?include_method=true", timeout=10, trust_env=False).content - ) - self.items = (*self.items, *(object_value(item) for item in list_value(payload["requests"]))) - return self.items - - def carrying(self, marker: str) -> tuple[str, ...]: - return tuple(str(item["path"]) for item in self.read() if marker in json.dumps(item.get("body"))) - - -def _owned_config(directory: Path, url: str, embed: tuple[str, str], chat: tuple[str, str]) -> Path: - shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) - litellm_settings: Final = object_value(shared["litellm_settings"]) - config: Final = { - **shared, - "model_list": [ - { - "model_name": embed[0], - "litellm_params": { - "model": f"openai/{_EMBEDDING_MODEL}", - "api_base": f"{url}/{embed[1]}/v1", - "api_key": _PROVIDER_KEY, - }, - "model_info": dict(_EMBED_ONLY), - }, - { - "model_name": chat[0], - "litellm_params": { - "model": f"openai/{_CHAT_MODEL}", - "api_base": f"{url}/{chat[1]}/v1", - "api_key": _PROVIDER_KEY, - }, - }, - ], - "litellm_settings": {**litellm_settings, "responses": {"background_mode": {"polling_via_cache": "all"}}}, - } - path: Final = directory / "gate-config.yaml" - path.write_text(yaml.safe_dump(config, sort_keys=False)) - return path def _owned_embedding(owned: Gateway, model: str, marker: str) -> httpx.Response: return owned.request("POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"}) -def _gated_expectation(route: str, model: str) -> dict[str, JsonValue]: - spec: Final = _ROUTES[route] - return ( - _anthropic_error(model, "/v1/embeddings") - if spec.anthropic_shape - else _openai_error(spec.endpoint, model, "/v1/embeddings") - ) - - -def _assert_gated_burst(results: tuple[tuple[str, httpx.Response], ...], model: str) -> tuple[str, ...]: - assert [(response.status_code, JSON_OBJECT.validate_json(response.content)) for _route, response in results] == [ - (400, _gated_expectation(route, model)) for route, _response in results - ], [response.text for _route, response in results] - return tuple(response.headers["x-litellm-call-id"] for _route, response in results) - - -@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -def test_yaml_embedding_only_model_stays_gated_through_an_upstream_pause_and_a_worker_kill( - gateway: Gateway, tmp_path: Path -) -> None: - suffix: Final = uuid.uuid4().hex[:12] - embed: Final = (f"gate-yaml-embed-{suffix}", f"gate-embed-{suffix}") - chat: Final = (f"gate-yaml-chat-{suffix}", f"gate-chat-{suffix}") - embedding: Final = JsonResponse( - content_type="application/json", body=JSON_OBJECT.validate_json(_embedding_reply().body) - ) - chat_completion: Final = JsonResponse( - content_type="application/json", - body={ - "id": f"chatcmpl-{suffix}", - "object": "chat.completion", - "created": 1, - "model": _CHAT_MODEL, - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": _SERVED_TEXT}, "finish_reason": "stop"} - ], - "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, - }, - ) - with owned_upstream(tmp_path) as slot: - upstream: Final = slot.process - assert upstream is not None - url: Final = slot.url - _scripted(url, embed[1], {"POST /v1/embeddings": embedding}) - _scripted(url, chat[1], {"POST /v1/chat/completions": chat_completion}) - config: Final = _owned_config(tmp_path, url, embed, chat) - with owned_proxy_process( - gateway, tmp_path, {"INTEGRATION_UPSTREAM_URL": url}, config=config, workers=2 - ) as owned: - proxy: Final = owned.gateway - eventually( - lambda: proxy.request("POST", "/v1/chat/completions", _chat_body(chat[0], "probe")).status_code, - lambda status: status == 200, - seconds=30, - ) - workers: Final = eventually(owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) - observed: Final = _Observed(url) - background_marker: Final = _marker() - background: Final = proxy.request( - "POST", "/v1/responses", {"model": embed[0], "input": background_marker, "background": True} - ) - assert background.status_code == 400, background.text - assert background.json() == _openai_error("/v1/responses", embed[0], "/v1/embeddings"), background.text - burst: Final = tuple((_BURST_ROUTES[index % len(_BURST_ROUTES)], _marker()) for index in range(18)) - with paused(upstream), ThreadPoolExecutor(max_workers=9) as pool: - futures: Final = tuple( - (route, pool.submit(_send, proxy, _ROUTES[route], embed[0], marker)) for route, marker in burst - ) - paused_results: Final = tuple((route, future.result(timeout=15)) for route, future in futures) - gated_ids: Final = _assert_gated_burst(paused_results, embed[0]) - served_markers: Final = tuple(_marker() for _ in range(8)) - with ThreadPoolExecutor(max_workers=8) as pool: - served: Final = tuple(pool.map(partial(_owned_embedding, proxy, embed[0]), served_markers)) - assert [response.status_code for response in served] == [200] * 8, [item.text for item in served] - embedded_paths: Final = eventually( - lambda: tuple(observed.carrying(marker) for marker in served_markers), - lambda paths: all(len(entry) >= 1 for entry in paths), - seconds=20, - ) - assert embedded_paths == ((f"/{embed[1]}/v1/embeddings",),) * 8, embedded_paths - assert [observed.carrying(marker) for _route, marker in burst] == [()] * 18 - assert observed.carrying(background_marker) == () - assert [_failure_row(call_id)["status"] for call_id in gated_ids] == ["failure"] * 18 - with owned.replacing_worker(min(workers)) as kill: - kill() - after_kill: Final = tuple( - ("chat_v1", _send(proxy, _ROUTES["chat_v1"], embed[0], _marker())) for _ in range(4) - ) - _assert_gated_burst(after_kill, embed[0]) - assert _owned_embedding(proxy, embed[0], _marker()).status_code == 200 - respawned: Final = tuple( - ( - "chat_v1", - proxy.request("POST", "/v1/chat/completions", _chat_body(embed[0], _marker()), headers=_CLOSE), - ) - for _ in range(6) - ) - _assert_gated_burst(respawned, embed[0]) - - @dataclass(frozen=True, slots=True) class _Routed: gateway: Gateway @@ -1473,7 +902,7 @@ def test_router_default_fallbacks_answer_for_an_embedding_only_group(routed: _Ro response: Final = _send(routed.gateway, _ROUTES[route], routed.names["embed"], marker) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert routed.upstream.targets(marker) == (_DEFAULT_FALLBACK_ROUTES[route],) * 2 + assert routed.upstream.targets(marker) == (_DEFAULT_FALLBACK_ROUTES[route],) _OPT_OUTS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( @@ -1506,7 +935,7 @@ def test_model_group_alias_to_an_embedding_only_group_falls_back_to_the_default( ) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert routed.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) @_ROUTED_TIMEOUT @@ -1817,7 +1246,7 @@ def test_request_fallbacks_naming_a_chat_deployment_id_answer_for_an_embedding_o ) assert response.status_code == 200, response.text assert _SERVED_TEXT in response.text, response.text - assert addressed.upstream.targets(marker) == ("/v1/chat/completions", "/v1/chat/completions") + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) @_ROUTED_TIMEOUT @@ -1843,7 +1272,7 @@ def test_request_fallbacks_naming_an_alias_only_a_wildcard_serves_answer_for_an_ "/v1/chat/completions", {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["wildcard_alias"]]}, ) - _assert_answered_after_one_miss(addressed.upstream, response, marker) + _assert_fallback_answered(addressed.upstream, response, marker) @_ROUTED_TIMEOUT @@ -1854,7 +1283,7 @@ def test_yaml_fallbacks_through_a_group_with_no_deployments_answer_for_an_embedd response: Final = addressed.gateway.request( "POST", "/v1/chat/completions", _chat_body(addressed.names["ghost_primary"], marker) ) - _assert_answered_after_one_miss(addressed.upstream, response, marker) + _assert_fallback_answered(addressed.upstream, response, marker) @_ROUTED_TIMEOUT @@ -1874,7 +1303,100 @@ def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rej for marker in markers ) expected: Final = _openai_error("/v1/chat/completions", name, "/v1/embeddings") - assert [(response.status_code, JSON_OBJECT.validate_json(response.content)) for response in responses] == [ - (400, expected) - ] * len(markers), [response.text for response in responses] + for response in responses: + _assert_error(response, expected) assert [addressed.upstream.calls(marker) for marker in markers] == [()] * len(markers) + + +@pytest.mark.parametrize("metadata", ({"model_group": 7}, {"model_group": ["wrong-group"]})) +def test_rejection_with_invalid_metadata_is_logged_without_blocking_embeddings( + rig: _Rig, metadata: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + model: Final = rig.names["embedding_only"] + rejected: Final = _chat(rig, model, marker, metadata=metadata) + _assert_rejected_chat(rig, rejected, marker, model) + embedded_marker: Final = _marker() + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": embedded_marker, "encoding_format": "float"} + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(embedded_marker) == ("/v1/embeddings",) + + +def test_existing_non_chat_metadata_on_a_chat_backend_remains_advisory(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy( + scenario, + rig.api_base, + "legacy-audio", + _single({"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + + +def test_a_mixed_group_skips_its_restricted_embedding_deployment(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name: Final = f"gate-mixed-{uuid.uuid4().hex[:12]}" + embed: Final = _create( + scenario.gateway, + name, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base, "order": 0}, + _EMBED_ONLY, + ) + scenario.cleanups.callback(scenario.delete_model, embed) + chat: Final = _create( + scenario.gateway, name, {"model": f"openai/{_CHAT_MODEL}", "api_base": rig.api_base, "order": 1}, {} + ) + scenario.cleanups.callback(scenario.delete_model, chat) + _served_by_every_worker(rig.gateway, frozenset({embed, chat})) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + calls: Final = rig.upstream.calls(marker) + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in calls] == [_CHAT_MODEL] + + +def test_local_rate_limit_with_only_unsupported_fallbacks_keeps_the_original_429(rig: _Rig) -> None: + limited: Final = rig.names["chat_rate_limited"] + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: [rig.names["embedding_only"]]}]} + ) + first: Final = _chat(rig, limited, _marker(), key=key) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = _chat(rig, limited, marker, key=key) + assert second.status_code == 429, second.text + assert "Supported endpoints" not in second.text, second.text + assert rig.upstream.calls(marker) == () + + +def test_responses_websocket_uses_a_chat_fallback_without_calling_the_embedding_provider( + rig: _Rig, tmp_path: Path +) -> None: + marker: Final = _marker() + certificate: Final = write_self_signed_cert(tmp_path) + with ( + responses_ws_peer.responses_peer( + certificate, lambda identity: responses_ws_peer.events(identity, _SERVED_TEXT) + ) as peer, + owned_proxy_process(rig.gateway, tmp_path, {"SSL_CERT_FILE": str(certificate[0])}, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + fallback: Final = scenario.model(model=f"openai/{_CHAT_MODEL}", api_base=peer.url, api_key=_PROVIDER_KEY) + primary: Final = scenario.model( + model=f"openai/{_EMBEDDING_MODEL}", api_base=f"{rig.api_base}/{marker}", model_info=_EMBED_ONLY + ) + key: Final = scenario.key(router_settings={"fallbacks": [{primary: [fallback]}]}) + (received,) = asyncio.run( + responses_ws_peer.session(str(owned.gateway.client.base_url), key, primary, (marker,)) + ) + assert responses_ws_peer.completed(received)["status"] == "completed" + assert any(event.get("delta") == _SERVED_TEXT for event in received), received + (forwarded,) = responses_ws_peer.drain(peer.frames) + assert marker in json.dumps(forwarded), forwarded + assert rig.upstream.calls(marker) == () diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 51278607e40..5e6b074a17e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -19,7 +19,6 @@ from httpx import Response import litellm from litellm.proxy.proxy_server import app -from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from tests._master_key import MASTER_KEY @@ -2110,97 +2109,6 @@ class TestResponsesWSFirstFrameModelAuth: ) -class TestResponsesWSDeclaredNonChatModel: - @pytest.mark.asyncio - @pytest.mark.parametrize("requested_model", ["embed-only", "embed-alias"]) - @pytest.mark.parametrize("use_query_model", [False, True]) - @pytest.mark.parametrize( - "model_info, rejected", - [ - ({"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, True), - ({"mode": "embedding"}, False), - ], - ) - async def test_session_closes_on_a_model_declared_non_chat_before_routing( - self, requested_model: str, use_query_model: bool, model_info: dict[str, object], rejected: bool - ): - from litellm.proxy.response_api_endpoints.endpoints import ( - responses_websocket_endpoint, - ) - - router: Final = litellm.Router( - model_list=[ - { - "model_name": "embed-only", - "litellm_params": {"model": "sagemaker/openai/embed-endpoint", "aws_region_name": "us-east-1"}, - "model_info": model_info, - } - ] - ) - ws = MagicMock() - ws.headers = {} - ws.query_params = {} - ws.scope = {"headers": []} - ws.url = "ws://testserver/v1/responses" - ws.accept = AsyncMock() - ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "response.create", "model": requested_model, "input": []}) - ) - ws.send_text = AsyncMock() - ws.close = AsyncMock() - key_aliases: Final = {"embed-alias": "embed-only"} - - async def alias_resolving_pre_call(**kwargs): - return {"model": key_aliases.get(kwargs["model"], kwargs["model"]), "litellm_metadata": {}}, MagicMock() - - processor = MagicMock() - processor.common_processing_pre_call_logic = AsyncMock(side_effect=alias_resolving_pre_call) - - async def fake_llm_call(): - return None - - with ( - patch( # test-quality-ok: first-frame model auth needs a live key table and has its own tests above - "litellm.proxy.response_api_endpoints.endpoints._enforce_responses_ws_first_frame_model_auth", - new_callable=AsyncMock, - ), - patch( # test-quality-ok: the pre-call processor needs a live proxy; whether the session reaches it is under test - "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", - return_value=processor, - ), - patch( # test-quality-ok: routing is the seam that would reach the provider - "litellm.proxy.route_llm_request.route_request", - new_callable=AsyncMock, - return_value=fake_llm_call(), - ) as mock_route_request, - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the session reads the proxy's router global - patch( # test-quality-ok: the session reads the proxy's logging global; the failure hook releases reservations - "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", - new_callable=AsyncMock, - ) as mock_failure_hook, - ): - await responses_websocket_endpoint( - websocket=ws, - model=requested_model if use_query_model else None, - user_api_key_dict=MagicMock(team_id=None), - ) - - if not rejected: - mock_route_request.assert_awaited_once() - ws.close.assert_not_awaited() - mock_failure_hook.assert_not_awaited() - return - mock_failure_hook.assert_awaited_once() - assert isinstance(mock_failure_hook.await_args.kwargs["original_exception"], ProxyUnsupportedEndpointError) - assert mock_failure_hook.await_args.kwargs["request_data"]["model"] == "embed-only" - frame = json.loads(ws.send_text.await_args.args[0]) - assert frame["type"] == "error" - assert frame["error"]["type"] == "invalid_request_error" - assert frame["error"]["message"] == ( - "/v1/responses is not supported by model=embed-only. Supported endpoints: /v1/embeddings" - ) - ws.close.assert_awaited_once_with(code=1008, reason="Pre-call error") - mock_route_request.assert_not_awaited() class TestReadWSModelFromFirstFrameErrors: diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 2e67e3857d3..c302bee3060 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -73,7 +73,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_spend_cache_key, team_member_budget_entity_id, ) -from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError from litellm.proxy.utils import ProxyLogging from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment @@ -713,70 +712,8 @@ class TestProxyBaseLLMRequestProcessing: assert returned_data["metadata"]["tags"] == ["existing-tag"] tag_budget_check.assert_not_awaited() - @staticmethod - def _router_with_embedding_only_model() -> Router: - return Router( - model_list=[ - { - "model_name": "embed-only", - "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-test"}, - "model_info": {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, - }, - ] - ) - @pytest.mark.asyncio - @pytest.mark.parametrize( - "requested_model, key_aliases", - [("embed-only", {}), ("embed-alias", {"embed-alias": "embed-only"})], - ) - async def test_common_processing_pre_call_logic_rejects_chat_on_embedding_only_model_before_pre_call_hooks( - self, monkeypatch, requested_model, key_aliases - ): - processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( - monkeypatch, - request_body={"model": requested_model, "messages": [{"role": "user", "content": "hello"}]}, - guardrail_tags=[], - ) - with pytest.raises(ProxyUnsupportedEndpointError) as exc_info: - await processing_obj.common_processing_pre_call_logic( - request=mock_request, - general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases=key_aliases), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=mock_proxy_config, - route_type="acompletion", - llm_router=self._router_with_embedding_only_model(), - ) - - assert exc_info.value.code == "400" - assert exc_info.value.message == ( - "/v1/chat/completions is not supported by model=embed-only. Supported endpoints: /v1/embeddings" - ) - mock_proxy_logging_obj.pre_call_hook.assert_not_awaited() - - @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_lets_embeddings_through_on_embedding_only_model(self, monkeypatch): - processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, _ = self._guardrail_tag_budget_harness( - monkeypatch, - request_body={"model": "embed-only", "input": ["hello"]}, - guardrail_tags=[], - route="/v1/embeddings", - ) - - returned_data, _ = await processing_obj.common_processing_pre_call_logic( - request=mock_request, - general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=mock_proxy_config, - route_type="aembedding", - llm_router=self._router_with_embedding_only_model(), - ) - - assert returned_data["model"] == "embed-only" - mock_proxy_logging_obj.pre_call_hook.assert_awaited_once() @pytest.mark.asyncio async def test_common_processing_pre_call_logic_skips_guardrail_tag_budget_check_for_zero_cost_model( @@ -6846,6 +6783,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: return processor.data, logging_obj mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] with patch.object( @@ -6887,6 +6825,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: ) mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = None with patch.object( @@ -6926,6 +6865,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: ) mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] with patch.object( @@ -6953,56 +6893,6 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == "gpt-4" - @pytest.mark.asyncio - async def test_fallback_skips_a_model_that_rejects_the_endpoint(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError - from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError - - processor = ProxyBaseLLMRequestProcessing(data={"model": "chat-rl"}) - attempted: list[str] = [] - - async def mock_pre_call_logic(**kwargs): - model = processor.data.get("model") - attempted.append(model) - if model == "chat-rl": - raise ProxyRateLimitError(detail="RPM limit exceeded", headers={"retry-after": "30"}) - if model == "embed-only": - raise ProxyUnsupportedEndpointError( - message="/v1/chat/completions is not supported by model=embed-only", - type="invalid_request_error", - param="model", - code=400, - ) - return processor.data, MagicMock() - - mock_router = MagicMock() - mock_router.fallbacks = [{"chat-rl": ["embed-only", "chat-ctl"]}] - - with patch.object( - processor, - "common_processing_pre_call_logic", - side_effect=mock_pre_call_logic, - ): - await processor._pre_call_with_fallbacks( - request=MagicMock(), - general_settings={}, - proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), - version=None, - proxy_config=MagicMock(), - user_model=None, - user_temperature=None, - user_request_timeout=None, - user_max_tokens=None, - user_api_base=None, - model="chat-rl", - route_type="acompletion", - llm_router=mock_router, - ) - - assert attempted == ["chat-rl", "embed-only", "chat-ctl"] - assert processor.data["model"] == "chat-ctl" @pytest.mark.asyncio async def test_fallback_uses_key_level_router_settings(self): @@ -7020,6 +6910,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: return processor.data, MagicMock() mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() @@ -7067,6 +6958,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: ) mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] with patch.object( @@ -7111,6 +7003,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: raise ValueError("unexpected auth failure on fallback") mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] with patch.object( @@ -7220,6 +7113,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: return processor.data, MagicMock() mock_router = MagicMock() + mock_router.endpoint_error_for_request.return_value = None mock_router.fallbacks = [{primary_model: [fallback_model]}] with patch.object( diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 626de363e33..d05599cf73b 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -7,16 +7,8 @@ from unittest.mock import MagicMock from fastapi import HTTPException -import litellm -from litellm import Router -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap -from litellm.types.router import RouterGeneralSettings from litellm.proxy.route_llm_request import ( - CHAT_SHAPED_ENDPOINTS, - CHAT_SHAPED_MODES, ProxyModelNotFoundError, - ProxyUnsupportedEndpointError, - endpoint_unsupported_by_model, route_request, ) @@ -1532,38 +1524,6 @@ async def test_route_request_read_through_recovers_model_created_on_sibling_repl assert table.find_many_wheres[0] == {"model_name": model_name} -@pytest.mark.asyncio -async def test_route_request_checks_the_endpoints_of_a_model_loaded_by_read_through(monkeypatch): - import litellm.proxy.proxy_server as proxy_server - - model_name: Final = "embed-only-on-sibling-replica" - row: Final = SimpleNamespace( - model_id=f"{model_name}-id", - model_name=model_name, - litellm_params={"model": "openai/text-embedding-3-small", "api_key": "fake"}, - model_info={"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]}, - blocked=False, - ) - router: Final = litellm.Router( - model_list=[{"model_name": "some-other-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] - ) - fake_prisma, table = _fake_prisma_client_with_models([row]) - monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) - monkeypatch.setattr(proxy_server, "store_model_in_db", True) - monkeypatch.setattr(proxy_server, "llm_router", router) - - with pytest.raises(ProxyUnsupportedEndpointError) as exc_info: - await route_request( - data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]}, - llm_router=router, - user_model=None, - route_type="acompletion", - ) - - assert exc_info.value.message == ( - f"/v1/chat/completions is not supported by model={model_name}. Supported endpoints: /v1/embeddings" - ) - assert table.find_many_wheres == [{"model_name": model_name}] @pytest.mark.asyncio @@ -1807,416 +1767,6 @@ async def test_route_request_router_settings_override_skips_null_fields(): assert "model_group_retry_policy" not in call_kwargs -CHAT_SHAPED_ROUTES: Final = ( - ("acompletion", "/v1/chat/completions"), - ("atext_completion", "/v1/completions"), - ("anthropic_messages", "/v1/messages"), - ("aresponses", "/v1/responses"), - ("agenerate_content", "generateContent"), - ("agenerate_content_stream", "streamGenerateContent"), - ("acreate_interaction", "/v1beta/interactions"), - ("_aresponses_websocket", "/v1/responses"), - ("acompact_responses", "/v1/responses/compact"), -) - -CHAT_SHAPED_ROUTE_TYPES: Final = tuple(route_type for route_type, _ in CHAT_SHAPED_ROUTES) - -EMBEDDING_ONLY: Final = {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"]} - - -def _embedding_deployment(model_name: str, model_info: dict[str, object]) -> dict[str, object]: - return { - "model_name": model_name, - "litellm_params": {"model": "sagemaker/openai/embed-endpoint", "aws_region_name": "us-east-1"}, - "model_info": model_info, - } - - -def _router_with(*deployments: dict[str, object], **router_kwargs: object) -> Router: - return Router(model_list=list(deployments), **router_kwargs) - - -def _unsupported( - route_type: str, model: object, router: Router | None, **request: object -) -> ProxyUnsupportedEndpointError | None: - return endpoint_unsupported_by_model( - route_type=route_type, data={"model": model, **request}, llm_router=router, team_id=None - ) - - -@pytest.mark.parametrize("route_type, endpoint", CHAT_SHAPED_ROUTES) -def test_endpoint_unsupported_rejects_chat_routes_on_embedding_only_model(route_type, endpoint): - error: Final = _unsupported( - route_type, "embed-only", _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) - ) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert error.code == "400" - assert error.param == "model" - assert error.type == "invalid_request_error" - assert error.message == f"{endpoint} is not supported by model=embed-only. Supported endpoints: /v1/embeddings" - - -@pytest.mark.parametrize("route_type", ["aembedding", "aimage_generation", "aspeech", "_arealtime"]) -def test_endpoint_unsupported_leaves_non_chat_routes_alone(route_type): - router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) - - assert _unsupported(route_type, "embed-only", router) is None - - -@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) -@pytest.mark.parametrize( - "model_info", - [ - {"mode": "embedding"}, - {}, - {"mode": "embedding", "supported_endpoints": None}, - {"mode": "embedding", "supported_endpoints": "/v1/embeddings"}, - {"mode": 1, "supported_endpoints": ["/v1/embeddings"]}, - {"mode": "", "supported_endpoints": ["/v1/embeddings"]}, - {"mode": " ", "supported_endpoints": ["/v1/embeddings"]}, - ], -) -def test_endpoint_unsupported_keeps_todays_behavior_without_a_valid_declaration(route_type, model_info): - router: Final = _router_with(_embedding_deployment("embed-model", model_info)) - - assert _unsupported(route_type, "embed-model", router) is None - - -@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) -@pytest.mark.parametrize( - "model_info", - [ - {"supported_endpoints": ["/v1/embeddings"]}, - {"mode": "chat", "supported_endpoints": ["/v1/batch"]}, - {"mode": "completion", "supported_endpoints": ["/fine_tuning"]}, - {"mode": "responses", "supported_endpoints": ["/v1/realtime"]}, - {"mode": " Chat ", "supported_endpoints": ["/v1/embeddings"]}, - ], -) -def test_endpoint_unsupported_needs_a_non_chat_mode_to_reject(route_type, model_info): - router: Final = _router_with(_embedding_deployment("advisory", model_info)) - - assert _unsupported(route_type, "advisory", router) is None - - -@pytest.mark.parametrize("route_type", CHAT_SHAPED_ROUTE_TYPES) -@pytest.mark.parametrize( - "declared", - [ - ["/v1/messages"], - ["/v1/responses"], - ["/v1/chat/completions"], - ["/chat/completions"], - ["/responses"], - ["/v1/completions"], - ["/completions"], - ["/v1beta/interactions"], - ["generateContent"], - ["streamGenerateContent"], - ["/v1/responses/compact"], - ], -) -def test_endpoint_unsupported_allows_every_chat_route_when_any_chat_endpoint_is_declared(route_type, declared): - router: Final = _router_with( - _embedding_deployment("chat-model", {"mode": "embedding", "supported_endpoints": declared}) - ) - - assert _unsupported(route_type, "chat-model", router) is None - - -def test_endpoint_unsupported_allows_a_group_with_one_deployment_that_declares_nothing(): - router: Final = _router_with( - _embedding_deployment("mixed", EMBEDDING_ONLY), - _embedding_deployment("mixed", {"mode": "embedding"}), - ) - - assert _unsupported("acompletion", "mixed", router) is None - - -def test_endpoint_unsupported_names_every_endpoint_the_group_declares(): - router: Final = _router_with( - _embedding_deployment("embed-group", EMBEDDING_ONLY), - _embedding_deployment( - "embed-group", {"mode": "embedding", "supported_endpoints": ["/v1/rerank", "/v1/embeddings"]} - ), - ) - - error: Final = _unsupported("acompletion", "embed-group", router) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert error.message == ( - "/v1/chat/completions is not supported by model=embed-group. Supported endpoints: /v1/embeddings, /v1/rerank" - ) - - -def test_endpoint_unsupported_says_none_for_an_empty_declared_list(): - router: Final = _router_with(_embedding_deployment("nothing", {"mode": "embedding", "supported_endpoints": []})) - - error: Final = _unsupported("aresponses", "nothing", router) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert error.message == "/v1/responses is not supported by model=nothing. Supported endpoints: none" - - -def test_endpoint_unsupported_follows_a_model_group_alias(): - router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), - model_group_alias={"embed-alias": "embed-only"}, - ) - - error: Final = _unsupported("anthropic_messages", "embed-alias", router) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert error.message == "/v1/messages is not supported by model=embed-alias. Supported endpoints: /v1/embeddings" - - -UNRELATED_WILDCARD_DEPLOYMENT: Final = { - "model_name": "openai/*", - "litellm_params": {"model": "openai/*", "api_key": "test-key"}, -} - - -@pytest.mark.parametrize("model", ["embed-1", "sagemaker/openai/embed-endpoint"]) -@pytest.mark.parametrize("extra_deployments", [(), (UNRELATED_WILDCARD_DEPLOYMENT,)]) -def test_endpoint_unsupported_rejects_chat_sent_to_the_deployment_id_or_backend_model(model, extra_deployments): - router: Final = _router_with( - _embedding_deployment("embed-only", {**EMBEDDING_ONLY, "id": "embed-1"}), - CHAT_FALLBACK_DEPLOYMENT, - *extra_deployments, - ) - - error: Final = _unsupported("acompletion", model, router) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert error.message == ( - f"/v1/chat/completions is not supported by model={model}. Supported endpoints: /v1/embeddings" - ) - - -@pytest.mark.parametrize( - "extra_deployments, router_kwargs", - [ - ((), {"router_general_settings": RouterGeneralSettings(pass_through_all_models=True)}), - (({"model_name": "sagemaker/*", "litellm_params": {"model": "openai/*", "api_key": "test-key"}},), {}), - ], -) -def test_endpoint_unsupported_leaves_a_backend_model_to_the_route_that_serves_it(extra_deployments, router_kwargs): - router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), *extra_deployments, **router_kwargs - ) - - assert _unsupported("acompletion", "sagemaker/openai/embed-endpoint", router) is None - - -@pytest.mark.parametrize("model", ["unknown-model", "", None]) -def test_endpoint_unsupported_ignores_models_without_deployments(model): - router: Final = _router_with(_embedding_deployment("embed-only", EMBEDDING_ONLY)) - - assert _unsupported("acompletion", model, router) is None - - -def test_endpoint_unsupported_skips_the_check_without_a_router(): - assert _unsupported("acompletion", "embed-only", None) is None - - -CHAT_FALLBACK_DEPLOYMENT: Final = { - "model_name": "chat-fallback", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"}, - "model_info": {"id": "chat-fallback-id"}, -} - -EMBEDDING_FALLBACK_DEPLOYMENT: Final = _embedding_deployment( - "embed-fallback", {**EMBEDDING_ONLY, "id": "embed-fallback-id"} -) - - -@pytest.mark.parametrize( - "router_kwargs, request_kwargs", - [ - ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {}), - ({"default_fallbacks": ["chat-fallback"]}, {}), - ({}, {"fallbacks": ["chat-fallback"]}), - ({}, {"fallbacks": [{"embed-only": ["chat-fallback"]}]}), - ({}, {"router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), - ({"fallbacks": [{"embed-only": ["embed-fallback", "missing-group", "chat-fallback"]}]}, {}), - ({}, {"fallbacks": [{"model": "chat-fallback", "temperature": 0.2}]}), - ({"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "chat-fallback"}}, {}), - ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {}), - ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"metadata": {"_disable_fallbacks": True}}), - ({}, {"fallbacks": {"chat-fallback": ["embed-fallback"]}}), - ({}, {"fallbacks": [{"embed-only": [{"model": "embed-fallback"}]}]}), - ({}, {"fallbacks": ["chat-fallback-id"]}), - ({"fallbacks": [{"embed-only": ["missing-group"]}]}, {}), - ({}, {"fallbacks": ["embed-fallback", "missing-group"]}), - ({"fallbacks": [{"embed-only": ["missing-group"]}, {"missing-group": ["chat-fallback"]}]}, {}), - ({}, {"fallbacks": ["openai/gpt-4o-mini"]}), - ( - {"fallbacks": [{"embed-only": ["chat-alias"]}], "model_group_alias": {"chat-alias": "openai/gpt-4o-mini"}}, - {}, - ), - ( - {"fallbacks": [{"embed-only": ["chat-fallback"]}, {"embed-fallback": ["embed-only"]}]}, - {"metadata": {"model_group": "embed-fallback"}}, - ), - ( - {"fallbacks": [{"embed-only": ["chat-fallback"]}, {"embed-fallback": ["embed-only"]}]}, - {"litellm_metadata": {"pre_routing_selected_model": "embed-fallback"}}, - ), - ], -) -def test_endpoint_unsupported_lets_a_configured_fallback_answer(router_kwargs, request_kwargs): - router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), - CHAT_FALLBACK_DEPLOYMENT, - EMBEDDING_FALLBACK_DEPLOYMENT, - **router_kwargs, - ) - - assert _unsupported("acompletion", "embed-only", router, **request_kwargs) is None - - -@pytest.mark.parametrize( - "router_kwargs, request_kwargs", - [ - ({"fallbacks": [{"other-group": ["chat-fallback"]}]}, {}), - ({"fallbacks": [{"embed-only": ["chat-fallback"]}]}, {"disable_fallbacks": True}), - ({"default_fallbacks": ["chat-fallback"]}, {"fallbacks": []}), - ( - {"fallbacks": [{"embed-only": ["chat-fallback"]}]}, - {"router_settings_override": {"fallbacks": [{"other-group": ["chat-fallback"]}]}}, - ), - ({}, {"fallbacks": [], "router_settings_override": {"fallbacks": [{"embed-only": ["chat-fallback"]}]}}), - ({"fallbacks": [{"embed-only": ["embed-fallback"]}]}, {}), - ({"fallbacks": [{"embed-only": ["embed-only"]}]}, {}), - ({"default_fallbacks": ["embed-fallback"]}, {}), - ({}, {"fallbacks": ["embed-fallback-id"]}), - ({}, {"fallbacks": [{"model": "embed-fallback", "temperature": 0.2}]}), - ({"fallbacks": [{"embed-only": ["embed-alias"]}], "model_group_alias": {"embed-alias": "embed-fallback"}}, {}), - ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["embed-only"]}]}, {}), - ({"fallbacks": [{"embed-only": ["embed-fallback"]}, {"embed-fallback": ["chat-fallback"]}]}, {"fallbacks": []}), - ], -) -def test_endpoint_unsupported_rejects_when_no_fallback_would_answer(router_kwargs, request_kwargs): - router: Final = _router_with( - _embedding_deployment("embed-only", EMBEDDING_ONLY), - CHAT_FALLBACK_DEPLOYMENT, - EMBEDDING_FALLBACK_DEPLOYMENT, - **router_kwargs, - ) - - assert isinstance( - _unsupported("acompletion", "embed-only", router, **request_kwargs), ProxyUnsupportedEndpointError - ) - - -def _catalog_key_declaring_non_chat() -> str: - return min( - key - for key, entry in GetModelCostMap.loaded_model_cost_map().items() - if entry.get("litellm_provider") == "openai" - and entry.get("mode") not in CHAT_SHAPED_MODES - and isinstance(endpoints := entry.get("supported_endpoints"), list) - and CHAT_SHAPED_ENDPOINTS.isdisjoint(endpoints) - ) - - -def _catalog_declaration(catalog_key: str) -> dict[str, object]: - entry: Final = GetModelCostMap.loaded_model_cost_map()[catalog_key] - return {"mode": entry["mode"], "supported_endpoints": entry["supported_endpoints"]} - - -@pytest.mark.parametrize( - "declaration, with_key, rejected", - [ - pytest.param("catalog", True, False, id="echo-of-the-catalog-row"), - pytest.param("catalog-mode-typed-endpoints", True, True, id="endpoints-typed-over-an-echo"), - pytest.param("typed-mode-catalog-endpoints", True, False, id="endpoints-left-as-echoed"), - pytest.param("typed", True, True, id="both-typed-over-an-echo"), - pytest.param("catalog", False, True, id="catalog-values-typed-by-hand"), - ], -) -def test_endpoint_unsupported_ignores_endpoints_copied_from_the_loaded_catalog(declaration, with_key, rejected): - catalog_key: Final = _catalog_key_declaring_non_chat() - catalog: Final = _catalog_declaration(catalog_key) - declared: Final = { - "catalog": catalog, - "catalog-mode-typed-endpoints": {**catalog, "supported_endpoints": ["/v1/embeddings"]}, - "typed-mode-catalog-endpoints": {**catalog, "mode": "embedding"}, - "typed": EMBEDDING_ONLY, - }[declaration] - router: Final = _router_with( - { - "model_name": "catalog-row", - "litellm_params": {"model": catalog_key, "api_key": "test-key"}, - "model_info": {**({"key": catalog_key} if with_key else {}), **declared}, - } - ) - - assert litellm.model_cost[catalog_key]["supported_endpoints"] == declared["supported_endpoints"] - assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected - - -def _catalog_key_serving_chat_by_mode() -> str: - return min( - key - for key, entry in GetModelCostMap.loaded_model_cost_map().items() - if entry.get("mode") == "chat" and "supported_endpoints" not in entry - ) - - -def _catalog_key_serving_chat_by_endpoints() -> str: - return min( - key - for key, entry in GetModelCostMap.loaded_model_cost_map().items() - if entry.get("mode") not in CHAT_SHAPED_MODES - and isinstance(endpoints := entry.get("supported_endpoints"), list) - and not CHAT_SHAPED_ENDPOINTS.isdisjoint(endpoints) - ) - - -@pytest.mark.parametrize( - "catalog_key, rejected", - [ - pytest.param(_catalog_key_serving_chat_by_mode, False, id="key-of-a-chat-mode-catalog-entry"), - pytest.param(_catalog_key_serving_chat_by_endpoints, False, id="key-of-a-catalog-entry-listing-chat"), - pytest.param(_catalog_key_declaring_non_chat, True, id="key-of-a-non-chat-catalog-entry"), - pytest.param(lambda: "not-in-the-catalog", True, id="key-missing-from-the-catalog"), - ], -) -def test_endpoint_unsupported_ignores_non_chat_values_on_a_row_whose_catalog_entry_serves_chat(catalog_key, rejected): - router: Final = _router_with( - { - "model_name": "catalog-row", - "litellm_params": {"model": "openai/catalog-row", "api_key": "test-key"}, - "model_info": {"key": catalog_key(), **EMBEDDING_ONLY}, - } - ) - - assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected - - -class _LookupCountingRouter(Router): - def __init__(self, **router_kwargs: object) -> None: - self.looked_up: list[str | None] = [] - super().__init__(**router_kwargs) - self.looked_up.clear() - - def get_model_list(self, model_name: str | None = None, team_id: str | None = None): - self.looked_up.append(model_name) - return super().get_model_list(model_name=model_name, team_id=team_id) - - -def test_endpoint_unsupported_looks_up_each_repeated_fallback_group_once(): - fallback_groups: Final = ("embed-a", "embed-b", "embed-c") - router: Final = _LookupCountingRouter( - model_list=[_embedding_deployment(name, EMBEDDING_ONLY) for name in ("embed-only", *fallback_groups)] - ) - - error: Final = _unsupported("acompletion", "embed-only", router, fallbacks=list(fallback_groups) * 5) - - assert isinstance(error, ProxyUnsupportedEndpointError) - assert sorted(name for name in router.looked_up if name in fallback_groups) == sorted(fallback_groups) def _router_with_defaults(default_litellm_params: dict[str, object]): diff --git a/tests/unit/router_utils/test_endpoint_eligibility.py b/tests/unit/router_utils/test_endpoint_eligibility.py new file mode 100644 index 00000000000..82d161743bb --- /dev/null +++ b/tests/unit/router_utils/test_endpoint_eligibility.py @@ -0,0 +1,76 @@ +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.router_utils.endpoint_eligibility import UnsupportedEndpointError, eligible_deployments, endpoint_selection + +_DEPLOYMENT: Final = TypeAdapter(list[dict[str, object]]) + + +def _deployment(info: Mapping[str, object], identity: str = "embedding") -> dict[str, object]: + return _DEPLOYMENT.validate_python( + [ + { + "model_name": "shared", + "litellm_params": {"model": "openai/test-model"}, + "model_info": {"id": identity, **info}, + } + ] + )[0] + + +@pytest.mark.parametrize("call_type", ["acompletion", "anthropic_messages", "aresponses", "acreate_interaction"]) +def test_explicit_restriction_excludes_an_embedding_deployment_from_a_mixed_group(call_type: str) -> None: + embedding: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + chat: Final = _deployment({"supported_endpoints": ["/v1/chat/completions"]}, "chat") + with endpoint_selection(call_type): + result: Final = eligible_deployments("shared", (embedding, chat)) + assert result == (chat,) + + +@pytest.mark.parametrize( + "info", + [ + {"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]}, + {"enforce_supported_endpoints": False, "supported_endpoints": ["/v1/embeddings"]}, + {"enforce_supported_endpoints": True}, + {"enforce_supported_endpoints": True, "supported_endpoints": None}, + {"enforce_supported_endpoints": True, "supported_endpoints": "invalid"}, + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/messages"]}, + ], +) +def test_advisory_or_incomplete_metadata_preserves_routing(info: Mapping[str, object]) -> None: + deployment: Final = _deployment(info) + with endpoint_selection("acompletion"): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) + + +def test_unsupported_error_explains_the_requested_operation_and_available_endpoints() -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection("acreate_interaction"): + result: Final = eligible_deployments("shared", (deployment,)) + assert isinstance(result, UnsupportedEndpointError) + assert result.status_code == 400 + assert "/v1beta/interactions" in result.message + assert "/v1/embeddings" in result.message + + +@pytest.mark.parametrize("call_type", ["aembedding", "aspeech", "_arealtime", "aget_responses"]) +def test_chat_restriction_does_not_block_other_operations(call_type: str) -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection(call_type): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) + + +def test_nested_non_chat_calls_do_not_inherit_the_outer_chat_restriction() -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection("acompletion"): + with endpoint_selection("aembedding"): + embedding_result: Final = eligible_deployments("shared", (deployment,)) + chat_result: Final = eligible_deployments("shared", (deployment,)) + assert embedding_result == (deployment,) + assert isinstance(chat_result, UnsupportedEndpointError) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6f6eae1f98e..abab8ad19eb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -49940,6 +49940,11 @@ export interface components { discoverable?: boolean | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** + * Enforce Supported Endpoints + * @default false + */ + enforce_supported_endpoints: boolean; /** Id */ id: string | null; /** Input Cost Per Character */ From 85011cdf9f23d29909f9b191041325f1e68b73b7 Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 17:49:56 -0700 Subject: [PATCH 15/21] fix(router): skip unresolved local rate-limit fallbacks --- litellm/proxy/route_llm_request.py | 3 +-- litellm/router.py | 4 ++-- .../routing/test_non_chat_model_endpoint_gate.py | 14 +++++++++++--- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index b1d0c72f97c..e5917780147 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -11,7 +11,6 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request -from litellm.router_utils.endpoint_eligibility import UnsupportedEndpointError # Client-supplied params that make the router or the call path fabricate a # failure or a delay instead of calling the provider. The ``mock_testing_*`` @@ -362,7 +361,7 @@ def endpoint_unsupported_by_model( data: Mapping[str, object], llm_router: LitellmRouter | None, team_id: str | None, -) -> UnsupportedEndpointError | None: +) -> litellm.BadRequestError | None: model: Final = data.get("model") if not isinstance(model, str) or llm_router is None: return None diff --git a/litellm/router.py b/litellm/router.py index 64b5ec216ec..374c881b3e3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13869,11 +13869,11 @@ class Router: def endpoint_error_for_request( self, model: str, call_type: str, request_kwargs: Mapping[str, object] - ) -> UnsupportedEndpointError | None: + ) -> litellm.BadRequestError | None: try: with endpoint_selection(call_type, model=model): self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) - except UnsupportedEndpointError as error: + except litellm.BadRequestError as error: return error return None diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 89a9d1c14f0..fa582b0432c 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1360,17 +1360,25 @@ def test_a_mixed_group_skips_its_restricted_embedding_deployment(rig: _Rig) -> N assert [JSON_OBJECT.validate_json(call.body)["model"] for call in calls] == [_CHAT_MODEL] -def test_local_rate_limit_with_only_unsupported_fallbacks_keeps_the_original_429(rig: _Rig) -> None: +@pytest.mark.parametrize( + "fallbacks, expected_status", + [(("embedding_only",), 429), (("missing",), 429), (("missing", "chat_target"), 200)], +) +def test_local_rate_limit_skips_unusable_fallbacks(rig: _Rig, fallbacks: tuple[str, ...], expected_status: int) -> None: limited: Final = rig.names["chat_rate_limited"] + targets: Final = tuple(rig.names[name] if name != "missing" else _marker() for name in fallbacks) with rig.gateway.scenario() as scenario: key: Final = scenario.key( - model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: [rig.names["embedding_only"]]}]} + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: list(targets)}]} ) first: Final = _chat(rig, limited, _marker(), key=key) assert first.status_code == 200, first.text marker: Final = _marker() second: Final = _chat(rig, limited, marker, key=key) - assert second.status_code == 429, second.text + assert second.status_code == expected_status, second.text + if expected_status == 200: + _assert_served(rig, second, marker, "/v1/chat/completions") + return assert "Supported endpoints" not in second.text, second.text assert rig.upstream.calls(marker) == () From e32ef13c0e3e4ee686d3487c24da9b600a4d9aaa Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 17:55:42 -0700 Subject: [PATCH 16/21] fix(router): ignore caller-provided endpoint selector values --- litellm/router.py | 5 +++-- .../routing/test_non_chat_model_endpoint_gate.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 374c881b3e3..6498cf88ce8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8686,8 +8686,9 @@ class Router: Handler for making a call to the .completion()/.embeddings()/etc. functions. """ model_group: Final = kwargs.get("model") - endpoint_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python( - kwargs.get("original_generic_function") or original_function + generic_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("original_generic_function")) + endpoint_function: Final = ( + generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) ) with endpoint_selection( getattr(endpoint_function, "__name__", None), diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index fa582b0432c..3b2e27d69a3 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1410,3 +1410,18 @@ def test_responses_websocket_uses_a_chat_fallback_without_calling_the_embedding_ (forwarded,) = responses_ws_peer.drain(peer.frames) assert marker in json.dumps(forwarded), forwarded assert rig.upstream.calls(marker) == () + + +def test_client_routing_fields_cannot_disable_an_explicit_restriction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy(scenario, rig.api_base, "restricted-chat", _single(_EMBED_ONLY)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat( + rig, + model, + marker, + original_generic_function="caller-value", + additional_drop_params=["original_generic_function"], + ) + _assert_rejected_chat(rig, response, marker, model) From adba6f48ab32c731fde6d76c3a50abb6c200257f Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 18:02:27 -0700 Subject: [PATCH 17/21] fix(router): normalize declared chat endpoint prefixes --- litellm/router_utils/endpoint_eligibility.py | 11 +++++++++-- .../test_embedding_only_supported_endpoints_e2e.py | 2 +- .../routing/test_non_chat_model_endpoint_gate.py | 14 ++++++++++++++ .../unit/router_utils/test_endpoint_eligibility.py | 11 +++++++++++ 4 files changed, 35 insertions(+), 3 deletions(-) diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py index 85db045d6a5..112ccaf535a 100644 --- a/litellm/router_utils/endpoint_eligibility.py +++ b/litellm/router_utils/endpoint_eligibility.py @@ -28,7 +28,13 @@ CHAT_ENDPOINT_BY_CALL: Final[Mapping[str, str]] = MappingProxyType( "acompact_responses": "/v1/responses/compact", } ) -CHAT_ENDPOINTS: Final = frozenset((*CHAT_ENDPOINT_BY_CALL.values(), "/chat/completions", "/completions", "/responses")) + + +def _endpoint_path(endpoint: str) -> str: + return endpoint.removeprefix("/openai").removeprefix("/v1beta").removeprefix("/v1") + + +CHAT_ENDPOINTS: Final = frozenset(_endpoint_path(endpoint) for endpoint in CHAT_ENDPOINT_BY_CALL.values()) _Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object]) _REQUEST_SELECTION: Final[ContextVar[tuple[str | None, str | None]]] = ContextVar( "router_request_endpoint", default=(None, None) @@ -72,7 +78,8 @@ def non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | Non if not declaration.enforce_supported_endpoints or declaration.supported_endpoints is None: return None endpoints: Final = frozenset(declaration.supported_endpoints) - return endpoints if endpoints.isdisjoint(CHAT_ENDPOINTS) else None + normalized: Final = frozenset(_endpoint_path(endpoint) for endpoint in endpoints) + return endpoints if normalized.isdisjoint(CHAT_ENDPOINTS) else None def eligible_deployments( diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py index 1f3ba3d05db..899b421b218 100644 --- a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -88,7 +88,7 @@ class TestEmbeddingOnlySupportedEndpoints: ) @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") - @meta(_subject(Route.CHAT_COMPLETIONS)) + @meta(dataclasses.replace(_subject(Route.CHAT_COMPLETIONS), models=(EMBEDDING_BACKEND, CHAT_FALLBACK_BACKEND))) def test_chat_completions_with_a_fallback_are_answered_by_the_fallback( self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients ) -> None: diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 3b2e27d69a3..97db60224b8 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1425,3 +1425,17 @@ def test_client_routing_fields_cannot_disable_an_explicit_restriction(rig: _Rig) additional_drop_params=["original_generic_function"], ) _assert_rejected_chat(rig, response, marker, model) + + +def test_a_compaction_endpoint_declaration_keeps_the_model_available_for_compaction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy( + scenario, + rig.api_base, + "compact-declaration", + _single({"supported_endpoints": ["/responses/compact"], "enforce_supported_endpoints": True}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = rig.gateway.request("POST", "/responses/compact", {"model": model, "input": marker}) + _assert_served(rig, response, marker, "/v1/responses/compact") diff --git a/tests/unit/router_utils/test_endpoint_eligibility.py b/tests/unit/router_utils/test_endpoint_eligibility.py index 82d161743bb..e129bf69b5c 100644 --- a/tests/unit/router_utils/test_endpoint_eligibility.py +++ b/tests/unit/router_utils/test_endpoint_eligibility.py @@ -74,3 +74,14 @@ def test_nested_non_chat_calls_do_not_inherit_the_outer_chat_restriction() -> No chat_result: Final = eligible_deployments("shared", (deployment,)) assert embedding_result == (deployment,) assert isinstance(chat_result, UnsupportedEndpointError) + + +@pytest.mark.parametrize( + "endpoint", + ["/messages", "/interactions", "/responses/compact", "/openai/v1/responses/compact", "/openai/v1/chat/completions"], +) +def test_chat_endpoint_declarations_keep_their_meaning_across_prefixes(endpoint: str) -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": [endpoint]}) + with endpoint_selection("acompletion"): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) From 6456dfc372870bd2361a98fb540a1e763df877e2 Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 18:10:26 -0700 Subject: [PATCH 18/21] test(router): cover preflight error-value behavior --- tests/unit/test_router/test_router.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 50a9d98385d..eaeaf7e7b1f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -11285,6 +11285,24 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) +def test_endpoint_preflight_returns_unavailable_model_errors_as_values() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "available", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "unit-provider-key"}, + } + ], + num_retries=0, + ) + missing: Final = "unregistered-fallback" + error: Final = router.endpoint_error_for_request(missing, "acompletion", {}) + assert isinstance(error, litellm.BadRequestError) + assert error.status_code == 400 + assert error.model == missing + assert router.endpoint_error_for_request("available", "acompletion", {}) is None + + def test_routing_groups_legacy_property_forwards_to_public_attribute(): router = _router_with_two_deployments([False, False]) From 7e527c39e705d2890f7bb3c1dd72daf4409f0e6c Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 18:24:23 -0700 Subject: [PATCH 19/21] test(router): derive replay metadata defaults from model info --- tests/unit/test_router_model_cost_isolation.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d8625ff9da3..6c01785771d 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -2206,10 +2206,8 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: f"the rebuild changed or dropped a field the boot registration wrote: " f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }" ) - assert {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} == { - "db_model": False, - "member_auto_router": False, - } + replay_defaults: Final = {field: rebuilt[field] for field in set(rebuilt) - set(at_boot)} + assert replay_defaults.items() <= ModelInfo().model_dump(exclude_none=True).items() assert "member_auto_router" not in litellm.model_cost["gpt-4o"] assert router.model_list finally: From 2ba59fb46055ddb8c157a7f78d310f6ae87fdf0f Mon Sep 17 00:00:00 2001 From: Mateo Date: Fri, 16 Oct 2026 18:44:07 -0700 Subject: [PATCH 20/21] fix(router): preserve fallback aliases and check synchronous text calls --- litellm/proxy/common_request_processing.py | 22 +++++++++- litellm/router.py | 13 +++--- .../test_non_chat_model_endpoint_gate.py | 44 +++++++++++++++++++ tests/unit/test_router/test_router.py | 33 ++++++++++++++ 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 45c98a6c0db..d4f1457c1bc 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2519,6 +2519,17 @@ class ProxyBaseLLMRequestProcessing: original_model, fallback_models, ) + raw_router_settings: Final = _DECODED_JSON.validate_python( + _REQUEST_FIELDS.validate_python(rate_limited_data).get("router_settings_override") + ) + group_alias: Final = _DECODED_JSON.validate_python( + _REQUEST_FIELDS.validate_python(raw_router_settings).get("model_group_alias") + if isinstance(raw_router_settings, Mapping) + else None + ) + key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime + object, user_api_key_dict.aliases + ) try: for fallback_model in fallback_models: if fallback_model == original_model: @@ -2534,8 +2545,17 @@ class ProxyBaseLLMRequestProcessing: self.data = independent_snapshot(pristine) self.data["model"] = fallback_model try: + probe_data = { + **_REQUEST_FIELDS.validate_python(pristine), + "model": resolve_served_model( + fallback_model, + _DECODED_JSON.validate_python(litellm.model_alias_map), + key_aliases, + group_alias, + ), + } if ( - endpoint_unsupported_by_model(route_type, self.data, llm_router, user_api_key_dict.team_id) + endpoint_unsupported_by_model(route_type, probe_data, llm_router, user_api_key_dict.team_id) is not None ): continue diff --git a/litellm/router.py b/litellm/router.py index 6498cf88ce8..3a0058b0f92 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5534,12 +5534,13 @@ class Router: kwargs.setdefault("metadata", {}).update({"model_group": model}) # pick the one that is available (lowest TPM/RPM) - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with endpoint_selection("text_completion", model=model): + deployment: Final = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) data: Final = deployment["litellm_params"].copy() for k, v in self.default_litellm_params.items(): diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index 97db60224b8..c78eb7bddad 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -1439,3 +1439,47 @@ def test_a_compaction_endpoint_declaration_keeps_the_model_available_for_compact _served_by_every_worker(rig.gateway, frozenset(identities)) response: Final = rig.gateway.request("POST", "/responses/compact", {"model": model, "input": marker}) _assert_served(rig, response, marker, "/v1/responses/compact") + + +@pytest.mark.parametrize("alias_source", ["key_alias", "key_router", "team_router"]) +def test_local_rate_limit_fallbacks_resolve_saved_aliases(rig: _Rig, tmp_path: Path, alias_source: str) -> None: + suffix: Final = uuid.uuid4().hex[:12] + names: Final = {label: f"gate-fallback-alias-{label}-{suffix}" for label in ("primary", "target", "alias")} + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["primary"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(names["target"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "fallbacks": [{names["primary"]: [names["alias"]]}], + }, + } + path: Final = tmp_path / "fallback-alias.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + aliases: Final = {names["alias"]: names["target"]} + with ( + owned_proxy_process(rig.gateway, tmp_path, {}, config=path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + team: Final = ( + scenario.team(router_settings={"model_group_alias": aliases}) if alias_source == "team_router" else None + ) + key: Final = scenario.key( + team_id=team, + aliases=aliases if alias_source == "key_alias" else {}, + router_settings={"model_group_alias": aliases} if alias_source == "key_router" else None, + model_rpm_limit={names["primary"]: 1}, + ) + first: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], _marker()), key=key + ) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], marker), key=key + ) + _assert_fallback_answered(rig.upstream, second, marker) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index eaeaf7e7b1f..9c93ad3c200 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -11285,6 +11285,39 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) +def test_synchronous_text_completion_rejects_a_restricted_model_before_the_provider() -> None: + model: Final = "text-completion-openai/unit-text-model" + router: Final = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": model, + "api_base": "https://sync-router.invalid/v1", + "api_key": "unit-provider-key", + }, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + } + ], + num_retries=0, + ) + with respx.mock(assert_all_called=False) as upstream: + provider: Final = upstream.post("https://sync-router.invalid/v1/completions").respond( + 200, + json={ + "id": "cmpl-unit", + "object": "text_completion", + "created": 1, + "model": "unit-text-model", + "choices": [{"text": "provider called", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + with pytest.raises(litellm.BadRequestError, match="/v1/completions is not supported"): + router.text_completion(model=model, prompt="test") + assert provider.call_count == 0 + + def test_endpoint_preflight_returns_unavailable_model_errors_as_values() -> None: router: Final = Router( model_list=[ From 06409e2f6ef447d9fa3f6df1e848fd9f5ff8fad9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 13:54:48 -0700 Subject: [PATCH 21/21] test(routing): rewrite the WebSocket model with a pre-call callback, since main no longer runs guardrails at socket admission --- .../test_non_chat_model_endpoint_gate.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py index c78eb7bddad..346cc60bc26 100644 --- a/tests/integration/routing/test_non_chat_model_endpoint_gate.py +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -972,14 +972,17 @@ def test_model_group_alias_to_an_embedding_only_group_keeps_serving_embeddings(r assert routed.upstream.targets(marker) == ("/v1/embeddings",) -_REWRITER: Final = """from litellm.integrations.custom_guardrail import CustomGuardrail +_REWRITER: Final = """from litellm.integrations.custom_logger import CustomLogger -class WebSocketModelRewriter(CustomGuardrail): +class WebSocketModelRewriter(CustomLogger): async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): if call_type == "_aresponses_websocket" and data.get("model") == "__SOURCE__": data["model"] = "__TARGET__" return data + + +rewriter = WebSocketModelRewriter() """ @@ -1000,16 +1003,7 @@ def _rewritten_config(directory: Path, api_base: str, source: str, target: str, _routed_entry(target, api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), ], "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, - "guardrails": [ - { - "guardrail_name": f"{module}-guardrail", - "litellm_params": { - "guardrail": f"{module}.WebSocketModelRewriter", - "mode": "pre_call", - "default_on": True, - }, - } - ], + "litellm_settings": {**object_value(shared["litellm_settings"]), "callbacks": [f"{module}.rewriter"]}, } path: Final = directory / "gate-rewritten-config.yaml" path.write_text(yaml.safe_dump(config, sort_keys=False))