mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
* feat(anthropic): workload identity federation and pluggable identity sources Backend half of #38818 (internal copy of the fork PR #38013), rebuilt as one commit on top of litellm_internal_staging without the dashboard changes. Deployments on anthropic/ without a static api_key can exchange an OIDC workload assertion for a short-lived sk-ant-oat01 token through a shared RFC 7523 JWT-bearer engine. The assertion comes from a mounted token file, an env token, a LiteLLM-signed issuer, or Keycloak, chosen per deployment, per named credential, or through ANTHROPIC_IDENTITY_SOURCE. The federation fields are server-owned: refused inline in request bodies and on POST /model/new, proxy-admin only on credentials, and the token exchange is pinned to api.anthropic.com unless LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS adds a host. GET /credentials/{name}/jwks exports the public key set of a LiteLLM-signed credential for the Claude Console. The OpenAI federation trio from #39613 rides along on the backend side with the same server-owned handling. Fixes #28607 Resolves LIT-6107 Co-authored-by: derhornspieler <15236687+derhornspieler@users.noreply.github.com> * fix(anthropic): let batch-result downloads mint from deployment params and accept host:port allowlist entries The files handler enabled workload identity on batch-result downloads but never received the deployment's litellm_params, so a deployment authenticating through a named credential could only mint from process-wide env vars. It now threads litellm_params through to the auth header the way the batch retrieve path already does. LITELLM_ANTHROPIC_WIF_ALLOWED_HOSTS entries written as host:port were read by urlsplit as a scheme, so the allowlist kept the raw entry while the exchange compared bare hostnames and refused the gateway. Entries are now parsed as network locations whether or not they carry a scheme. * fix(types): move the WIF kwargs key sets to a leaf module so the kwargs funnel imports without a cycle * test(anthropic): pin case-insensitive matching of WIF exchange-host allowlist entries * fix(anthropic): end workload identity federation errors without a period so the router suffix reads cleanly * fix(proxy): decrypt stored litellm_params before the WIF write gate * fix(proxy): hide WIF secret references from /health output * fix(proxy): keep the proxy error shape on credential endpoint refusals * fix(proxy): hide identity token file paths from /health output * feat(openai): honor per-deployment workload identity on non-chat routes Image generation, audio speech and transcription, moderation, text completions, files, batches, and fine-tuning built their OpenAI client from the env vars alone, so a keyless deployment carrying openai_identity_provider_id, openai_service_account_id, and openai_identity_token_file failed with the api_key error on every route other than chat, Responses, and embeddings. Those handlers now build their client through the shared workload identity resolver, and the managed batch file-access credential extraction keeps the trio * fix(openai): cache workload identity clients and let deployment identity outrank process-wide keys Non-chat handlers built a fresh SDK client per call, so files, batches, fine-tuning, text completion and sync moderation exchanged the OIDC token on every request. The builders now reuse the cached workload identity client per identity, base, timeout, retries and organization; static-key clients are still built fresh. A process-wide OPENAI_API_KEY (or litellm.api_key / litellm.openai_key) no longer silently overrides a deployment's openai_identity_* params. An explicit distinct api_key on the deployment or call still wins. * fix(openai): keep shared workload identity clients on the process http session The workload identity client is cached and shared across requests, so it must never borrow the transport a single caller passed in. Static-key clients keep using the caller's http client; federated clients always ride the process-wide session. Also makes the http client helpers on BaseOpenAILLM public since the builder now needs them from outside the class. * test(openai): pin partial deployment identity beating the process key * test(openai): type the managed-object invoke callback * fix(anthropic): rename the federation workspace param so Bedrock's anthropic_workspace_id keeps working The Bedrock Claude Platform route already reads anthropic_workspace_id from optional_params, so banning that spelling as a server-owned federation parameter broke a pre-existing client capability. The federation field is now anthropic_federation_workspace_id (env ANTHROPIC_FEDERATION_WORKSPACE_ID), which restores the base branch's behavior for Bedrock callers, drops the Bedrock-specific hint from the refusal message, and deletes the unconditional ban constant that no longer had a reader * fix(auth): share one exchanged token across workers reading the same assertion Anthropic accepts each identity assertion exactly once, so two uvicorn workers reading the same token file both minting from it means the second exchange is denied with jti_reused. Minted tokens now land in a per-user 0700 cache directory guarded by a file lock, so workers on the same host reuse one exchange until the token expires or the assertion rotates. A 401 is only retried when the re-read assertion actually differs, and the denial hint explains jti_reused. LITELLM_TOKEN_EXCHANGE_CACHE_DIR moves the cache and an empty value disables it * feat(openai): honor deployment workload identity on image edits, vector stores, containers, videos, and evals The httpx-based OpenAI handlers built their Authorization header from the deployment api_key or OPENAI_API_KEY alone, so a keyless deployment carrying openai_identity_provider_id, openai_service_account_id, and openai_identity_token_file sent "Bearer None". They now resolve the bearer through resolve_openai_bearer_token, which exchanges the deployment's identity token the same way the SDK-based routes do. A deployment that carries only part of the identity trio now raises a 500 naming the missing field instead of silently falling back to the process key, and the SDK client builders take a factory for the static-key http client so the federated branch never touches litellm.client_session. * fix: keep anthropic federation from being shadowed or leaked An empty or whitespace-only ANTHROPIC_API_KEY counted as set, so a federated deployment sent an empty x-api-key on every call instead of minting a token. Blank values now read as unset, and a real static key on a federated deployment logs once that it outranks federation and nothing is being federated. The exchange-host allowlist matched hostnames only, so a second process on another port of an allowed host was trusted with the workload's identity token. An entry that names a port now trusts that port alone, while a bare host still trusts every port. The shared token store exists so the workers reading one projected token file do not each spend its single-use jti. A source that mints its own assertion per exchange shares nothing with another worker, so it no longer writes a live token to disk for a lookup that can never hit. * feat(openai): resolve per-deployment workload identity for the assistants routes OpenAIAssistantsAPI built its OpenAI and AsyncOpenAI clients from a locals() scrape, so a deployment carrying openai_identity_provider_id, openai_service_account_id and openai_identity_token_file never exchanged a token. Thread litellm_params through the class and build both clients with the shared workload identity builders. * fix: unlink a staged token file a failed write leaves behind The 401 denial hint now also says federation ignores ANTHROPIC_WORKSPACE_ID, which the Bedrock Claude platform provider already reads. * refactor: move anthropic jwks derivation behind a provider-owned tagged union * feat(openai): resolve per-deployment workload identity for image variations Build the image variations client through the shared workload identity builders instead of scraping locals(), and give the async path the async http session it was missing. Let a federation trio serve the model-less batch fallback, and hold the standard LRU size of exchanged tokens. * fix: unlink the staged token file when its write fails at close A buffered write only reaches the disk when the handle closes, so a full disk surfaces at close and left the staging file behind holding a usable token. * fix(anthropic): close the staging descriptor before writing the shared token file * fix(wif): judge federation writes by what they set, not what is stored The admin gate read the stored deployment, so a team admin lost edit, delete and Test Connection on any deployment carrying federation params. It now returns early unless the submitted fields touch the federation surface, and a Test Connection probe that points the deployment at its own api_base is still refused, with the 403 no longer wrapped into a 500 The rest of the same review pass: POST /model/new refuses only a blocking value of `blocked`, so a client that always sends `blocked: false` is not turned away; a request body can no longer pick which federated identity to mint as by naming a stored credential; an advisory refresh the executor refuses disarms the entry instead of wedging the identity until the follower timeout; the static-key shadow warning resolves its env fallback inside the cache instead of once per request; credential writes drop nulls before storing them; the token exchange validates the endpoint URL before reading an assertion and keeps refusing redirects across a client heal; /health hides every server-owned federation field from non-admins; and the async create_file and create_batch paths say which setting is missing when the provider resolves no URL * style(tests): keep two blank lines before the batch_cost_is_final block * test(batches): type the workload identity fallback test fixtures * fix(moderation): stop validating unrelated kwargs on the sync workload identity path * fix(auth): refuse server-owned workload identity keys in array form and inside a client router config * fix(openai): keep a deployment's own key ahead of workload identity on every route Thread the caller's own api_key into the identity resolver for embedding, transcription, moderation, image generation, and assistants so a deployment key equal to the process key still wins. A partial identity trio now falls back to a configured static key, unset max_retries keeps the SDK default, and /v1/decisions plus OpenAI count_tokens resolve the exchanged bearer for identity-only deployments * fix(containers): pass deployment params to container file auth headers Container file routes built their auth headers without the deployment's litellm_params, so a keyless deployment with an OpenAI workload identity sent Bearer None on upload, list, retrieve, delete and content calls * fix(openai): keep process-wide key precedence and move token exchange off the event loop Restore the merge base order where an explicit or process-wide OpenAI key wins over federation, so exporting OPENAI_API_KEY keeps working the way it did before. Run the header-building validate_environment calls (image edit, videos, containers, vector stores, vector store files) and the decisions and System One request preparation in a worker thread, so a slow token exchange no longer stalls the loop. Federate realtime websocket sessions, live sessions, realtime client secrets and transcription sessions, and the OpenAI moderation guardrail when no key is set. * fix(openai): answer 401 for a partial identity trio and federate realtime health checks A deployment with part of the workload identity trio and no static key now fails as an auth error naming the missing params, so the router does not retry it as a server error. The realtime health check sends the exchanged bearer for a keyless federated OpenAI deployment, and API token counting reads the same static keys as every other route before minting. * fix(openai): keep local token counting without an identity trio and federate the moderation guardrail from its own config * fix(openai): skip a partial-trio moderation guardrail and let vault keys win in token counting * test(integration): refuse caller-chosen identity params in user_config and bracketed form keys * fix(openai): send the honored secret manager key from SDK client builders A deployment whose OPENAI_API_KEY lives only in the secret manager now sends that key on image, assistants, files, batches, and other SDK-built routes instead of failing with no api_key, matching what the workload identity resolver already treats as a static key * test(openai): check the token exchange runs off the event loop instead of timing the loop --------- Co-authored-by: derhornspieler <15236687+derhornspieler@users.noreply.github.com>
1178 lines
45 KiB
Python
1178 lines
45 KiB
Python
import asyncio
|
|
import json
|
|
import time
|
|
from collections.abc import Awaitable, Callable, Mapping
|
|
from pathlib import Path
|
|
from types import TracebackType
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
import litellm
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
from litellm.llms.openai.workload_identity import _workload_identity_auth
|
|
from litellm.models.credentials import CredentialItem
|
|
from litellm.realtime_api import main as realtime_main
|
|
from litellm.realtime_api.main import _with_resolved_session_model
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
|
|
@pytest.fixture
|
|
def local_model_cost_map(monkeypatch):
|
|
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
|
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
|
litellm.get_model_info.cache_clear()
|
|
yield
|
|
litellm.get_model_info.cache_clear()
|
|
|
|
|
|
class FakeLogging:
|
|
def update_from_kwargs(self, **kwargs):
|
|
pass
|
|
|
|
def pre_call(self, **kwargs):
|
|
pass
|
|
|
|
|
|
class _CallCapture:
|
|
def __init__(self) -> None:
|
|
self.values: Mapping[str, object] | None = None
|
|
|
|
|
|
def test_resolves_top_level_session_model():
|
|
resolved = _with_resolved_session_model({"model": "alias/gpt-realtime"}, "gpt-realtime")
|
|
assert resolved == {"model": "gpt-realtime"}
|
|
|
|
|
|
def test_session_without_model_is_returned_unchanged():
|
|
session = {"type": "realtime", "audio": {"input": {}}}
|
|
assert _with_resolved_session_model(session, "gpt-realtime") == session
|
|
|
|
|
|
def test_does_not_clobber_flat_transcription_model():
|
|
"""The nested transcription model is a different model than the realtime
|
|
conversation model and must not be overwritten with the routing model."""
|
|
resolved = _with_resolved_session_model(
|
|
{"model": "gpt-4o-realtime-preview", "input_audio_transcription": {"model": "whisper-1"}},
|
|
"gpt-4o-realtime-preview",
|
|
)
|
|
assert resolved["input_audio_transcription"]["model"] == "whisper-1"
|
|
|
|
|
|
def test_does_not_clobber_nested_audio_transcription_model():
|
|
resolved = _with_resolved_session_model(
|
|
{
|
|
"model": "gpt-4o-realtime-preview",
|
|
"audio": {"input": {"transcription": {"model": "whisper-1"}}},
|
|
},
|
|
"gpt-4o-realtime-preview",
|
|
)
|
|
assert resolved["audio"]["input"]["transcription"]["model"] == "whisper-1"
|
|
|
|
|
|
def test_original_session_is_not_mutated():
|
|
session = {"model": "alias/gpt-realtime"}
|
|
_with_resolved_session_model(session, "gpt-realtime")
|
|
assert session == {"model": "alias/gpt-realtime"}
|
|
|
|
|
|
def _run_client_secret(session, model, monkeypatch):
|
|
captured = {}
|
|
|
|
async def mock_handler(**kwargs):
|
|
captured.update(kwargs)
|
|
return object()
|
|
|
|
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
|
return model, "openai", None, api_base
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(
|
|
realtime_main.base_llm_http_handler,
|
|
"async_realtime_client_secret_handler",
|
|
mock_handler,
|
|
)
|
|
|
|
asyncio.run(
|
|
realtime_main.acreate_realtime_client_secret.__wrapped__(
|
|
model=model,
|
|
session=session,
|
|
litellm_logging_obj=FakeLogging(),
|
|
)
|
|
)
|
|
return captured
|
|
|
|
|
|
def test_client_secret_session_model_takes_priority_over_top_level(monkeypatch):
|
|
"""Backwards-compatible ordering: an explicit session.model wins over the
|
|
top-level model, matching the proxy's own resolution order."""
|
|
captured = _run_client_secret(
|
|
session={"model": "gpt-realtime-session"},
|
|
model="gpt-realtime-top-level",
|
|
monkeypatch=monkeypatch,
|
|
)
|
|
assert captured["model"] == "gpt-realtime-session"
|
|
assert captured["request_data"]["session"]["model"] == "gpt-realtime-session"
|
|
|
|
|
|
async def _hanging_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]:
|
|
await asyncio.sleep(30)
|
|
return "", ""
|
|
|
|
|
|
async def _thread_offloaded_hanging_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]:
|
|
from litellm.litellm_core_utils.asyncify import asyncify
|
|
|
|
await asyncify(time.sleep)(30)
|
|
return "", ""
|
|
|
|
|
|
async def _instant_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]:
|
|
return "token-abc", "resolved-project"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_credential_resolution_returns_the_resolved_token_and_project():
|
|
assert await realtime_main._resolve_vertex_access_token_bounded(
|
|
credentials="fake-credentials",
|
|
project_id="fake-project",
|
|
resolver=_instant_resolver,
|
|
timeout_seconds=5,
|
|
) == ("token-abc", "resolved-project")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_credential_resolution_times_out_instead_of_hanging():
|
|
"""Regression for the realtime accept-then-silence hang: a stalled Google
|
|
OAuth token refresh used to block the vertex branch unbounded (minutes of
|
|
zero frames for the client). It must raise promptly and name the timeout."""
|
|
start = time.monotonic()
|
|
with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"):
|
|
await realtime_main._resolve_vertex_access_token_bounded(
|
|
credentials="fake-credentials",
|
|
project_id="fake-project",
|
|
resolver=_hanging_resolver,
|
|
timeout_seconds=0.05,
|
|
)
|
|
assert time.monotonic() - start < 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_credential_resolution_bounds_a_thread_offloaded_refresh():
|
|
"""The real stall is a blocking google-auth refresh that runs in a worker
|
|
thread via asyncify, not a plain awaitable sleep. A timeout that only bounds
|
|
cancellable awaits would leave that shape hanging, so bound the shape the
|
|
proxy actually runs."""
|
|
start = time.monotonic()
|
|
with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"):
|
|
await realtime_main._resolve_vertex_access_token_bounded(
|
|
credentials="fake-credentials",
|
|
project_id="fake-project",
|
|
resolver=_thread_offloaded_hanging_resolver,
|
|
timeout_seconds=0.05,
|
|
)
|
|
assert time.monotonic() - start < 5
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_meta_realtime_dispatches_to_base_handler_with_meta_config(monkeypatch: pytest.MonkeyPatch):
|
|
from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig
|
|
|
|
captured: dict[str, object] = {}
|
|
|
|
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
|
return model.removeprefix("meta/"), "meta", None, api_base
|
|
|
|
async def mock_async_realtime(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime)
|
|
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="meta/muse-voice-transcribe-1.0",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
query_params={"model": "meta/muse-voice-transcribe-1.0", "intent": "transcription"},
|
|
)
|
|
|
|
assert isinstance(captured["provider_config"], MetaRealtimeConfig)
|
|
assert captured["model"] == "muse-voice-transcribe-1.0"
|
|
assert captured["query_params"] == {"model": "muse-voice-transcribe-1.0", "intent": "transcription"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_vertex_branch_resolves_credentials_under_a_bound(monkeypatch):
|
|
"""The wiring half of the regression: the vertex branch of _arealtime must
|
|
go through the bounded resolver, so a hung token refresh surfaces as a
|
|
prompt error there rather than as an accepted-then-silent websocket."""
|
|
|
|
async def hanging_token_refresh(**kwargs):
|
|
await asyncio.sleep(30)
|
|
|
|
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
|
return model, "vertex_ai", None, api_base
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", hanging_token_refresh)
|
|
monkeypatch.setattr(realtime_main, "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 0.05)
|
|
|
|
start = time.monotonic()
|
|
with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="gemini-live-2.5-flash",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
vertex_credentials="fake-credentials",
|
|
vertex_project="fake-project",
|
|
vertex_location="us-central1",
|
|
)
|
|
assert time.monotonic() - start < 5
|
|
|
|
|
|
def test_client_secret_forwards_nested_transcription_model_untouched(monkeypatch):
|
|
captured = _run_client_secret(
|
|
session={
|
|
"model": "gpt-4o-realtime-preview",
|
|
"input_audio_transcription": {"model": "whisper-1"},
|
|
},
|
|
model=None,
|
|
monkeypatch=monkeypatch,
|
|
)
|
|
session = captured["request_data"]["session"]
|
|
assert session["model"] == "gpt-4o-realtime-preview"
|
|
assert session["input_audio_transcription"]["model"] == "whisper-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_live_session_dispatches_to_openai_live_handler(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
captured: Final = _CallCapture()
|
|
|
|
def mock_get_llm_provider(
|
|
model: str,
|
|
api_base: str | None,
|
|
api_key: str | None,
|
|
custom_llm_provider: str | None = None,
|
|
) -> tuple[str, str, str | None, str | None]:
|
|
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
|
|
|
class FakeLiveHandler:
|
|
async def async_live_session(self, **kwargs: object) -> None:
|
|
captured.values = kwargs
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "openai_live_sessions", FakeLiveHandler())
|
|
start_frame: Final[dict[str, object]] = {
|
|
"type": "session.start",
|
|
"session": {"model": "live-alias", "instructions": "Stay brief"},
|
|
}
|
|
websocket: Final = MagicMock()
|
|
logging_obj: Final = FakeLogging()
|
|
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="live-alias",
|
|
websocket=websocket,
|
|
litellm_logging_obj=logging_obj,
|
|
live_session_start=start_frame,
|
|
)
|
|
|
|
assert captured.values is not None
|
|
assert captured.values["model"] == "gpt-live-1"
|
|
assert captured.values["websocket"] is websocket
|
|
assert captured.values["logging_obj"] is logging_obj
|
|
assert captured.values["session_start"] == start_frame
|
|
assert captured.values["api_base"] == "https://live-gateway.example/v1"
|
|
assert captured.values["api_key"] == "resolved-key"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_live_session_rejects_matching_guardrail_before_dispatch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
class RecordingGuardrail(CustomGuardrail):
|
|
def __init__(self) -> None:
|
|
super().__init__(
|
|
guardrail_name="live-blocker",
|
|
event_hook=GuardrailEventHooks.realtime_input_transcription,
|
|
default_on=True,
|
|
)
|
|
self.request_data: dict[str, object] | None = None
|
|
|
|
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
|
|
self.request_data = data
|
|
return super().should_run_guardrail(data, event_type)
|
|
|
|
class FakeLiveHandler:
|
|
def __init__(self) -> None:
|
|
self.called = False
|
|
|
|
async def async_live_session(self, **kwargs: object) -> None:
|
|
self.called = True
|
|
|
|
guardrail: Final = RecordingGuardrail()
|
|
live_handler: Final = FakeLiveHandler()
|
|
|
|
def mock_get_llm_provider(
|
|
model: str,
|
|
api_base: str | None,
|
|
api_key: str | None,
|
|
custom_llm_provider: str | None = None,
|
|
) -> tuple[str, str, str | None, str | None]:
|
|
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler)
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="live-alias",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
litellm_metadata={"session_id": "live-session"},
|
|
live_session_start={"type": "session.start", "session": {"model": "live-alias"}},
|
|
)
|
|
|
|
assert guardrail.request_data == {"litellm_metadata": {"session_id": "live-session"}, "stream": True}
|
|
assert live_handler.called is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_live_session_rejects_streaming_only_pre_call_guardrail(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
class FakeLiveHandler:
|
|
def __init__(self) -> None:
|
|
self.called = False
|
|
|
|
async def async_live_session(self, **kwargs: object) -> None:
|
|
self.called = True
|
|
|
|
guardrail: Final = CustomGuardrail(
|
|
guardrail_name="streaming-only-blocker",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
default_on=True,
|
|
stream_scope="streaming",
|
|
)
|
|
live_handler: Final = FakeLiveHandler()
|
|
|
|
def mock_get_llm_provider(
|
|
model: str,
|
|
api_base: str | None,
|
|
api_key: str | None,
|
|
custom_llm_provider: str | None = None,
|
|
) -> tuple[str, str, str | None, str | None]:
|
|
return "gpt-live-1", "openai", "resolved-key", "https://live-gateway.example/v1"
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "openai_live_sessions", live_handler)
|
|
monkeypatch.setattr(litellm, "callbacks", [guardrail])
|
|
|
|
with pytest.raises(ValueError, match="Guardrails are not supported on OpenAI Live sessions"):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="live-alias",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
litellm_metadata={"session_id": "live-session"},
|
|
live_session_start={"type": "session.start", "session": {"model": "live-alias"}},
|
|
)
|
|
|
|
assert live_handler.called is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_live_session_rejects_non_openai_provider(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
def mock_get_llm_provider(
|
|
model: str,
|
|
api_base: str | None,
|
|
api_key: str | None,
|
|
custom_llm_provider: str | None = None,
|
|
) -> tuple[str, str, str | None, str | None]:
|
|
return "claude-live", "anthropic", None, None
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
|
|
with pytest.raises(ValueError, match="anthropic"):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="anthropic/claude-live",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
live_session_start={"type": "session.start", "session": {"model": "claude-live"}},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_without_live_start_preserves_openai_realtime_dispatch(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
captured: Final = _CallCapture()
|
|
|
|
def mock_get_llm_provider(
|
|
model: str,
|
|
api_base: str | None,
|
|
api_key: str | None,
|
|
custom_llm_provider: str | None = None,
|
|
) -> tuple[str, str, str | None, str | None]:
|
|
return "gpt-realtime", "openai", "resolved-key", "https://realtime-gateway.example/v1"
|
|
|
|
class FakeRealtimeHandler:
|
|
async def async_realtime(self, **kwargs: object) -> None:
|
|
captured.values = kwargs
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "openai_realtime", FakeRealtimeHandler())
|
|
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="gpt-realtime",
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
)
|
|
|
|
assert captured.values is not None
|
|
assert captured.values["model"] == "gpt-realtime"
|
|
assert captured.values["api_key"] == "resolved-key"
|
|
assert captured.values["api_base"] == "https://realtime-gateway.example/v1"
|
|
|
|
|
|
class _CapturingConnect:
|
|
def __init__(self) -> None:
|
|
self.url: str | None = None
|
|
self.kwargs: dict[str, object] = {}
|
|
|
|
def __call__(self, url: str, **kwargs: object) -> "_CapturingConnect":
|
|
self.url = url
|
|
self.kwargs = kwargs
|
|
return self
|
|
|
|
async def __aenter__(self) -> MagicMock:
|
|
return MagicMock()
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc: BaseException | None,
|
|
tb: TracebackType | None,
|
|
) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_resolves_stored_credentials(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(
|
|
credential_name="azure-rt",
|
|
credential_values={
|
|
"api_key": "sk-from-credential",
|
|
"api_base": "https://example.openai.azure.com",
|
|
"api_version": "2025-04-01-preview",
|
|
},
|
|
credential_info={},
|
|
)
|
|
],
|
|
)
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-realtime",
|
|
custom_llm_provider="azure",
|
|
api_key=None,
|
|
realtime_protocol="beta",
|
|
model_params={"model": "azure/gpt-realtime", "litellm_credential_name": "azure-rt"},
|
|
)
|
|
assert connect.kwargs["additional_headers"] == {"api-key": "sk-from-credential"}
|
|
assert connect.url is not None
|
|
assert connect.url.startswith("wss://example.openai.azure.com")
|
|
assert "api-version=2025-04-01-preview" in connect.url
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("custom_llm_provider", "model", "expected_url"),
|
|
[
|
|
("xai", "grok-voice-latest", "wss://api.x.ai/v1/realtime?model=grok-voice-latest"),
|
|
("openai", "gpt-realtime", "wss://api.openai.com/v1/realtime?model=gpt-realtime"),
|
|
],
|
|
)
|
|
async def test_bearer_health_check_sends_stored_credential_as_bearer_token(
|
|
monkeypatch, custom_llm_provider: str, model: str, expected_url: str
|
|
):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"credential_list",
|
|
[
|
|
CredentialItem(
|
|
credential_name="voice-key",
|
|
credential_values={"api_key": "sk-from-credential"},
|
|
credential_info={},
|
|
)
|
|
],
|
|
)
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model=model,
|
|
custom_llm_provider=custom_llm_provider,
|
|
api_key=None,
|
|
model_params={"model": f"{custom_llm_provider}/{model}", "litellm_credential_name": "voice-key"},
|
|
)
|
|
assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer sk-from-credential"}
|
|
assert connect.url == expected_url
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_probes_ga_transcription_url_for_transcription_model(local_model_cost_map):
|
|
"""Regression for LIT-6240: transcription-only models (mode audio_transcription
|
|
in the cost map) are GA-only and 400 on the beta path, so the health probe
|
|
must hit /openai/v1/realtime?intent=transcription like real calls do."""
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-realtime-whisper",
|
|
custom_llm_provider="azure",
|
|
api_key="fake-key",
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_version="2025-04-01-preview",
|
|
)
|
|
assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?intent=transcription"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_stays_on_ga_when_deployment_registration_overwrites_mode(
|
|
local_model_cost_map, monkeypatch
|
|
):
|
|
"""In a live proxy, Router._register_deployment_in_model_cost writes the
|
|
operator's deployment model_info (mode: realtime) over the catalog entry for
|
|
azure/gpt-realtime-whisper, so mode alone misreads the model as speech-capable
|
|
and the probe regresses to the beta path. supported_endpoints survives that
|
|
registration and must keep the probe on the GA transcription path."""
|
|
polluted = {**litellm.model_cost["azure/gpt-realtime-whisper"], "mode": "realtime"}
|
|
monkeypatch.setitem(litellm.model_cost, "azure/gpt-realtime-whisper", polluted)
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-realtime-whisper",
|
|
custom_llm_provider="azure",
|
|
api_key="fake-key",
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_version="2025-04-01-preview",
|
|
)
|
|
assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?intent=transcription"
|
|
|
|
|
|
def test_transcription_only_detection_falls_back_to_mode(local_model_cost_map):
|
|
"""azure/whisper-1 declares mode audio_transcription but no supported_endpoints,
|
|
so only the mode signal can classify it as transcription-only."""
|
|
assert realtime_main._is_transcription_only_realtime_model("whisper-1", "azure") is True
|
|
|
|
|
|
def test_transcription_only_detection_rejects_speech_model(local_model_cost_map):
|
|
assert realtime_main._is_transcription_only_realtime_model("gpt-realtime-mini", "azure") is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_probes_the_ga_upstream_for_an_unconfigured_speech_model(monkeypatch):
|
|
monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False)
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-4o-realtime-preview",
|
|
custom_llm_provider="azure",
|
|
api_key="fake-key",
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_version="2024-10-01-preview",
|
|
)
|
|
assert connect.url == "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-4o-realtime-preview"
|
|
|
|
|
|
_AZURE_BETA_HEALTH_URL: Final = (
|
|
"wss://my-endpoint.openai.azure.com/openai/realtime"
|
|
"?api-version=2024-10-01-preview&deployment=gpt-4o-realtime-preview"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_honors_deployment_realtime_protocol(monkeypatch):
|
|
monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False)
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-4o-realtime-preview",
|
|
custom_llm_provider="azure",
|
|
api_key="fake-key",
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_version="2024-10-01-preview",
|
|
model_params={"realtime_protocol": "beta"},
|
|
)
|
|
assert connect.url == _AZURE_BETA_HEALTH_URL
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_health_check_honors_env_realtime_protocol(monkeypatch):
|
|
monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta")
|
|
connect = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gpt-4o-realtime-preview",
|
|
custom_llm_provider="azure",
|
|
api_key="fake-key",
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_version="2024-10-01-preview",
|
|
)
|
|
assert connect.url == _AZURE_BETA_HEALTH_URL
|
|
|
|
|
|
class _ConnectThatStopsAfterCapturingTheUrl:
|
|
url: str | None = None
|
|
|
|
def __call__(self, url: str, **kwargs: object) -> "_ConnectThatStopsAfterCapturingTheUrl":
|
|
self.url = url
|
|
return self
|
|
|
|
async def __aenter__(self) -> None:
|
|
raise RuntimeError("backend url captured, nothing to bridge")
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc: BaseException | None,
|
|
tb: TracebackType | None,
|
|
) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_azure_ai_on_a_foundry_host_connects_to_the_azure_openai_realtime_route():
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="azure_ai/gpt-realtime-mini",
|
|
websocket=MagicMock(),
|
|
api_base="https://my-project.services.ai.azure.com",
|
|
api_key="fake-key",
|
|
litellm_logging_obj=FakeLogging(),
|
|
)
|
|
assert connect.url == "wss://my-project.services.ai.azure.com/openai/v1/realtime?model=gpt-realtime-mini"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_routes_a_bare_model_name_to_its_explicit_provider():
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="my-realtime",
|
|
websocket=MagicMock(close=AsyncMock()),
|
|
custom_llm_provider="openai",
|
|
api_base="https://realtime-gateway.example/v1",
|
|
api_key="fake-key",
|
|
query_params={"model": "my-realtime"},
|
|
litellm_logging_obj=FakeLogging(),
|
|
)
|
|
assert connect.url == "wss://realtime-gateway.example/v1/realtime?model=my-realtime"
|
|
|
|
|
|
class _GatewayThatRecordsTheUrl:
|
|
url: str | None = None
|
|
|
|
def __call__(self, request: httpx.Request) -> httpx.Response:
|
|
self.url = str(request.url)
|
|
return httpx.Response(200, json={"value": "ek_test"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("route", "route_kwargs", "path"),
|
|
[
|
|
(realtime_main.acreate_realtime_client_secret, {}, "client_secrets"),
|
|
(realtime_main.acreate_realtime_transcription_session, {}, "transcription_sessions"),
|
|
(realtime_main.arealtime_calls, {"openai_ephemeral_key": "ek_test", "sdp_body": b"v=0"}, "calls"),
|
|
],
|
|
)
|
|
async def test_realtime_http_routes_send_a_bare_model_name_to_its_explicit_provider(
|
|
route: Callable[..., Awaitable[object]], route_kwargs: Mapping[str, object], path: str
|
|
) -> None:
|
|
gateway: Final = _GatewayThatRecordsTheUrl()
|
|
await route.__wrapped__(
|
|
model="my-realtime",
|
|
custom_llm_provider="openai",
|
|
api_base="https://realtime-gateway.example/v1",
|
|
api_key="fake-key",
|
|
litellm_logging_obj=FakeLogging(),
|
|
client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)),
|
|
**route_kwargs,
|
|
)
|
|
assert gateway.url == f"https://realtime-gateway.example/v1/realtime/{path}"
|
|
|
|
|
|
class _ClientWebSocketWithHeaders:
|
|
def __init__(self, headers: tuple[tuple[bytes, bytes], ...]) -> None:
|
|
self.scope: Final = {"headers": headers}
|
|
|
|
|
|
_GA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=())
|
|
_BETA_CLIENT: Final = _ClientWebSocketWithHeaders(headers=((b"openai-beta", b"realtime=v1"),))
|
|
|
|
|
|
async def _azure_backend_url_dialed_for(websocket: _ClientWebSocketWithHeaders, **kwargs: object) -> str | None:
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="azure/gpt-realtime",
|
|
websocket=websocket,
|
|
api_base="https://my-endpoint.openai.azure.com",
|
|
api_key="fake-key",
|
|
litellm_logging_obj=FakeLogging(),
|
|
**kwargs,
|
|
)
|
|
return connect.url
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_azure_ga_client_without_beta_header_dials_the_ga_upstream(monkeypatch):
|
|
monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False)
|
|
assert (
|
|
await _azure_backend_url_dialed_for(_GA_CLIENT)
|
|
== "wss://my-endpoint.openai.azure.com/openai/v1/realtime?model=gpt-realtime"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_azure_beta_header_client_keeps_the_beta_upstream(monkeypatch):
|
|
monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False)
|
|
assert await _azure_backend_url_dialed_for(_BETA_CLIENT) == (
|
|
"wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_azure_explicit_beta_protocol_wins_over_a_ga_client(monkeypatch):
|
|
monkeypatch.delenv("LITELLM_AZURE_REALTIME_PROTOCOL", raising=False)
|
|
assert await _azure_backend_url_dialed_for(_GA_CLIENT, realtime_protocol="beta") == (
|
|
"wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypatch):
|
|
monkeypatch.setenv("LITELLM_AZURE_REALTIME_PROTOCOL", "beta")
|
|
assert await _azure_backend_url_dialed_for(_GA_CLIENT) == (
|
|
"wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime"
|
|
)
|
|
|
|
|
|
async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None):
|
|
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig
|
|
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
|
|
|
|
captured: dict[str, object] = {}
|
|
|
|
def mock_get_llm_provider(model, api_base, api_key, custom_llm_provider: str | None = None):
|
|
return model.removeprefix("vertex_ai/"), "vertex_ai", None, api_base
|
|
|
|
async def mock_token_resolver(**kwargs):
|
|
return "access-token", kwargs["project_id"]
|
|
|
|
async def mock_async_realtime(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider)
|
|
monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", mock_token_resolver)
|
|
monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime)
|
|
monkeypatch.setattr(litellm, "vertex_location", None)
|
|
monkeypatch.delenv("VERTEXAI_LOCATION", raising=False)
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model=model,
|
|
websocket=MagicMock(),
|
|
litellm_logging_obj=FakeLogging(),
|
|
query_params={"model": model, "intent": "transcription"},
|
|
vertex_credentials="fake-credentials",
|
|
vertex_project="proj-1",
|
|
vertex_location=vertex_location,
|
|
)
|
|
provider_config = captured["provider_config"]
|
|
assert isinstance(provider_config, (VertexAIRealtimeConfig, VertexChirpRealtimeConfig))
|
|
return provider_config, captured["model"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_routes_chirp_models_to_the_speech_to_text_backend(monkeypatch):
|
|
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig
|
|
|
|
provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", None)
|
|
assert isinstance(provider_config, VertexChirpRealtimeConfig)
|
|
assert model == "chirp_3"
|
|
assert provider_config.get_complete_url(None, model) == "us-speech.googleapis.com"
|
|
assert provider_config.validate_environment({}, model, "https://us-speech.googleapis.com") == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_routes_chirp_models_to_the_configured_speech_region(monkeypatch):
|
|
provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", "europe-west4")
|
|
assert provider_config.get_complete_url(None, model) == "europe-west4-speech.googleapis.com"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_keeps_gemini_live_on_the_vertex_realtime_websocket(monkeypatch):
|
|
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
|
|
|
|
provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/gemini-live-2.5-flash", None)
|
|
assert isinstance(provider_config, VertexAIRealtimeConfig)
|
|
assert provider_config.get_complete_url(None, model).startswith("wss://us-central1-aiplatform.googleapis.com/")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_realtime_health_check_names_the_batch_mode_for_chirp_models():
|
|
with pytest.raises(ValueError, match="mode audio_transcription"):
|
|
await realtime_main._realtime_health_check(model="chirp_3", custom_llm_provider="vertex_ai", api_key=None)
|
|
|
|
|
|
class _ClosableGaClientWebSocket:
|
|
def __init__(self) -> None:
|
|
self.scope: Final = {"headers": ()}
|
|
|
|
async def close(self, code: int = 1000, reason: str = "") -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_openai_forwards_the_intent_query_param_to_the_upstream_url():
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="openai/gpt-realtime",
|
|
websocket=_ClosableGaClientWebSocket(),
|
|
api_key="fake-key",
|
|
query_params={"model": "openai/gpt-realtime", "intent": "chat"},
|
|
litellm_logging_obj=FakeLogging(),
|
|
)
|
|
assert connect.url == "wss://api.openai.com/v1/realtime?model=gpt-realtime&intent=chat"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model", "api_base", "client_query_params", "expected_backend_url"),
|
|
[
|
|
(
|
|
"openai/gpt-live-transcribe",
|
|
"https://api.openai.com/",
|
|
{"model": "my-transcribe-alias", "intent": "transcription"},
|
|
"wss://api.openai.com/v1/realtime?intent=transcription",
|
|
),
|
|
(
|
|
"openai/gpt-live-transcribe",
|
|
"https://eu.api.openai.com/",
|
|
{"model": "my-transcribe-alias", "intent": "transcription"},
|
|
"wss://eu.api.openai.com/v1/realtime?intent=transcription",
|
|
),
|
|
(
|
|
"openai/gpt-live-transcribe",
|
|
"https://api.openai.com/",
|
|
{"model": "my-realtime-alias"},
|
|
"wss://api.openai.com/v1/realtime?model=gpt-live-transcribe",
|
|
),
|
|
(
|
|
"openai/gateway-transcribe-alias",
|
|
"https://gateway.example/",
|
|
{"model": "my-transcribe-alias", "intent": "transcription"},
|
|
"wss://gateway.example/v1/realtime?model=gateway-transcribe-alias&intent=transcription",
|
|
),
|
|
(
|
|
"xai/grok-voice",
|
|
"https://api.x.ai/v1",
|
|
{"model": "my-voice-alias", "intent": "transcription"},
|
|
"wss://api.x.ai/v1/realtime?model=grok-voice&intent=transcription",
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_arealtime_drops_model_from_the_upstream_url_only_for_transcription_sessions_on_openai_hosts(
|
|
model, api_base, client_query_params, expected_backend_url
|
|
):
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheUrl()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model=model,
|
|
websocket=_ClosableGaClientWebSocket(),
|
|
api_base=api_base,
|
|
api_key="fake-key",
|
|
litellm_logging_obj=FakeLogging(),
|
|
query_params=client_query_params,
|
|
)
|
|
assert connect.url == expected_backend_url
|
|
|
|
|
|
async def _vertex_health_check_connect_for(monkeypatch: pytest.MonkeyPatch, api_base: str | None) -> _CapturingConnect:
|
|
async def fake_token_resolver(
|
|
credentials: object, project_id: str | None, custom_llm_provider: str
|
|
) -> tuple[str, str]:
|
|
return "access-token", project_id or ""
|
|
|
|
monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", fake_token_resolver)
|
|
connect: Final = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main._realtime_health_check(
|
|
model="gemini-3.8-live",
|
|
custom_llm_provider="vertex_ai",
|
|
api_key=None,
|
|
api_base=api_base,
|
|
model_params={
|
|
"vertex_project": "proj-1",
|
|
"vertex_credentials": "fake-credentials",
|
|
"vertex_location": "us",
|
|
},
|
|
)
|
|
return connect
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_health_check_sends_no_tls_argument_to_a_plain_ws_api_base(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
connect: Final = await _vertex_health_check_connect_for(monkeypatch, "http://127.0.0.1:8080")
|
|
assert connect.url == "ws://127.0.0.1:8080/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
|
assert connect.kwargs["ssl"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vertex_health_check_keeps_tls_for_the_multi_region_host(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
connect: Final = await _vertex_health_check_connect_for(monkeypatch, None)
|
|
assert connect.url == (
|
|
"wss://aiplatform.us.rep.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
|
)
|
|
assert connect.kwargs["ssl"] is not None
|
|
|
|
|
|
BLOCKED_PHRASE = "XSECRETBLOCKTESTPHRASEX"
|
|
|
|
|
|
class PhraseBlockingGuardrail(CustomGuardrail):
|
|
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
|
for text in inputs.get("texts", []):
|
|
if BLOCKED_PHRASE in text:
|
|
raise ValueError("Content blocked: contains forbidden test phrase.")
|
|
return inputs
|
|
|
|
|
|
def _make_guardrail(event_hook=GuardrailEventHooks.pre_call):
|
|
return PhraseBlockingGuardrail(
|
|
guardrail_name="integration-test-guard",
|
|
event_hook=event_hook,
|
|
default_on=True,
|
|
)
|
|
|
|
|
|
async def _build_streaming(client_events, backend_ws):
|
|
client_ws = MagicMock()
|
|
input_queue: asyncio.Queue = asyncio.Queue()
|
|
|
|
async def send_text(data: str):
|
|
client_events.append(json.loads(data))
|
|
|
|
client_ws.send_text = send_text
|
|
client_ws.receive_text = input_queue.get
|
|
|
|
logging_obj: Final = MagicMock(
|
|
pre_call=MagicMock(),
|
|
async_success_handler=AsyncMock(),
|
|
success_handler=MagicMock(),
|
|
dispatch_success_handlers=AsyncMock(),
|
|
dispatch_failure_handlers=AsyncMock(),
|
|
model_call_details={},
|
|
)
|
|
|
|
streaming = RealTimeStreaming(
|
|
websocket=client_ws,
|
|
backend_ws=backend_ws,
|
|
logging_obj=logging_obj,
|
|
request_data={"guardrails": ["integration-test-guard"]},
|
|
)
|
|
return streaming, input_queue
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_voice_transcript_blocked_by_guardrail():
|
|
"""
|
|
Simulate a backend-side voice transcription event containing the blocked phrase.
|
|
Guardrail must block it - no response.create sent to OpenAI.
|
|
"""
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
guardrail = _make_guardrail(GuardrailEventHooks.realtime_input_transcription)
|
|
litellm.callbacks = [guardrail]
|
|
|
|
client_events: list[dict] = []
|
|
|
|
# Build the transcript event that would come from the OpenAI backend
|
|
transcript_event = json.dumps(
|
|
{
|
|
"type": "conversation.item.input_audio_transcription.completed",
|
|
"transcript": f"This is {BLOCKED_PHRASE} in my voice message",
|
|
"item_id": "item_integ_test",
|
|
}
|
|
).encode()
|
|
|
|
# Mock backend that delivers the transcript then closes
|
|
backend_ws = MagicMock()
|
|
backend_ws.recv = AsyncMock(
|
|
side_effect=[
|
|
transcript_event,
|
|
ConnectionClosed(None, None),
|
|
]
|
|
)
|
|
backend_ws.send = AsyncMock()
|
|
|
|
try:
|
|
streaming, _ = await _build_streaming(client_events, backend_ws)
|
|
await streaming.backend_to_client_send_messages()
|
|
|
|
event_types = [e.get("type") for e in client_events]
|
|
|
|
# 1. Error event must be sent to client
|
|
error_events = [e for e in client_events if e.get("type") == "error"]
|
|
assert len(error_events) >= 1, f"Expected guardrail error event, got: {event_types}"
|
|
assert error_events[0]["error"]["type"] == "guardrail_violation"
|
|
|
|
# 2. Check what was sent to backend.
|
|
# The guardrail may send response.cancel + conversation.item.create (block msg)
|
|
# + response.create (to speak the block message). That's acceptable.
|
|
# What we assert is that a response.cancel was sent (blocking the original).
|
|
sent_to_backend = [
|
|
json.loads(c.args[0]) for c in backend_ws.send.call_args_list if c.args and isinstance(c.args[0], str)
|
|
]
|
|
response_cancels = [e for e in sent_to_backend if e.get("type") == "response.cancel"]
|
|
assert len(response_cancels) >= 1 or len(sent_to_backend) == 0, (
|
|
f"Guardrail should have sent response.cancel or nothing, got: {sent_to_backend}"
|
|
)
|
|
|
|
# Note: The guardrail may or may not send transcript deltas; the error event
|
|
# (assertion #1) is the primary signal that the blocked content was handled.
|
|
|
|
finally:
|
|
litellm.callbacks = []
|
|
|
|
|
|
@pytest.fixture
|
|
def federated_deployment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> dict[str, str]:
|
|
token_file: Final = tmp_path / "subject_token.jwt"
|
|
token_file.write_text("subject-token")
|
|
for name in (
|
|
"OPENAI_API_KEY",
|
|
"OPENAI_BASE_URL",
|
|
"OPENAI_API_BASE",
|
|
"OPENAI_IDENTITY_PROVIDER_ID",
|
|
"OPENAI_SERVICE_ACCOUNT_ID",
|
|
"OPENAI_IDENTITY_TOKEN_FILE",
|
|
):
|
|
monkeypatch.delenv(name, raising=False)
|
|
for attribute in ("api_key", "openai_key", "api_base"):
|
|
monkeypatch.setattr(litellm, attribute, None)
|
|
_workload_identity_auth.cache_clear()
|
|
return {
|
|
"openai_identity_provider_id": "idp_realtime",
|
|
"openai_service_account_id": "user-realtime",
|
|
"openai_identity_token_file": str(token_file),
|
|
}
|
|
|
|
|
|
def _mock_token_exchange() -> None:
|
|
respx.post("https://auth.openai.com/oauth/token").mock(
|
|
return_value=httpx.Response(200, json={"access_token": "federated-bearer", "expires_in": 3600})
|
|
)
|
|
|
|
|
|
class _ConnectThatStopsAfterCapturingTheAuthorization:
|
|
authorization: str | None = None
|
|
|
|
def __call__(
|
|
self, url: str, *, additional_headers: Mapping[str, str], **kwargs: object
|
|
) -> "_ConnectThatStopsAfterCapturingTheAuthorization":
|
|
self.authorization = additional_headers.get("Authorization")
|
|
return self
|
|
|
|
async def __aenter__(self) -> None:
|
|
raise RuntimeError("backend headers captured, nothing to bridge")
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc: BaseException | None,
|
|
tb: TracebackType | None,
|
|
) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"live_session_start",
|
|
[None, {"type": "session.start", "session": {"model": "gpt-live-1"}}],
|
|
ids=["realtime", "live_session"],
|
|
)
|
|
@respx.mock
|
|
async def test_openai_websocket_sessions_send_the_federated_bearer_for_a_keyless_deployment(
|
|
federated_deployment: dict[str, str], live_session_start: Mapping[str, object] | None
|
|
) -> None:
|
|
_mock_token_exchange()
|
|
connect: Final = _ConnectThatStopsAfterCapturingTheAuthorization()
|
|
with patch("websockets.connect", connect):
|
|
await realtime_main._arealtime.__wrapped__(
|
|
model="gpt-realtime",
|
|
custom_llm_provider="openai",
|
|
websocket=MagicMock(close=AsyncMock()),
|
|
litellm_logging_obj=FakeLogging(),
|
|
live_session_start=live_session_start,
|
|
**federated_deployment,
|
|
)
|
|
assert connect.authorization == "Bearer federated-bearer"
|
|
|
|
|
|
class _GatewayThatRecordsTheAuthorization:
|
|
authorization: str | None = None
|
|
|
|
def __call__(self, request: httpx.Request) -> httpx.Response:
|
|
self.authorization = request.headers.get("Authorization")
|
|
return httpx.Response(200, json={"value": "ek_test"})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"route",
|
|
[realtime_main.acreate_realtime_client_secret, realtime_main.acreate_realtime_transcription_session],
|
|
ids=["client_secrets", "transcription_sessions"],
|
|
)
|
|
@respx.mock
|
|
async def test_openai_realtime_http_routes_send_the_federated_bearer_for_a_keyless_deployment(
|
|
federated_deployment: dict[str, str], route: Callable[..., Awaitable[object]]
|
|
) -> None:
|
|
_mock_token_exchange()
|
|
gateway: Final = _GatewayThatRecordsTheAuthorization()
|
|
await route.__wrapped__(
|
|
model="gpt-realtime",
|
|
custom_llm_provider="openai",
|
|
litellm_logging_obj=FakeLogging(),
|
|
client=AsyncHTTPHandler(transport=httpx.MockTransport(gateway)),
|
|
**federated_deployment,
|
|
)
|
|
assert gateway.authorization == "Bearer federated-bearer"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@respx.mock
|
|
async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_keyless_deployment(
|
|
federated_deployment: dict[str, str],
|
|
) -> None:
|
|
_mock_token_exchange()
|
|
connect: Final = _CapturingConnect()
|
|
with patch("websockets.connect", connect):
|
|
assert await realtime_main.realtime_health_check(
|
|
model="gpt-realtime",
|
|
custom_llm_provider="openai",
|
|
api_key=None,
|
|
model_params={"model": "openai/gpt-realtime", **federated_deployment},
|
|
)
|
|
assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer federated-bearer"}
|