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] 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