""" 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 Final from litellm import verbose_logger 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: super().__init__() self._access_token = access_token self._project = project self._location = location def _include_function_response_id(self) -> bool: return False # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ def get_complete_url( self, api_base: str | None, model: str, api_key: str | None = 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: Final = 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: str | None = 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: Final = { # Gemini Live native audio (OpenAI GA realtime default) is 24kHz PCM. "pcm16": "audio/pcm;rate=24000", "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: Final[list[GeminiResponseModalities]] = ["AUDIO"] full_model_path: Final = f"projects/{self._project}/locations/{self._location}/publishers/google/models/{model}" setup_config: Final[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 input transcript so guardrails can inspect user speech. "inputAudioTranscription": {}, # Return output transcript so clients can read what the model said. "outputAudioTranscription": {}, } return json.dumps({"setup": setup_config}) # ------------------------------------------------------------------ # Request translation # ------------------------------------------------------------------ def _vertex_model_path(self, model: str) -> str: """Return the fully-qualified Vertex AI model resource path.""" return f"projects/{self._project}/locations/{self._location}/publishers/google/models/{model}" def _build_vertex_ai_setup_config(self, model: str, session_params: dict) -> dict: """Build Vertex AI setup configuration with proper model path and defaults.""" # Normalize GA-remapped fields (``output_modalities``, nested # ``audio.input.transcription``, ``audio.input.turn_detection``) back to # their flat beta keys so ``map_openai_params`` picks them up. Without # this, GA clients' explicit modality / transcription / turn-detection # settings would be silently dropped because ``map_openai_params`` only # recognises the flat OpenAI-beta key names. session_params = self._normalize_session_payload_for_mapping(session_params) setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params) # Use full Vertex AI model path setup_config["model"] = self._vertex_model_path(model) # Add Vertex AI specific defaults if not provided generation_config: Final = setup_config.setdefault("generationConfig", {}) generation_config.setdefault("responseModalities", ["AUDIO"]) # Ensure Vertex defaults for realtimeInputConfig apply even when # the client provided a partial ``turn_detection`` (e.g. only # ``silence_duration_ms``). ``map_automatic_turn_detection`` sets # ``disabled=True`` whenever ``create_response`` is absent or # ``False``. Force ``disabled=False`` only when the client did # not explicitly request ``create_response: False`` — that path # is how transcription guardrails suppress automatic responses, # and overriding it here would silently bypass the guardrail. # Vertex Live has no "VAD on, no auto-response" mode, so callers # that need that behaviour must accept that VAD is off. client_turn_detection: Final = session_params.get("turn_detection") client_disabled_auto_response: Final = ( isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False ) realtime_input_config: Final = setup_config.setdefault("realtimeInputConfig", {}) automatic_detection: Final = realtime_input_config.setdefault("automaticActivityDetection", {}) if not client_disabled_auto_response: automatic_detection["disabled"] = False automatic_detection.setdefault("silenceDurationMs", 800) setup_config.setdefault("inputAudioTranscription", {}) setup_config.setdefault("outputAudioTranscription", {}) return self._finalize_gemini_live_setup(model, setup_config) def transform_realtime_request( self, message: str, model: str, session_configuration_request: str | None = None, ) -> list[str]: """ Translate OpenAI realtime client messages to Vertex AI format. On the first ``session.update`` (when no setup has been sent yet) the full ``BidiGenerateContentSetup`` is built with Vertex AI's model path and forwarded. Any later ``session.update`` is dropped: Vertex AI documents ``setup`` as the first-and-only client message, and a second ``setup`` closes the connection with a 1007 policy error. """ json_message: Final = json.loads(message) msg_type: Final = json_message.get("type") if msg_type == "session.update": if session_configuration_request is None: setup_config: Final = self._build_vertex_ai_setup_config(model, json_message.get("session") or {}) gemini_setup_msg: Final = json.dumps({"setup": setup_config}) verbose_logger.debug("Vertex AI Realtime: Sending initial setup with tools to backend") return [gemini_setup_msg] # A follow-up session.update can't be forwarded as a second setup # (Vertex Live closes the WebSocket with 1007). If this drop is # silencing the audio-transcription guardrail's create_response # disable, surface a warning so operators know the model will # auto-respond before the guardrail can gate it on Vertex AI. client_turn_detection = GeminiRealtimeConfig._extract_turn_detection(json_message.get("session") or {}) if isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False: verbose_logger.warning( "Vertex AI Realtime: Dropping subsequent session.update " "(turn_detection.create_response=False) — Vertex Live " "rejects a second setup message. Audio-transcription " "guardrails cannot suppress the model's auto-response on " "Vertex AI in non-deferred mode." ) else: verbose_logger.debug("Vertex AI Realtime: Ignoring session.update (setup already sent)") return [] return super().transform_realtime_request(message, model, session_configuration_request)