litellm/tests/unit/realtime_api/test_main.py
Mateo Wang b39eadcb82
feat(openai)!: honor per-deployment workload identity on non-chat routes (#39966)
* 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>
2026-10-16 22:25:15 +00:00

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"}