diff --git a/litellm/constants.py b/litellm/constants.py index f7ba55d8ce4..9f29baac562 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -969,7 +969,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = { "context_management": None, } -openai_compatible_endpoints: Final[list] = [ +openai_compatible_endpoints: Final[list[str]] = [ "api.perplexity.ai", "api.endpoints.anyscale.com/v1", "api.deepinfra.com/v1/openai", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 74d9344a172..a67c1bb7df8 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -61,6 +61,10 @@ def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool: return url_path == endpoint_path or url_path.startswith(endpoint_path + "/") +def is_openai_compatible_endpoint(api_base: str) -> bool: + return any(_endpoint_matches_api_base(endpoint, api_base) for endpoint in litellm.openai_compatible_endpoints) + + def is_non_openai_azure_model(model: str) -> bool: try: model_name: Final = model.split("/", 1)[1] diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f9225dc1a79..f566c8f4e02 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -59,8 +59,8 @@ if TYPE_CHECKING: else: ResponseText = str # Fallback for ResponseText import from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.get_llm_provider_logic import is_openai_compatible_endpoint from litellm.llms.openai.data_residency import infer_openai_data_residency -from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams @@ -2461,6 +2461,13 @@ def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ) +def _websocket_custom_llm_provider(custom_llm_provider: str | None, api_base: str | None) -> str | None: + is_openai_on_known_provider_host: Final = ( + custom_llm_provider == "openai" and api_base is not None and is_openai_compatible_endpoint(api_base) + ) + return None if is_openai_on_known_provider_host else custom_llm_provider + + @client async def _aresponses_websocket( model: str, @@ -2491,6 +2498,7 @@ async def _aresponses_websocket( dynamic_api_base, ) = litellm.get_llm_provider( model=model, + custom_llm_provider=_websocket_custom_llm_provider(litellm_params.custom_llm_provider, api_base), api_base=api_base, api_key=api_key, litellm_params=litellm_params, @@ -2512,22 +2520,18 @@ async def _aresponses_websocket( ) resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None - responses_api_provider_config: BaseResponsesAPIConfig | None = None - if _custom_llm_provider is not None: - responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( + responses_api_provider_config: Final = ( + None + if _custom_llm_provider is None + else ProviderConfigManager.get_provider_responses_api_config( model=resolved_model, provider=litellm.LlmProviders(_custom_llm_provider), api_base=resolved_api_base, ) - - resolved_api_key: Final = ( - dynamic_api_key - or litellm_params.api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") ) + resolved_api_key: Final = api_key or dynamic_api_key + # Extract params that we're passing explicitly to avoid duplicates in **kwargs _explicit_keys: Final = { "user_api_key_dict", diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index 18be6bb9abb..27d328369d3 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -27,14 +27,23 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) PROVIDER_MODEL: Final = "ws-peer-model" USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) +PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key" +PROXY_AZURE_KEY: Final = "proxy-environment-azure-key" _Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] +@dataclass(frozen=True, slots=True) +class Credentials: + authorization: str | None + api_key: str | None + + @dataclass(frozen=True, slots=True) class Peer: url: str paths: SimpleQueue[str] + credentials: SimpleQueue[Credentials] frames: SimpleQueue[dict[str, JsonValue]] @@ -75,10 +84,17 @@ def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]: async def _answer( connection: ServerConnection, paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], frames: SimpleQueue[dict[str, JsonValue]], script: _Script, ) -> None: - paths.put(connection.request.path if connection.request is not None else "") + request: Final = connection.request + paths.put(request.path if request is not None else "") + credentials.put( + Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) + if request is not None + else Credentials(None, None) + ) turns: Final = itertools.count(1) async for raw in connection: frame: Final = json.loads(raw) @@ -92,12 +108,15 @@ async def _answer( async def _serve( tls: ssl.SSLContext, paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], frames: SimpleQueue[dict[str, JsonValue]], ports: SimpleQueue[int], stop: asyncio.Event, script: _Script, ) -> None: - async with serve(lambda connection: _answer(connection, paths, frames, script), "127.0.0.1", 0, ssl=tls) as server: + async with serve( + lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls + ) as server: ports.put(next(iter(server.sockets)).getsockname()[1]) await stop.wait() @@ -107,16 +126,17 @@ def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterato loop: Final = asyncio.new_event_loop() stop: Final = asyncio.Event() paths: Final = SimpleQueue[str]() + credentials: Final = SimpleQueue[Credentials]() frames: Final = SimpleQueue[dict[str, JsonValue]]() ports: Final = SimpleQueue[int]() thread: Final = threading.Thread( target=loop.run_until_complete, - args=(_serve(server_context(*cert), paths, frames, ports, stop, script),), + args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), daemon=True, ) thread.start() try: - yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, frames) + yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) finally: loop.call_soon_threadsafe(stop.set) thread.join(timeout=10) @@ -180,7 +200,11 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("responses-ws"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy( + base, + tmp_path_factory.mktemp("responses-ws"), + {"SSL_CERT_FILE": str(cert[0]), "OPENAI_API_KEY": PROXY_OPENAI_KEY, "AZURE_API_KEY": PROXY_AZURE_KEY}, + ) as proxy: yield proxy @@ -253,3 +277,67 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a error: Final = object_value(object_value(row["metadata"])["error_information"]) assert str(error["error_code"]) == "429", error assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error + + +def _one_turn(candidate: Gateway, key: str, model: str) -> tuple[dict[str, JsonValue], ...]: + (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) + return events + + +def test_a_bare_model_declared_as_openai_sends_its_deployment_key_to_the_provider_socket( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + deployment_key: Final = f"sk-deployment-{uuid.uuid4().hex}" + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=PROVIDER_MODEL, custom_llm_provider="openai", api_base=peer.url, api_key=deployment_key + ) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + assert peer.credentials.get_nowait().authorization == f"Bearer {deployment_key}" + + +def test_an_openai_deployment_without_its_own_key_uses_the_proxy_openai_key( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url, api_key=None) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + assert peer.credentials.get_nowait().authorization == f"Bearer {PROXY_OPENAI_KEY}" + + +def test_an_azure_deployment_without_its_own_key_gets_the_azure_key_and_never_the_openai_key( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{PROVIDER_MODEL}", api_base=peer.url.removesuffix("/v1"), api_key=None + ) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + credentials: Final = peer.credentials.get_nowait() + assert credentials.api_key == PROXY_AZURE_KEY + assert PROXY_OPENAI_KEY not in (credentials.authorization or "") + + +def test_a_declared_openai_provider_wins_over_another_providers_model_prefix( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + gateway_model: Final = f"anthropic/{PROVIDER_MODEL}" + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model(model=gateway_model, custom_llm_provider="openai", api_base=peer.url) + key: Final = scenario.key(models=[model]) + events: Final = _one_turn(candidate, key, model) + assert _completed(events)["status"] == "completed" + (forwarded,) = _drain(peer.frames) + assert forwarded["model"] == gateway_model + (row,) = eventually( + lambda: read_rows( + 'SELECT custom_llm_provider FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (hashlib.sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert row["custom_llm_provider"] == "openai", row diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 9f9deb2ca0f..3625bbb7c04 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,6 +8,7 @@ Tests that: """ import json +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock @@ -2805,6 +2806,178 @@ class TestWebSocketChunkTypes: assert "Complete reasoning" in serialized +class TestResponsesWebSocketCredentials: + @staticmethod + async def _upstream_handshake_headers(**call_kwargs: object) -> Mapping[str, str]: + from asyncio import Future, get_running_loop + from types import TracebackType + from unittest.mock import AsyncMock, patch + + from litellm.responses import main as responses_main + + captured_headers: Final[Future[Mapping[str, str]]] = get_running_loop().create_future() + + class FakeConnect: + def __init__( + self, + _url: str, + *, + additional_headers: Mapping[str, str], + **_kwargs: object, + ) -> None: + captured_headers.set_result(dict(additional_headers)) + + async def __aenter__(self) -> None: + raise RuntimeError("stop at upstream handshake") + + async def __aexit__( + self, + _exc_type: type[BaseException] | None, + _exc: BaseException | None, + _tb: TracebackType | None, + ) -> None: + return None + + websocket: Final = MagicMock() + websocket.close = AsyncMock() + with patch("websockets.connect", FakeConnect): + await responses_main._aresponses_websocket.__wrapped__( + websocket=websocket, + litellm_logging_obj=MagicMock(), + **call_kwargs, + ) + + assert captured_headers.done(), "the request never reached the native upstream websocket" + return captured_headers.result() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "api_base,explicit_key,global_key,expected_key", + [ + ("https://upstream.invalid/v1", "deployment-key", None, "deployment-key"), + ("https://upstream.invalid/v1", "deployment-key", "global-key", "deployment-key"), + ("https://upstream.invalid/v1", None, "global-key", "global-key"), + ], + ) + async def test_openai_websocket_authorization( + self, + monkeypatch: pytest.MonkeyPatch, + api_base: str, + explicit_key: str | None, + global_key: str | None, + expected_key: str, + ) -> None: + import litellm + + monkeypatch.setattr(litellm, "api_key", global_key) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("LITELLM_RUST", "false") + + headers: Final = await self._upstream_handshake_headers( + model="gpt-6-astra", + custom_llm_provider="openai", + api_key=explicit_key, + api_base=api_base, + ) + + assert headers["Authorization"] == f"Bearer {expected_key}" + + @pytest.mark.asyncio + async def test_azure_websocket_without_its_own_key_never_gets_the_openai_key( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import litellm + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "azure_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "openai-global-key") + monkeypatch.setenv("AZURE_API_KEY", "azure-env-key") + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + monkeypatch.setenv("LITELLM_RUST", "false") + + headers: Final = await self._upstream_handshake_headers( + model="gpt-6-astra", + custom_llm_provider="azure", + api_base="https://resource.openai.azure.com", + ) + + assert headers["api-key"] == "azure-env-key" + assert "openai-global-key" not in " ".join(headers.values()) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider,api_base,deployment_key,expected_key", + [ + (None, "https://api.deepseek.com/v1", "deployment-key", "deployment-key"), + ("openai", "https://api.deepseek.com/v1", "deployment-key", "deployment-key"), + ("openai", "https://api.deepseek.com/v1", None, "deepseek-env-key"), + ("openai", "https://api.mistral.ai/v1", "deployment-key", "deployment-key"), + ], + ) + async def test_deployment_on_a_known_provider_host_completes_through_that_provider( + self, + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: str | None, + api_base: str, + deployment_key: str | None, + expected_key: str, + ) -> None: + from unittest.mock import AsyncMock, patch + + import httpx + import respx + + import litellm + from litellm.responses import main as responses_main + + chunks: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + stream: Final = "".join( + f"data: {json.dumps({'id': 'chatcmpl-1', 'object': 'chat.completion.chunk', 'created': 1, 'model': 'astra-chat', **chunk})}\n\n" + for chunk in chunks + ) + + def refuse_native_socket(*_args: object, **_kwargs: object) -> None: + raise ConnectionRefusedError("this host has no Responses WebSocket") + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "openai-env-key") + monkeypatch.setenv("DEEPSEEK_API_KEY", "deepseek-env-key") + monkeypatch.setenv("MISTRAL_API_KEY", "mistral-env-key") + monkeypatch.setenv("LITELLM_RUST", "false") + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.close = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) + + with respx.mock(assert_all_called=False) as upstream, patch("websockets.connect", refuse_native_socket): + chat_completions: Final = upstream.post(f"{api_base}/chat/completions").mock( + return_value=httpx.Response( + 200, text=stream + "data: [DONE]\n\n", headers={"content-type": "text/event-stream"} + ) + ) + await responses_main._aresponses_websocket.__wrapped__( + model="astra-chat", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=deployment_key, + first_message=json.dumps({"type": "response.create", "model": "astra-chat", "input": "Reply OK"}), + websocket=websocket, + litellm_logging_obj=MagicMock(), + ) + + sent_types: Final = [json.loads(call.args[0]).get("type") for call in websocket.send_text.await_args_list] + assert chat_completions.called, f"the request never reached {api_base}/chat/completions, client saw {sent_types}" + assert chat_completions.calls.last.request.headers["Authorization"] == f"Bearer {expected_key}" + assert "response.completed" in sent_types + + class TestNativeWebSocketUrlConstruction: """Test that native WebSocket URLs include the model query parameter.