scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #34658 from BerriAI/litellm_azure_realtime_entra_id_auth

Browse source 
fix(azure/realtime): authenticate realtime websocket with Azure AD token when no api-key
This commit is contained in:
Mateo Wang 2026-08-26 11:44:38 -07:00 • committed by GitHub
commit 767e6015af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 261 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/llms/azure/realtime/handler.py
View file

@ -4,6 +4,8 @@ This file contains the calling Azure OpenAI's `/openai/realtime` endpoint.
This requires websockets, and is currently only supported on LiteLLM Proxy.
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, cast
from litellm._logging import _redact_string, verbose_proxy_logger
@ -30,6 +32,21 @@ async def forward_messages(client_ws: Any, backend_ws: Any):
class AzureOpenAIRealtime(AzureChatCompletion):
@staticmethod
def get_auth_headers(api_key: str | None, azure_ad_token: str | None) -> Mapping[str, str]:
"""
Build the websocket handshake auth headers, preferring a static api-key and falling back to
an Azure AD (Entra ID) bearer token. Never sends both.
"""
if api_key:
return MappingProxyType({"api-key": api_key})
if azure_ad_token:
return MappingProxyType({"Authorization": f"Bearer {azure_ad_token}"})
raise ValueError(
"Missing Azure credentials for the realtime endpoint. Set an api_key, or configure Azure AD auth "
"(azure_ad_token, tenant_id/client_id/client_secret, or a managed identity)"
)
def _construct_url(
self,
api_base: str,
@ -117,13 +134,13 @@ class AzureOpenAIRealtime(AzureChatCompletion):
query_params=query_params,
)
auth_headers: Final = self.get_auth_headers(api_key=api_key, azure_ad_token=azure_ad_token)
try:
ssl_context: Final = get_shared_realtime_ssl_context()
async with websockets.connect(
url,
additional_headers={
"api-key": api_key,
},
additional_headers=auth_headers,
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
ssl=ssl_context,
) as backend_ws:

29
litellm/realtime_api/main.py
View file

@ -2,6 +2,8 @@
import asyncio
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, Literal, cast
import litellm
@ -29,6 +31,7 @@ from litellm.utils import ProviderConfigManager
from ..litellm_core_utils.get_litellm_params import get_litellm_params
from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from ..llms.azure.common_utils import get_azure_ad_token
from ..llms.azure.realtime.handler import AzureOpenAIRealtime
from ..llms.bedrock.realtime.handler import BedrockRealtime
from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context
@ -44,6 +47,7 @@ bedrock_realtime: Final = BedrockRealtime()
xai_realtime: Final = XAIRealtime()
vertex_llm_base: Final = VertexBase()
base_llm_http_handler = BaseLLMHTTPHandler()
_EMPTY_MODEL_PARAMS: Final[Mapping[str, Any]] = MappingProxyType({})
def _with_resolved_session_model(session: dict[str, Any], model_name: str) -> dict[str, Any]:
@ -411,13 +415,16 @@ async def _arealtime(
if realtime_protocol is None and (query_params or {}).get("intent") == "transcription":
realtime_protocol = "GA"
realtime_protocol = realtime_protocol or "beta"
resolved_azure_ad_token: Final = (
None if api_key else get_azure_ad_token(GenericLiteLLMParams(**kwargs, azure_ad_token=azure_ad_token))
)
await azure_realtime.async_realtime(
model=model,
websocket=websocket,
api_base=api_base,
api_key=api_key,
api_version=api_version,
azure_ad_token=None,
azure_ad_token=resolved_azure_ad_token,
client=None,
timeout=timeout,
logging_obj=litellm_logging_obj,
@ -550,6 +557,17 @@ async def _arealtime(
raise ValueError(f"Unsupported model: {model}")
def _realtime_health_check_auth_headers(
custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, Any]
) -> Mapping[str, str | None]:
if custom_llm_provider != "azure":
return MappingProxyType({"api-key": api_key})
return azure_realtime.get_auth_headers(
api_key=api_key,
azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))),
)
async def _realtime_health_check(
model: str,
custom_llm_provider: str,
@ -578,6 +596,11 @@ async def _realtime_health_check(
import websockets
url: str | None = None
auth_headers: Final = _realtime_health_check_auth_headers(
custom_llm_provider=custom_llm_provider,
api_key=api_key,
model_params=model_params or _EMPTY_MODEL_PARAMS,
)
if custom_llm_provider == "azure":
url = azure_realtime._construct_url(
api_base=api_base or "",
@ -627,9 +650,7 @@ async def _realtime_health_check(
ssl_context = get_shared_realtime_ssl_context()
async with websockets.connect(
url,
additional_headers={
"api-key": api_key,
},
additional_headers=auth_headers,
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
ssl=ssl_context,
):

216
tests/test_litellm/llms/azure/realtime/test_azure_realtime_handler.py
View file

@ -559,3 +559,219 @@ async def test_async_realtime_default_maintains_backwards_compatibility():
mock_realtime_streaming.call_args.kwargs["backend_uses_beta_protocol"]
is True
)
class _DummyAsyncContextManager:
def __init__(self, value):
self.value = value
async def __aenter__(self):
return self.value
async def __aexit__(self, exc_type, exc, tb):
return None
@pytest.mark.asyncio
async def test_async_realtime_uses_bearer_token_when_no_api_key():
"""
Entra ID-only Azure realtime deployments have no static api-key, so the handshake must
authenticate with `Authorization: Bearer <azure_ad_token>` and must not send `api-key`.
Regression test for https://github.com/BerriAI/litellm/issues/34654
"""
from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime
handler = AzureOpenAIRealtime()
mock_backend_ws = AsyncMock()
with (
patch(
"websockets.connect",
return_value=_DummyAsyncContextManager(mock_backend_ws),
) as mock_ws_connect,
patch( # test-quality-ok: handler owns the streaming loop, only the handshake headers are under test
"litellm.llms.azure.realtime.handler.RealTimeStreaming"
) as mock_realtime_streaming,
):
mock_realtime_streaming.return_value.bidirectional_forward = AsyncMock()
await handler.async_realtime(
model="gpt-realtime-whisper",
websocket=AsyncMock(),
logging_obj=MagicMock(),
api_base="https://my-endpoint.openai.azure.com",
api_key=None,
api_version="2024-10-01-preview",
azure_ad_token="my-entra-token",
)
headers = mock_ws_connect.call_args.kwargs["additional_headers"]
assert headers == {"Authorization": "Bearer my-entra-token"}
def test_get_auth_headers_prefers_api_key_and_never_sends_both():
from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime
assert AzureOpenAIRealtime.get_auth_headers(api_key="test-key", azure_ad_token="my-entra-token") == {
"api-key": "test-key"
}
def test_get_auth_headers_without_credentials_raises():
from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime
with pytest.raises(ValueError, match="Missing Azure credentials"):
AzureOpenAIRealtime.get_auth_headers(api_key=None, azure_ad_token=None)
@pytest.mark.asyncio
async def test_arealtime_resolves_azure_ad_token_when_no_api_key(monkeypatch):
"""
`_arealtime` must resolve an Azure AD token (managed identity, service principal, etc.)
and forward it to the handler when the deployment has no api_key.
Regression test for https://github.com/BerriAI/litellm/issues/34654
"""
from litellm.realtime_api import main as realtime_main
mock_async_realtime = AsyncMock()
monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime))
monkeypatch.setattr(
realtime_main,
"get_llm_provider",
lambda model, api_base=None, api_key=None: (
"gpt-realtime-whisper",
"azure",
None,
"https://my-endpoint.openai.azure.com",
),
)
monkeypatch.delenv("AZURE_API_KEY", raising=False)
captured_params = {}
def fake_get_azure_ad_token(litellm_params):
captured_params["tenant_id"] = litellm_params.get("tenant_id")
return "my-entra-token"
monkeypatch.setattr(realtime_main, "get_azure_ad_token", fake_get_azure_ad_token)
await realtime_main._arealtime(
model="azure/gpt-realtime-whisper",
websocket=MagicMock(),
api_version="2024-10-01-preview",
litellm_logging_obj=MagicMock(),
tenant_id="my-tenant",
client_id="my-client",
client_secret="my-secret",
)
assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "my-entra-token"
assert captured_params["tenant_id"] == "my-tenant"
@pytest.mark.asyncio
async def test_arealtime_does_not_resolve_azure_ad_token_when_api_key_present(monkeypatch):
from litellm.realtime_api import main as realtime_main
mock_async_realtime = AsyncMock()
monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime))
monkeypatch.setattr(
realtime_main,
"get_llm_provider",
lambda model, api_base=None, api_key=None: (
"gpt-realtime-whisper",
"azure",
"test-key",
"https://my-endpoint.openai.azure.com",
),
)
def fail_get_azure_ad_token(litellm_params):
raise AssertionError("should not resolve an AD token when an api_key is configured")
monkeypatch.setattr(realtime_main, "get_azure_ad_token", fail_get_azure_ad_token)
await realtime_main._arealtime(
model="azure/gpt-realtime-whisper",
websocket=MagicMock(),
api_key="test-key",
api_version="2024-10-01-preview",
litellm_logging_obj=MagicMock(),
)
assert mock_async_realtime.call_args.kwargs["azure_ad_token"] is None
@pytest.mark.asyncio
async def test_realtime_health_check_uses_bearer_token_when_no_api_key(monkeypatch):
"""
An Entra ID-only realtime deployment must also pass its realtime health check.
Regression test for https://github.com/BerriAI/litellm/issues/34654
"""
from litellm.realtime_api import main as realtime_main
connect_calls = []
monkeypatch.setattr(
realtime_main,
"get_azure_ad_token",
lambda litellm_params: "my-entra-token",
)
def fake_connect(url, **kwargs):
connect_calls.append(kwargs)
return _DummyAsyncContextManager(MagicMock())
monkeypatch.setattr("websockets.connect", fake_connect)
assert (
await realtime_main._realtime_health_check(
model="gpt-realtime-whisper",
custom_llm_provider="azure",
api_key=None,
api_base="https://my-endpoint.openai.azure.com",
api_version="2024-10-01-preview",
model_params={"tenant_id": "my-tenant"},
)
is True
)
assert connect_calls[0]["additional_headers"] == {"Authorization": "Bearer my-entra-token"}
@pytest.mark.asyncio
async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch):
"""
The router binds a deployment's `azure_ad_token` to `_arealtime`'s named parameter rather than
**kwargs, so it must still reach the handler.
Regression test for https://github.com/BerriAI/litellm/issues/34654
"""
from litellm.realtime_api import main as realtime_main
mock_async_realtime = AsyncMock()
monkeypatch.setattr(realtime_main, "azure_realtime", MagicMock(async_realtime=mock_async_realtime))
monkeypatch.setattr(
realtime_main,
"get_llm_provider",
lambda model, api_base=None, api_key=None: (
"gpt-realtime-whisper",
"azure",
None,
"https://my-endpoint.openai.azure.com",
),
)
monkeypatch.delenv("AZURE_API_KEY", raising=False)
monkeypatch.setattr(realtime_main.litellm, "api_key", None)
await realtime_main._arealtime(
model="azure/gpt-realtime-whisper",
websocket=MagicMock(),
api_version="2024-10-01-preview",
azure_ad_token="deployment-entra-token",
litellm_logging_obj=MagicMock(),
)
assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "deployment-entra-token"