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):