mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
feat(vertex_ai): Vertex AI Gemini Live via unified /realtime endpoint (#22153)
* feat(vertex_ai): add Vertex AI Gemini Live support via unified /realtime endpoint Adds VertexAIRealtimeConfig which translates the OpenAI Realtime WebSocket protocol to Vertex AI BidiGenerateContent. Supports voice in/voice out (16 kHz mic → 24 kHz speaker) and text in/text out through the proxy's /realtime endpoint. Key changes: - New litellm/llms/vertex_ai/realtime/transformation.py with VertexAIRealtimeConfig - Builds correct wss:// URL (regional + global) - OAuth2 Bearer token auth (not API key) - Full model path (projects/.../publishers/google/models/...) - Ignores session.update (Vertex AI only accepts one setup message) - realtime_api/main.py: vertex_ai branch resolves OAuth token + constructs config - llm_http_handler.py: auto-sends session setup before bidirectional_forward - gemini/realtime/transformation.py: fix crashes on empty turnComplete events - realtime_streaming.py: try/except guard so bad messages don't kill the loop - proxy_server.py: add missing websockets.exceptions import * docs: add vertex_realtime to sidebars * fix: drop unknown event types in Gemini transform; add vertex_ai health check * fix: propagate UUID fallback IDs from transform_content_done_event to return_additional_content_done_events * fix: route guardrail backend sends through provider transform; fix str.strip misuse for model prefix * fix: handle Vertex AI full resource path in session.created; route guardrail block sends through _send_to_backend * fix: remove unused VertexBase in transformation.py; apply UUID fallback in return_additional_content_done_events
This commit is contained in:
Ishaan Jaff
Sameer Kankute
parent
6100380568
commit
eb8b74a14e
10 changed files with 768 additions and 45 deletions
203
docs/my-website/docs/providers/vertex_realtime.md
Normal file
203
docs/my-website/docs/providers/vertex_realtime.md
Normal file
|
|
@ -0,0 +1,203 @@
|
|||
# Vertex AI Gemini Live - Realtime API
|
||||
|
||||
Use Vertex AI's Gemini Live API (BidiGenerateContent) through LiteLLM's unified `/realtime` endpoint, which speaks the OpenAI Realtime protocol.
|
||||
|
||||
| Feature | Supported |
|
||||
|---------|-----------|
|
||||
| Proxy (`/realtime`) | ✅ |
|
||||
| Voice in / Voice out | ✅ |
|
||||
| Text in / Text out | ✅ |
|
||||
| Server VAD | ✅ |
|
||||
| Output transcription | ✅ |
|
||||
|
||||
## Setup
|
||||
|
||||
### 1. Auth
|
||||
|
||||
LiteLLM uses your Google Cloud credentials (OAuth2 Bearer token), not an API key.
|
||||
|
||||
```bash
|
||||
gcloud auth application-default login
|
||||
```
|
||||
|
||||
Or set a service-account key file:
|
||||
|
||||
```bash
|
||||
export GOOGLE_APPLICATION_CREDENTIALS=/path/to/sa-key.json
|
||||
```
|
||||
|
||||
### 2. Proxy config
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: vertex-gemini-live
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-2.0-flash-live-001
|
||||
vertex_project: your-gcp-project-id
|
||||
vertex_location: us-east4 # or any supported region, or "global"
|
||||
|
||||
general_settings:
|
||||
master_key: sk-your-key
|
||||
```
|
||||
|
||||
### 3. Start the proxy
|
||||
|
||||
```bash
|
||||
litellm --config config.yaml --port 4000
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Python (websockets)
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
import json
|
||||
import websockets
|
||||
|
||||
PROXY_URL = "ws://localhost:4000/realtime?model=vertex-gemini-live"
|
||||
API_KEY = "sk-your-key"
|
||||
|
||||
async def main():
|
||||
async with websockets.connect(
|
||||
PROXY_URL,
|
||||
additional_headers={"api-key": API_KEY},
|
||||
) as ws:
|
||||
# Wait for session.created
|
||||
event = json.loads(await ws.recv())
|
||||
print(f"session.created: {event['session']['id']}")
|
||||
|
||||
# Send a text message
|
||||
await ws.send(json.dumps({
|
||||
"type": "conversation.item.create",
|
||||
"item": {
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Say hello in one sentence."}],
|
||||
},
|
||||
}))
|
||||
|
||||
# Collect the response
|
||||
async for raw in ws:
|
||||
ev = json.loads(raw)
|
||||
t = ev.get("type", "")
|
||||
if t == "response.text.delta":
|
||||
print(ev.get("delta", ""), end="", flush=True)
|
||||
elif t == "response.done":
|
||||
print("\n[done]")
|
||||
break
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
### Node.js
|
||||
|
||||
```js
|
||||
const WebSocket = require("ws");
|
||||
|
||||
const ws = new WebSocket(
|
||||
"ws://localhost:4000/realtime?model=vertex-gemini-live",
|
||||
{ headers: { "api-key": "sk-your-key" } }
|
||||
);
|
||||
|
||||
ws.on("open", () => {
|
||||
ws.send(JSON.stringify({
|
||||
type: "conversation.item.create",
|
||||
item: {
|
||||
type: "message",
|
||||
role: "user",
|
||||
content: [{ type: "input_text", text: "Say hello." }],
|
||||
},
|
||||
}));
|
||||
});
|
||||
|
||||
ws.on("message", (data) => {
|
||||
const ev = JSON.parse(data);
|
||||
if (ev.type === "response.text.delta") process.stdout.write(ev.delta);
|
||||
if (ev.type === "response.done") ws.close();
|
||||
});
|
||||
```
|
||||
|
||||
### OpenAI SDK (Python)
|
||||
|
||||
```python
|
||||
import asyncio
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
client = AsyncOpenAI(
|
||||
base_url="http://localhost:4000",
|
||||
api_key="sk-your-key",
|
||||
)
|
||||
|
||||
async def main():
|
||||
async with client.beta.realtime.connect(
|
||||
model="vertex-gemini-live"
|
||||
) as conn:
|
||||
await conn.session.update(session={"modalities": ["text"]})
|
||||
|
||||
await conn.conversation.item.create(
|
||||
item={
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Say hello."}],
|
||||
}
|
||||
)
|
||||
|
||||
async for event in conn:
|
||||
if event.type == "response.text.delta":
|
||||
print(event.delta, end="", flush=True)
|
||||
elif event.type == "response.done":
|
||||
print()
|
||||
break
|
||||
|
||||
asyncio.run(main())
|
||||
```
|
||||
|
||||
## Voice in / Voice out
|
||||
|
||||
For a complete voice example see [`voice_realtime_test.py`](https://github.com/BerriAI/litellm/blob/main/voice_realtime_test.py).
|
||||
|
||||
Key settings for audio:
|
||||
- Microphone input: **16 kHz** PCM16 (`audio/pcm;rate=16000`)
|
||||
- Speaker output: **24 kHz** PCM16 (Vertex AI returns audio at 24 kHz)
|
||||
- Server VAD is enabled by default with 800 ms silence threshold
|
||||
|
||||
```python
|
||||
# session.update with server VAD — the proxy ignores this for Vertex AI
|
||||
# because VAD is already configured in the initial setup message.
|
||||
await ws.send(json.dumps({
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"modalities": ["audio"],
|
||||
"turn_detection": {"type": "server_vad", "silence_duration_ms": 800},
|
||||
},
|
||||
}))
|
||||
```
|
||||
|
||||
## Supported OpenAI Realtime Events
|
||||
|
||||
**Client → Proxy (→ Vertex AI)**
|
||||
|
||||
| OpenAI event | Notes |
|
||||
|---|---|
|
||||
| `input_audio_buffer.append` | Forwarded as `realtime_input.audio` |
|
||||
| `conversation.item.create` | Forwarded as `realtime_input.text` |
|
||||
| `session.update` | Silently ignored — Vertex AI does not support mid-session reconfiguration |
|
||||
| `response.create` | Silently ignored — Vertex AI responds automatically after each turn |
|
||||
|
||||
**Vertex AI → Proxy (→ Client)**
|
||||
|
||||
| OpenAI event emitted | Vertex AI source |
|
||||
|---|---|
|
||||
| `session.created` | Synthesized after `setupComplete` |
|
||||
| `response.text.delta` | `serverContent.modelTurn.parts[].text` |
|
||||
| `response.audio.delta` | `serverContent.modelTurn.parts[].inlineData` |
|
||||
| `response.audio_transcript.delta` | `serverContent.outputTranscription.text` |
|
||||
| `conversation.item.input_audio_transcription.completed` | `serverContent.inputTranscription.text` |
|
||||
| `response.done` | `serverContent.turnComplete` |
|
||||
|
||||
## Limitations
|
||||
|
||||
- `session.update` is not forwarded (Vertex AI only accepts one setup message per connection).
|
||||
- Tool calling / function calling is not yet supported.
|
||||
- Audio transcription requires `outputAudioTranscription: {}` to be set in the initial setup (done automatically by LiteLLM).
|
||||
|
|
@ -758,6 +758,7 @@ const sidebars = {
|
|||
"providers/vertex_batch",
|
||||
"providers/vertex_ocr",
|
||||
"providers/vertex_ai_agent_engine",
|
||||
"providers/vertex_realtime",
|
||||
]
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -145,7 +145,9 @@ class RealTimeStreaming:
|
|||
except (json.JSONDecodeError, AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
def _collect_user_input_from_backend_event(self, event_obj: dict) -> None:
|
||||
def _collect_user_input_from_backend_event(
|
||||
self, event_obj: Union[dict, OpenAIRealtimeEvents]
|
||||
) -> None:
|
||||
"""Extract user voice transcription from backend events for spend logging."""
|
||||
try:
|
||||
event_type = event_obj.get("type", "")
|
||||
|
|
@ -162,7 +164,7 @@ class RealTimeStreaming:
|
|||
pass
|
||||
|
||||
def _collect_tool_calls_from_response_done(
|
||||
self, event_obj: dict
|
||||
self, event_obj: Union[dict, OpenAIRealtimeEvents]
|
||||
) -> None:
|
||||
"""Extract function_call items from response.done events for spend logging."""
|
||||
try:
|
||||
|
|
@ -211,6 +213,23 @@ class RealTimeStreaming:
|
|||
## SYNC LOGGING
|
||||
executor.submit(self.logging_obj.success_handler(self.messages))
|
||||
|
||||
async def _send_to_backend(self, message: str) -> None:
|
||||
"""Send a message to the backend WebSocket.
|
||||
|
||||
If a provider_config is set the message is first passed through
|
||||
transform_realtime_request so that provider-specific translation
|
||||
(e.g. dropping session.update for Vertex AI) is applied even for
|
||||
guardrail-injected messages.
|
||||
"""
|
||||
if self.provider_config:
|
||||
transformed = self.provider_config.transform_realtime_request(
|
||||
message, self.model, self.session_configuration_request
|
||||
)
|
||||
for msg in transformed:
|
||||
await self.backend_ws.send(msg)
|
||||
else:
|
||||
await self.backend_ws.send(message)
|
||||
|
||||
def _has_realtime_guardrails(self) -> bool:
|
||||
"""Return True if any callback is registered for realtime_input_transcription."""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -276,9 +295,9 @@ class RealTimeStreaming:
|
|||
safe_msg = str(e) or "I'm sorry, that request was blocked by the content filter."
|
||||
# Cancel any in-flight response before speaking the warning.
|
||||
# This handles the race where create_response fired before we could intercept.
|
||||
await self.backend_ws.send(json.dumps({"type": "response.cancel"}))
|
||||
# Ask OpenAI to speak the warning — TTS audio plays naturally in the client
|
||||
await self.backend_ws.send(
|
||||
await self._send_to_backend(json.dumps({"type": "response.cancel"}))
|
||||
# Ask the model to speak the warning — TTS audio plays naturally in the client
|
||||
await self._send_to_backend(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "response.create",
|
||||
|
|
@ -333,7 +352,7 @@ class RealTimeStreaming:
|
|||
## GUARDRAIL: inject create_response=false on session.created
|
||||
if isinstance(event, dict) and event.get("type") == "session.created":
|
||||
if self._has_realtime_guardrails():
|
||||
await self.backend_ws.send(
|
||||
await self._send_to_backend(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
|
|
@ -362,7 +381,7 @@ class RealTimeStreaming:
|
|||
transcript, item_id=event.get("item_id")
|
||||
)
|
||||
if not blocked:
|
||||
await self.backend_ws.send(
|
||||
await self._send_to_backend(
|
||||
json.dumps({"type": "response.create"})
|
||||
)
|
||||
continue
|
||||
|
|
@ -383,7 +402,7 @@ class RealTimeStreaming:
|
|||
# set create_response=false so the LLM never auto-responds
|
||||
# before our guardrail has a chance to run.
|
||||
if self._has_realtime_guardrails():
|
||||
await self.backend_ws.send(
|
||||
await self._send_to_backend(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
|
|
@ -416,7 +435,7 @@ class RealTimeStreaming:
|
|||
)
|
||||
if not blocked:
|
||||
# Clean — trigger LLM response
|
||||
await self.backend_ws.send(
|
||||
await self._send_to_backend(
|
||||
json.dumps({"type": "response.create"})
|
||||
)
|
||||
return True
|
||||
|
|
@ -437,7 +456,13 @@ class RealTimeStreaming:
|
|||
raw_response = await self.backend_ws.recv() # type: ignore[assignment]
|
||||
|
||||
if self.provider_config:
|
||||
await self._handle_provider_config_message(raw_response)
|
||||
try:
|
||||
await self._handle_provider_config_message(raw_response)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"Error processing backend message, skipping: {e}"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
handled = await self._handle_raw_backend_message(raw_response)
|
||||
if handled:
|
||||
|
|
|
|||
|
|
@ -4678,6 +4678,14 @@ class BaseLLMHTTPHandler:
|
|||
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
|
||||
ssl=ssl_context,
|
||||
) as backend_ws:
|
||||
# Auto-send session setup if the provider requires it
|
||||
# (e.g. Gemini/Vertex AI Live needs a `setup` message before any realtime_input)
|
||||
_session_config: Optional[str] = None
|
||||
if provider_config.requires_session_configuration():
|
||||
_session_config = provider_config.session_configuration_request(model)
|
||||
if _session_config:
|
||||
await backend_ws.send(_session_config)
|
||||
|
||||
realtime_streaming = RealTimeStreaming(
|
||||
websocket,
|
||||
cast(ClientConnection, backend_ws),
|
||||
|
|
@ -4685,6 +4693,8 @@ class BaseLLMHTTPHandler:
|
|||
provider_config,
|
||||
model,
|
||||
)
|
||||
if _session_config:
|
||||
realtime_streaming.session_configuration_request = _session_config
|
||||
await realtime_streaming.bidirectional_forward()
|
||||
|
||||
except websockets.exceptions.InvalidStatusCode as e: # type: ignore
|
||||
|
|
|
|||
|
|
@ -226,35 +226,46 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
message_str = str(message)
|
||||
raise ValueError(f"Invalid JSON message: {message_str}")
|
||||
|
||||
## HANDLE SESSION UPDATE ##
|
||||
messages: List[str] = []
|
||||
if "type" in json_message and json_message["type"] == "session.update":
|
||||
msg_type = json_message.get("type")
|
||||
|
||||
## HANDLE SESSION UPDATE — translate to Gemini setup; no realtime_input needed ##
|
||||
if msg_type == "session.update":
|
||||
client_session_configuration_request = self.map_openai_params(
|
||||
optional_params={}, non_default_params=json_message["session"]
|
||||
)
|
||||
client_session_configuration_request["model"] = f"models/{model}"
|
||||
|
||||
messages.append(
|
||||
json.dumps(
|
||||
{
|
||||
"setup": client_session_configuration_request,
|
||||
}
|
||||
)
|
||||
json.dumps({"setup": client_session_configuration_request})
|
||||
)
|
||||
# elif session_configuration_request is None:
|
||||
# default_session_configuration_request = self.session_configuration_request(model)
|
||||
# messages.append(default_session_configuration_request)
|
||||
return messages
|
||||
|
||||
## HANDLE response.create — Gemini responds automatically; nothing to forward ##
|
||||
if msg_type == "response.create":
|
||||
return []
|
||||
|
||||
## HANDLE INPUT AUDIO BUFFER ##
|
||||
if (
|
||||
"type" in json_message
|
||||
and json_message["type"] == "input_audio_buffer.append"
|
||||
):
|
||||
if msg_type == "input_audio_buffer.append":
|
||||
realtime_input_dict["audio"] = HttpxBlobType(
|
||||
mimeType=self.get_audio_mime_type(), data=json_message["audio"]
|
||||
)
|
||||
## HANDLE conversation.item.create — extract actual user text ##
|
||||
elif msg_type == "conversation.item.create":
|
||||
item = json_message.get("item", {})
|
||||
content_list = item.get("content", [])
|
||||
text_parts = [
|
||||
c.get("text", "")
|
||||
for c in content_list
|
||||
if isinstance(c, dict) and c.get("type") == "input_text"
|
||||
]
|
||||
text = " ".join(filter(None, text_parts))
|
||||
if not text:
|
||||
return []
|
||||
realtime_input_dict["text"] = text
|
||||
else:
|
||||
realtime_input_dict["text"] = message
|
||||
# Unknown/unsupported OpenAI event type — drop silently rather than
|
||||
# forwarding raw JSON as text input to the model.
|
||||
return []
|
||||
|
||||
if len(realtime_input_dict) != 1:
|
||||
raise ValueError(
|
||||
|
|
@ -301,9 +312,17 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
if _system_instruction is not None and isinstance(_system_instruction, str):
|
||||
session["instructions"] = _system_instruction
|
||||
if _model is not None and isinstance(_model, str):
|
||||
session["model"] = _model.strip(
|
||||
"models/"
|
||||
) # keep it consistent with how openai returns the model name
|
||||
# Normalise to bare model name for OpenAI compatibility.
|
||||
# Vertex AI uses a full resource path:
|
||||
# projects/{project}/locations/{location}/publishers/google/models/{model}
|
||||
# Google AI Studio uses:
|
||||
# models/{model}
|
||||
if "/models/" in _model:
|
||||
session["model"] = _model.split("/models/")[-1]
|
||||
elif _model.startswith("models/"):
|
||||
session["model"] = _model[len("models/"):]
|
||||
else:
|
||||
session["model"] = _model
|
||||
|
||||
return OpenAIRealtimeStreamSessionEvents(
|
||||
type="session.created",
|
||||
|
|
@ -435,7 +454,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
if "text" in part:
|
||||
delta += part["text"]
|
||||
elif "inlineData" in part:
|
||||
delta += part["inlineData"]["data"]
|
||||
delta += part["inlineData"].get("data", "")
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error transforming content delta events: {e}, got message: {message}"
|
||||
|
|
@ -466,10 +485,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
delta = "".join([delta_chunk["delta"] for delta_chunk in delta_chunks])
|
||||
else:
|
||||
delta = ""
|
||||
if current_output_item_id is None or current_response_id is None:
|
||||
raise ValueError(
|
||||
"current_output_item_id and current_response_id cannot be None for a 'done' event."
|
||||
)
|
||||
if current_output_item_id is None:
|
||||
current_output_item_id = "item_{}".format(uuid.uuid4())
|
||||
if current_response_id is None:
|
||||
current_response_id = "resp_{}".format(uuid.uuid4())
|
||||
if delta_type == "text":
|
||||
return OpenAIRealtimeResponseTextDone(
|
||||
type="response.text.done",
|
||||
|
|
@ -503,10 +522,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
- return response.content_part.done
|
||||
- return response.output_item.done
|
||||
"""
|
||||
if current_output_item_id is None or current_response_id is None:
|
||||
raise ValueError(
|
||||
"current_output_item_id and current_response_id cannot be None for a 'done' event."
|
||||
)
|
||||
if current_output_item_id is None:
|
||||
current_output_item_id = "item_{}".format(uuid.uuid4())
|
||||
if current_response_id is None:
|
||||
current_response_id = "resp_{}".format(uuid.uuid4())
|
||||
returned_items: List[OpenAIRealtimeEvents] = []
|
||||
|
||||
delta_done_event_text = cast(Optional[str], delta_done_event.get("text"))
|
||||
|
|
@ -644,10 +663,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
output_items: Optional[List[OpenAIRealtimeOutputItemDone]],
|
||||
session_configuration_request: Optional[str] = None,
|
||||
) -> OpenAIRealtimeDoneEvent:
|
||||
if current_conversation_id is None or current_response_id is None:
|
||||
raise ValueError(
|
||||
f"current_conversation_id and current_response_id must all be set for a 'done' event. Got=current_conversation_id: {current_conversation_id}, current_response_id: {current_response_id}"
|
||||
)
|
||||
if current_conversation_id is None:
|
||||
current_conversation_id = "conv_{}".format(uuid.uuid4())
|
||||
if current_response_id is None:
|
||||
current_response_id = "resp_{}".format(uuid.uuid4())
|
||||
|
||||
if session_configuration_request:
|
||||
session_configuration_request_dict: BidiGenerateContentSetup = json.loads(
|
||||
|
|
@ -758,9 +777,14 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
)
|
||||
returned_message = [transformed_content_done_event]
|
||||
|
||||
# Use IDs from the done event — transform_content_done_event may have
|
||||
# generated UUID fallbacks when the originals were None.
|
||||
resolved_item_id = transformed_content_done_event.get("item_id") or current_output_item_id
|
||||
resolved_response_id = transformed_content_done_event.get("response_id") or current_response_id
|
||||
|
||||
additional_items = self.return_additional_content_done_events(
|
||||
current_output_item_id=current_output_item_id,
|
||||
current_response_id=current_response_id,
|
||||
current_output_item_id=resolved_item_id,
|
||||
current_response_id=resolved_response_id,
|
||||
delta_done_event=transformed_content_done_event,
|
||||
delta_type=delta_type,
|
||||
)
|
||||
|
|
|
|||
0
litellm/llms/vertex_ai/realtime/__init__.py
Normal file
0
litellm/llms/vertex_ai/realtime/__init__.py
Normal file
159
litellm/llms/vertex_ai/realtime/transformation.py
Normal file
159
litellm/llms/vertex_ai/realtime/transformation.py
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
"""
|
||||
Vertex AI Realtime (BidiGenerateContent) config.
|
||||
|
||||
Extends GeminiRealtimeConfig but adapts the WSS URL and auth header for the
|
||||
Vertex AI endpoint instead of Google AI Studio.
|
||||
|
||||
URL pattern:
|
||||
wss://{location}-aiplatform.googleapis.com/ws/
|
||||
google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent
|
||||
|
||||
Auth: OAuth2 Bearer token (not an API key).
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import List, Optional
|
||||
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
|
||||
class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
||||
"""
|
||||
Realtime config for Vertex AI (BidiGenerateContent).
|
||||
|
||||
``access_token`` and ``project`` must be pre-resolved by the caller
|
||||
(they require async I/O) and injected at construction time.
|
||||
"""
|
||||
|
||||
def __init__(self, access_token: str, project: str, location: str) -> None:
|
||||
self._access_token = access_token
|
||||
self._project = project
|
||||
self._location = location
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# URL
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_complete_url(
|
||||
self, api_base: Optional[str], model: str, api_key: Optional[str] = None # noqa: ARG002
|
||||
) -> str:
|
||||
"""
|
||||
Build the Vertex AI Live WSS endpoint URL.
|
||||
|
||||
If *api_base* is provided it overrides the default aiplatform host,
|
||||
allowing enterprise / VPC-SC deployments to point at a custom gateway.
|
||||
"""
|
||||
if api_base:
|
||||
# Allow callers to supply a fully-qualified wss:// base URL.
|
||||
base = api_base.rstrip("/")
|
||||
base = base.replace("https://", "wss://").replace("http://", "ws://")
|
||||
return f"{base}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
|
||||
location = self._location
|
||||
if location == "global":
|
||||
host = "aiplatform.googleapis.com"
|
||||
else:
|
||||
host = f"{location}-aiplatform.googleapis.com"
|
||||
|
||||
return f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Auth headers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str, # noqa: ARG002
|
||||
api_key: Optional[str] = None, # noqa: ARG002
|
||||
) -> dict:
|
||||
"""
|
||||
Return headers with a Bearer token for Vertex AI.
|
||||
|
||||
``api_key`` is intentionally ignored — Vertex AI uses OAuth2 tokens,
|
||||
not API keys. The token was resolved at config-construction time.
|
||||
"""
|
||||
headers = dict(headers)
|
||||
headers["Authorization"] = f"Bearer {self._access_token}"
|
||||
if self._project:
|
||||
headers["x-goog-user-project"] = self._project
|
||||
return headers
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Audio MIME type — Vertex AI needs the sample rate in the MIME string
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_audio_mime_type(self, input_audio_format: str = "pcm16") -> str:
|
||||
mime_types = {
|
||||
"pcm16": "audio/pcm;rate=16000",
|
||||
"g711_ulaw": "audio/pcmu",
|
||||
"g711_alaw": "audio/pcma",
|
||||
}
|
||||
return mime_types.get(input_audio_format, "application/octet-stream")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Session setup message
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def session_configuration_request(self, model: str) -> str:
|
||||
"""
|
||||
Return the JSON setup message for Vertex AI Live.
|
||||
|
||||
Vertex AI requires the fully-qualified model path:
|
||||
``projects/{project}/locations/{location}/publishers/google/models/{model}``
|
||||
|
||||
Also enables automatic activity detection (server VAD) and output
|
||||
audio transcription so the proxy forwards transcript events.
|
||||
"""
|
||||
from litellm.types.llms.gemini import BidiGenerateContentSetup
|
||||
from litellm.types.llms.vertex_ai import GeminiResponseModalities
|
||||
|
||||
response_modalities: list[GeminiResponseModalities] = ["AUDIO"]
|
||||
full_model_path = (
|
||||
f"projects/{self._project}"
|
||||
f"/locations/{self._location}"
|
||||
f"/publishers/google/models/{model}"
|
||||
)
|
||||
setup_config: BidiGenerateContentSetup = {
|
||||
"model": full_model_path,
|
||||
"generationConfig": {"responseModalities": response_modalities},
|
||||
# Enable server-side VAD with sensible defaults for voice sessions.
|
||||
"realtimeInputConfig": {
|
||||
"automaticActivityDetection": {
|
||||
"disabled": False,
|
||||
"silenceDurationMs": 800,
|
||||
}
|
||||
},
|
||||
# Return output transcript so clients can read what the model said.
|
||||
"outputAudioTranscription": {},
|
||||
}
|
||||
return json.dumps({"setup": setup_config})
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Request translation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def transform_realtime_request(
|
||||
self,
|
||||
message: str,
|
||||
model: str,
|
||||
session_configuration_request: Optional[str] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Translate OpenAI realtime client messages to Vertex AI format.
|
||||
|
||||
``session.update`` is intentionally ignored (returns []) because
|
||||
Vertex AI only accepts a single ``setup`` message at the start of
|
||||
the connection — sending a second one causes a 1007 close error.
|
||||
The initial setup (sent automatically before bidirectional_forward)
|
||||
already includes AUDIO modality and server VAD, so there is nothing
|
||||
more to configure.
|
||||
"""
|
||||
json_message = json.loads(message)
|
||||
if json_message.get("type") == "session.update":
|
||||
# Do not forward as a second setup — Vertex AI rejects it.
|
||||
return []
|
||||
|
||||
return super().transform_realtime_request(
|
||||
message, model, session_configuration_request
|
||||
)
|
||||
|
|
@ -19,6 +19,8 @@ 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
|
||||
from ..llms.openai.realtime.handler import OpenAIRealtime
|
||||
from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
|
||||
from ..llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from ..llms.xai.realtime.handler import XAIRealtime
|
||||
from ..utils import client as wrapper_client
|
||||
|
||||
|
|
@ -26,6 +28,7 @@ azure_realtime = AzureOpenAIRealtime()
|
|||
openai_realtime = OpenAIRealtime()
|
||||
bedrock_realtime = BedrockRealtime()
|
||||
xai_realtime = XAIRealtime()
|
||||
vertex_llm_base = VertexBase()
|
||||
base_llm_http_handler = BaseLLMHTTPHandler()
|
||||
|
||||
|
||||
|
|
@ -215,6 +218,52 @@ async def _arealtime(
|
|||
timeout=timeout,
|
||||
query_params=query_params,
|
||||
)
|
||||
elif _custom_llm_provider == "vertex_ai":
|
||||
vertex_credentials = (
|
||||
kwargs.get("vertex_credentials")
|
||||
or kwargs.get("vertex_ai_credentials")
|
||||
or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
)
|
||||
vertex_project = (
|
||||
kwargs.get("vertex_project")
|
||||
or kwargs.get("vertex_ai_project")
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_location = (
|
||||
kwargs.get("vertex_location")
|
||||
or kwargs.get("vertex_ai_location")
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
|
||||
resolved_location = vertex_llm_base.get_vertex_region(
|
||||
vertex_region=vertex_location, model=model
|
||||
)
|
||||
|
||||
access_token, resolved_project = await vertex_llm_base._ensure_access_token_async(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
vertex_realtime_config = VertexAIRealtimeConfig(
|
||||
access_token=access_token,
|
||||
project=resolved_project,
|
||||
location=resolved_location,
|
||||
)
|
||||
|
||||
await base_llm_http_handler.async_realtime(
|
||||
model=model,
|
||||
websocket=websocket,
|
||||
logging_obj=litellm_logging_obj,
|
||||
provider_config=vertex_realtime_config,
|
||||
api_base=dynamic_api_base or litellm_params.api_base,
|
||||
api_key=None,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
headers=headers,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported model: {model}")
|
||||
|
||||
|
|
@ -261,6 +310,33 @@ async def _realtime_health_check(
|
|||
url = xai_realtime._construct_url(
|
||||
api_base=api_base or "https://api.x.ai/v1", query_params={"model": model}
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
vertex_location = litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
|
||||
resolved_location = vertex_llm_base.get_vertex_region(
|
||||
vertex_region=vertex_location, model=model
|
||||
)
|
||||
access_token, resolved_project = await vertex_llm_base._ensure_access_token_async(
|
||||
credentials=None,
|
||||
project_id=litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
vertex_realtime_config = VertexAIRealtimeConfig(
|
||||
access_token=access_token,
|
||||
project=resolved_project,
|
||||
location=resolved_location,
|
||||
)
|
||||
url = vertex_realtime_config.get_complete_url(api_base=api_base, model=model)
|
||||
ssl_context = get_shared_realtime_ssl_context()
|
||||
headers = vertex_realtime_config.validate_environment(
|
||||
headers={}, model=model, api_key=None
|
||||
)
|
||||
async with websockets.connect( # type: ignore
|
||||
url,
|
||||
additional_headers=headers,
|
||||
max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
|
||||
ssl=ssl_context,
|
||||
):
|
||||
return True
|
||||
else:
|
||||
raise ValueError(f"Unsupported model: {model}")
|
||||
ssl_context = get_shared_realtime_ssl_context()
|
||||
|
|
|
|||
|
|
@ -1066,7 +1066,8 @@
|
|||
"fine_tuning": true,
|
||||
"rag_ingest": true,
|
||||
"rag_query": true,
|
||||
"generateContent": true
|
||||
"generateContent": true,
|
||||
"realtime": true
|
||||
}
|
||||
},
|
||||
"gemini": {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,224 @@
|
|||
"""
|
||||
Unit tests for VertexAIRealtimeConfig.
|
||||
|
||||
Validates:
|
||||
- URL construction (regional and global)
|
||||
- Auth headers (Bearer token + project header)
|
||||
- Session setup message format
|
||||
- Full text-in / text-out round-trip via RealTimeStreaming with a mocked
|
||||
WebSocket pair (no real network calls)
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
import websockets.exceptions # registers websockets.exceptions on the websockets namespace
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config unit tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_complete_url_regional():
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="tok", project="my-proj", location="us-central1"
|
||||
)
|
||||
url = cfg.get_complete_url(api_base=None, model="gemini-2.0-flash-live-001")
|
||||
assert url == (
|
||||
"wss://us-central1-aiplatform.googleapis.com"
|
||||
"/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
|
||||
|
||||
def test_get_complete_url_global():
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="tok", project="my-proj", location="global"
|
||||
)
|
||||
url = cfg.get_complete_url(api_base=None, model="gemini-2.0-flash-live-001")
|
||||
assert url == (
|
||||
"wss://aiplatform.googleapis.com"
|
||||
"/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
|
||||
|
||||
def test_get_complete_url_custom_api_base():
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="tok", project="my-proj", location="us-central1"
|
||||
)
|
||||
url = cfg.get_complete_url(
|
||||
api_base="https://custom-gateway.example.com",
|
||||
model="gemini-2.0-flash-live-001",
|
||||
)
|
||||
assert url.startswith("wss://custom-gateway.example.com")
|
||||
assert "BidiGenerateContent" in url
|
||||
|
||||
|
||||
def test_validate_environment_sets_bearer_and_project():
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="mytoken", project="proj-123", location="us-central1"
|
||||
)
|
||||
headers = cfg.validate_environment(
|
||||
headers={}, model="gemini-2.0-flash-live-001", api_key=None
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer mytoken"
|
||||
assert headers["x-goog-user-project"] == "proj-123"
|
||||
|
||||
|
||||
def test_session_configuration_request_model_format():
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="tok", project="my-proj", location="us-central1"
|
||||
)
|
||||
raw = cfg.session_configuration_request("gemini-2.0-flash-live-001")
|
||||
parsed = json.loads(raw)
|
||||
assert parsed["setup"]["model"] == (
|
||||
"projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-001"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Round-trip test: text-in / text-out via RealTimeStreaming
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Minimal Gemini BidiGenerateContent message sequence:
|
||||
# server → setupComplete
|
||||
# client → conversation.item.create (OpenAI format, translated by config)
|
||||
# server → serverContent with modelTurn text delta
|
||||
# server → serverContent with generationComplete
|
||||
|
||||
SETUP_COMPLETE = json.dumps({"setupComplete": {}})
|
||||
|
||||
SERVER_TEXT_DELTA = json.dumps(
|
||||
{
|
||||
"serverContent": {
|
||||
"modelTurn": {
|
||||
"parts": [{"text": "Hello from Vertex AI!"}]
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# generationComplete fires RESPONSE_TEXT_DONE; turnComplete fires RESPONSE_DONE
|
||||
# They must be separate messages (the transformer processes one top-level key per message).
|
||||
SERVER_GENERATION_COMPLETE = json.dumps(
|
||||
{"serverContent": {"generationComplete": True}}
|
||||
)
|
||||
|
||||
SERVER_TURN_COMPLETE = json.dumps(
|
||||
{"serverContent": {"turnComplete": True}}
|
||||
)
|
||||
|
||||
# OpenAI-format text message the client sends
|
||||
CLIENT_TEXT_MESSAGE = json.dumps(
|
||||
{
|
||||
"type": "conversation.item.create",
|
||||
"item": {
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "Say hello"}],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_realtime_text_in_text_out():
|
||||
"""
|
||||
Simulate a full text-in / text-out session through RealTimeStreaming using
|
||||
VertexAIRealtimeConfig for message translation. All I/O is mocked.
|
||||
"""
|
||||
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
||||
|
||||
cfg = VertexAIRealtimeConfig(
|
||||
access_token="fake-token",
|
||||
project="fake-project",
|
||||
location="us-central1",
|
||||
)
|
||||
|
||||
# --- mock client WebSocket (FastAPI side) ---
|
||||
client_ws = MagicMock()
|
||||
client_ws.exceptions = MagicMock()
|
||||
client_ws.exceptions.ConnectionClosed = Exception
|
||||
|
||||
sent_to_client: list[str] = []
|
||||
|
||||
async def _client_send_text(data: str):
|
||||
sent_to_client.append(data)
|
||||
|
||||
client_ws.send_text = AsyncMock(side_effect=_client_send_text)
|
||||
|
||||
# Client sends one text message then raises to end the loop
|
||||
client_ws.receive_text = AsyncMock(
|
||||
side_effect=[CLIENT_TEXT_MESSAGE, Exception("client done")]
|
||||
)
|
||||
|
||||
# --- mock backend WebSocket (Vertex AI side) ---
|
||||
backend_ws = MagicMock()
|
||||
|
||||
upstream_messages = [
|
||||
SETUP_COMPLETE,
|
||||
SERVER_TEXT_DELTA,
|
||||
SERVER_GENERATION_COMPLETE,
|
||||
SERVER_TURN_COMPLETE,
|
||||
]
|
||||
|
||||
async def _backend_recv(decode=True): # noqa: ARG001
|
||||
if not upstream_messages:
|
||||
# Signal normal connection close so the loop exits cleanly
|
||||
raise websockets.exceptions.ConnectionClosedOK(None, None) # type: ignore[arg-type]
|
||||
return upstream_messages.pop(0)
|
||||
|
||||
backend_ws.recv = AsyncMock(side_effect=_backend_recv)
|
||||
|
||||
sent_to_backend: list[str] = []
|
||||
|
||||
async def _backend_send(data: str):
|
||||
sent_to_backend.append(data)
|
||||
|
||||
backend_ws.send = AsyncMock(side_effect=_backend_send)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_trace_id = "test-trace-id"
|
||||
logging_obj.pre_call = MagicMock()
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.success_handler = MagicMock()
|
||||
|
||||
streaming = RealTimeStreaming(
|
||||
websocket=client_ws,
|
||||
backend_ws=backend_ws,
|
||||
logging_obj=logging_obj,
|
||||
provider_config=cfg,
|
||||
model="gemini-2.0-flash-live-001",
|
||||
)
|
||||
|
||||
# Run backend→client forwarding for the three queued messages, then stop.
|
||||
# We don't run client_ack_messages here to avoid the blocking receive loop.
|
||||
await streaming.backend_to_client_send_messages()
|
||||
|
||||
# --- Assertions ---
|
||||
|
||||
# session.created should have been forwarded to client
|
||||
session_created_msgs = [
|
||||
m for m in sent_to_client if '"session.created"' in m
|
||||
]
|
||||
assert session_created_msgs, "Expected session.created to be sent to client"
|
||||
|
||||
# At least one text delta should have been forwarded
|
||||
text_delta_msgs = [
|
||||
m for m in sent_to_client if '"response.text.delta"' in m
|
||||
]
|
||||
assert text_delta_msgs, "Expected response.text.delta to be sent to client"
|
||||
|
||||
# Verify the delta contains the model's text
|
||||
delta_obj = json.loads(text_delta_msgs[0])
|
||||
assert "Hello from Vertex AI!" in delta_obj.get("delta", "")
|
||||
|
||||
# response.done should have been forwarded
|
||||
done_msgs = [m for m in sent_to_client if '"response.done"' in m]
|
||||
assert done_msgs, "Expected response.done to be sent to client"
|
||||
Loading…
Add table
Reference in a new issue