diff --git a/docs/my-website/docs/providers/vertex_realtime.md b/docs/my-website/docs/providers/vertex_realtime.md new file mode 100644 index 00000000000..00db682a0d7 --- /dev/null +++ b/docs/my-website/docs/providers/vertex_realtime.md @@ -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). diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 88d740b5188..b01bb53cfe7 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -758,6 +758,7 @@ const sidebars = { "providers/vertex_batch", "providers/vertex_ocr", "providers/vertex_ai_agent_engine", + "providers/vertex_realtime", ] }, { diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index d15d23f8eea..5d7a5bfe318 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -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: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7267532933d..b09a36be60f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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 diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 62329358e47..2e0e678e69f 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -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, ) diff --git a/litellm/llms/vertex_ai/realtime/__init__.py b/litellm/llms/vertex_ai/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py new file mode 100644 index 00000000000..eaa9844f108 --- /dev/null +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -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 + ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index f1cc9b1d977..e4c8f648190 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -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() diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 328398a296a..8834d8b19c0 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1066,7 +1066,8 @@ "fine_tuning": true, "rag_ingest": true, "rag_query": true, - "generateContent": true + "generateContent": true, + "realtime": true } }, "gemini": { diff --git a/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py new file mode 100644 index 00000000000..9145896647e --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -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"