Update app/main.py

app/main.py

@@ -1,7 +1,8 @@
+# app/main.py
 import os
 import json
 import asyncio
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, List
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.responses import JSONResponse
@@ -10,25 +11,31 @@ import websockets
 
 load_dotenv()
 
-app = FastAPI(title="Gemini Live WS Proxy", version="1.0.0")
+app = FastAPI(title="Gemini Live Native-Audio WS Proxy", version="2.0.0")
 
-# Gemini Live API WebSocket endpoint for BidiGenerateContent (v1beta)
-# (Official endpoint in the Live API WebSockets reference.)
+# Gemini Live API WebSocket endpoint (v1beta, BidiGenerateContent)
 GEMINI_LIVE_WS_URL = (
     "wss://generativelanguage.googleapis.com/ws/"
     "google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent"
 )
 
-DEFAULT_MODEL = os.getenv("GEMINI_MODEL", "models/gemini-2.5-flash-native-audio-preview-12-2025")
-DEFAULT_SYSTEM = os.getenv("GEMINI_SYSTEM_INSTRUCTION", "You are a helpful assistant.")
+API_KEY = os.getenv("GEMINI_API_KEY", "").strip()
+
+# Defaults (override via HF Space variables)
+DEFAULT_MODEL = os.getenv("GEMINI_MODEL", "models/gemini-2.0-flash-live-001")
+DEFAULT_SYSTEM = os.getenv("GEMINI_SYSTEM_INSTRUCTION", "You are a helpful assistant for a school coding club.")
 DEFAULT_TEMPERATURE = float(os.getenv("GEMINI_TEMPERATURE", "0.7"))
 DEFAULT_MAX_TOKENS = int(os.getenv("GEMINI_MAX_OUTPUT_TOKENS", "1024"))
 
-API_KEY = os.getenv("GEMINI_API_KEY", "").strip()
-if not API_KEY:
-    # Don't crash import-time on HF if they haven't set secrets yet;
-    # we will return a clear runtime error at connection time.
-    pass
+# Native-audio config defaults
+DEFAULT_VOICE = os.getenv("GEMINI_VOICE_NAME", "Kore")
+# input audio: most common is 16k PCM16 mono
+DEFAULT_INPUT_RATE = int(os.getenv("GEMINI_INPUT_AUDIO_RATE", "16000"))
+# output audio: docs commonly mention 24k PCM16
+DEFAULT_OUTPUT_RATE = int(os.getenv("GEMINI_OUTPUT_AUDIO_RATE", "24000"))
+
+# Debug passthrough (set to "1" to enable)
+DEBUG_GEMINI_RAW = os.getenv("DEBUG_GEMINI_RAW", "0").strip() == "1"
 
 
 @app.get("/health")
@@ -39,75 +46,59 @@ async def health():
             "ok": ok,
             "has_api_key": ok,
             "model": DEFAULT_MODEL,
+            "voice": DEFAULT_VOICE,
+            "input_rate": DEFAULT_INPUT_RATE,
+            "output_rate": DEFAULT_OUTPUT_RATE,
+            "debug_raw": DEBUG_GEMINI_RAW,
         }
     )
 
 
-def _safe_get_text_from_content(content: Dict[str, Any]) -> str:
-    """
-    Gemini Content format typically includes:
-      {"role": "...", "parts": [{"text": "..."} , ...]}
-    We concatenate any text parts we see.
-    """
+def _extract_text_parts(content: Dict[str, Any]) -> str:
     parts = content.get("parts") or []
-    out = []
+    out: List[str] = []
     for p in parts:
-        if isinstance(p, dict) and "text" in p and isinstance(p["text"], str):
+        if isinstance(p, dict) and isinstance(p.get("text"), str):
             out.append(p["text"])
     return "".join(out)
 
 
-async def _gemini_connect_and_setup(
-    model: str,
-    system_instruction: str,
-    temperature: float,
-    max_output_tokens: int,
-    response_modalities: Optional[list] = None,
-):
+def _extract_inline_audio_parts(content: Dict[str, Any]) -> List[Dict[str, str]]:
     """
-    Opens a websocket to Gemini Live API and sends the required initial setup message.
-    Clients should wait for setupComplete before sending further messages.
+    Returns list of {"mime": "...", "data": "base64..."} for any inlineData parts.
     """
-    headers = {
-        # Gemini API auth: x-goog-api-key header is required for requests. :contentReference[oaicite:2]{index=2}
-        "x-goog-api-key": API_KEY,
-    }
-
+    parts = content.get("parts") or []
+    out: List[Dict[str, str]] = []
+    for p in parts:
+        if not isinstance(p, dict):
+            continue
+        inline = p.get("inlineData")
+        if isinstance(inline, dict):
+            data = inline.get("data")
+            mime = inline.get("mimeType")
+            if isinstance(data, str) and isinstance(mime, str):
+                out.append({"mime": mime, "data": data})
+    return out
+
+
+async def _gemini_ws_connect(setup_payload: Dict[str, Any]):
+    headers = {"x-goog-api-key": API_KEY}
     ws = await websockets.connect(
         GEMINI_LIVE_WS_URL,
         extra_headers=headers,
-        max_size=8 * 1024 * 1024,  # allow larger payloads if needed later
+        max_size=16 * 1024 * 1024,
         ping_interval=20,
         ping_timeout=20,
     )
 
-    setup_payload = {
-        "setup": {
-            "model": model,
-            "generationConfig": {
-                "temperature": temperature,
-                "maxOutputTokens": max_output_tokens,
-                "responseModalities": response_modalities or ["TEXT"],
-            },
-            # Live API reference shows systemInstruction is Content; we send text-only Content.
-            # (Docs note text parts in system instruction.) :contentReference[oaicite:3]{index=3}
-            "systemInstruction": {
-                "role": "system",
-                "parts": [{"text": system_instruction}],
-            },
-        }
-    }
-
     await ws.send(json.dumps(setup_payload))
 
-    # Wait for setupComplete
+    # wait for setupComplete
     while True:
         raw = await ws.recv()
         msg = json.loads(raw)
         if "setupComplete" in msg:
             return ws
-        # Forward other early messages if they appear, but don't block setup forever.
-        # If Gemini returns an error-like structure, surface it.
         if "error" in msg:
             raise RuntimeError(f"Gemini setup error: {msg['error']}")
 
@@ -115,120 +106,192 @@ async def _gemini_connect_and_setup(
 @app.websocket("/ws")
 async def ws_proxy(client_ws: WebSocket):
     """
-    Client protocol (simple):
-      -> {"type":"text","text":"hello"}
-      -> {"type":"configure", "model": "...", "system_instruction": "...", "temperature": 0.7, "max_output_tokens": 1024}
-      -> {"type":"close"}
+    Client protocol (native-audio + VAD friendly):
+      -> {"type":"configure", "model": "...", "system_instruction": "...", "temperature": 0.7,
+          "max_output_tokens": 1024, "voice": "Kore", "input_rate": 16000}
+         (optional, must be first; else defaults are used)
 
-    Server sends:
+      -> {"type":"audio","data":"<base64 pcm16 mono>","rate":16000}
+         (send repeatedly while user is speaking)
+
+      -> {"type":"audio_end"}
+         (send when VAD decides user stopped speaking; triggers assistant response)
+
+      -> {"type":"text","text":"..."}  (optional helper; NOT the main mode for native audio)
+
+    Server -> client:
       <- {"type":"ready"}
-      <- {"type":"text_delta","text":"..."}   (streaming)
+      <- {"type":"text_delta","text":"..."}          (assistant text parts, if any)
+      <- {"type":"audio_delta","mime":"...","data":"..."} (assistant audio chunks)
       <- {"type":"turn_complete"}
-      <- {"type":"gemini_raw","message":{...}} (debug passthrough)
       <- {"type":"error","message":"..."}
+      <- {"type":"gemini_raw","message":{...}}       (only if DEBUG_GEMINI_RAW=1)
     """
     await client_ws.accept()
 
     if not API_KEY:
-        await client_ws.send_text(
-            json.dumps(
-                {
-                    "type": "error",
-                    "message": "Server missing GEMINI_API_KEY env var. Set it in your Space secrets.",
-                }
-            )
-        )
+        await client_ws.send_text(json.dumps({"type": "error", "message": "Missing GEMINI_API_KEY on server."}))
         await client_ws.close(code=1011)
         return
 
-    # Per-connection defaults (can be overridden by configure message)
-    model = DEFAULT_MODEL
-    system_instruction = DEFAULT_SYSTEM
-    temperature = DEFAULT_TEMPERATURE
-    max_output_tokens = DEFAULT_MAX_TOKENS
+    # --- Phase 1: accept optional configure before connecting to Gemini ---
+    cfg = {
+        "model": DEFAULT_MODEL,
+        "system_instruction": DEFAULT_SYSTEM,
+        "temperature": DEFAULT_TEMPERATURE,
+        "max_output_tokens": DEFAULT_MAX_TOKENS,
+        "voice": DEFAULT_VOICE,
+        "input_rate": DEFAULT_INPUT_RATE,
+    }
+
+    async def _wait_for_optional_config(timeout_s: float = 1.2):
+        try:
+            raw = await asyncio.wait_for(client_ws.receive_text(), timeout=timeout_s)
+        except asyncio.TimeoutError:
+            return
+        except Exception:
+            return
+
+        data = json.loads(raw)
+        if data.get("type") != "configure":
+            # if first message is not configure, we treat it as "not configure"
+            # and stash it for later by putting it into a queue (simple: handle inline)
+            return data
+
+        # apply config
+        if isinstance(data.get("model"), str) and data["model"].strip():
+            cfg["model"] = data["model"].strip()
+        if isinstance(data.get("system_instruction"), str) and data["system_instruction"].strip():
+            cfg["system_instruction"] = data["system_instruction"].strip()
+        if data.get("temperature") is not None:
+            try:
+                cfg["temperature"] = float(data["temperature"])
+            except Exception:
+                pass
+        if data.get("max_output_tokens") is not None:
+            try:
+                cfg["max_output_tokens"] = int(data["max_output_tokens"])
+            except Exception:
+                pass
+        if isinstance(data.get("voice"), str) and data["voice"].strip():
+            cfg["voice"] = data["voice"].strip()
+        if data.get("input_rate") is not None:
+            try:
+                cfg["input_rate"] = int(data["input_rate"])
+            except Exception:
+                pass
+
+        await client_ws.send_text(json.dumps({"type": "configured"}))
+        return None
+
+    first_non_config = await _wait_for_optional_config()
+
+    # --- Phase 2: connect to Gemini with native-audio setup ---
+    # NOTE: For native-audio models, AUDIO modality is required.
+    setup_payload = {
+        "setup": {
+            "model": cfg["model"],
+            "generationConfig": {
+                "temperature": cfg["temperature"],
+                "maxOutputTokens": cfg["max_output_tokens"],
+                "responseModalities": ["AUDIO"],
+                "speechConfig": {
+                    "voiceConfig": {
+                        "prebuiltVoiceConfig": {
+                            "voiceName": cfg["voice"],
+                        }
+                    }
+                },
+            },
+            # Enable transcripts so Scratch can display text while audio plays
+            "inputAudioTranscription": {},
+            "outputAudioTranscription": {},
+            "systemInstruction": {
+                "role": "system",
+                "parts": [{"text": cfg["system_instruction"]}],
+            },
+        }
+    }
 
-    gemini_ws = None
     stop_event = asyncio.Event()
+    gemini_ws = None
 
-    async def ensure_gemini():
-        nonlocal gemini_ws
-        if gemini_ws is None:
-            gemini_ws = await _gemini_connect_and_setup(
-                model=model,
-                system_instruction=system_instruction,
-                temperature=temperature,
-                max_output_tokens=max_output_tokens,
-                response_modalities=["TEXT"],
-            )
+    try:
+        gemini_ws = await _gemini_ws_connect(setup_payload)
+        await client_ws.send_text(json.dumps({"type": "ready"}))
+    except Exception as e:
+        await client_ws.send_text(json.dumps({"type": "error", "message": f"Gemini setup failed: {e}"}))
+        await client_ws.close(code=1011)
+        return
+
+    # If we consumed a non-config first message, we need to handle it.
+    pending_first = first_non_config
 
     async def forward_client_to_gemini():
-        """
-        Reads from your client WebSocket and sends appropriate Live API messages to Gemini.
-        Uses clientContent + turnComplete for clean text turns. :contentReference[oaicite:4]{index=4}
-        """
+        nonlocal pending_first
         try:
             while not stop_event.is_set():
-                raw = await client_ws.receive_text()
-                data = json.loads(raw)
-
-                msg_type = data.get("type")
-                if msg_type == "configure":
-                    # Allow config BEFORE Gemini connection is created.
-                    if gemini_ws is not None:
-                        await client_ws.send_text(
-                            json.dumps(
-                                {
-                                    "type": "error",
-                                    "message": "Cannot configure after session started. Open a new WS connection.",
-                                }
-                            )
-                        )
-                        continue
+                if pending_first is not None:
+                    data = pending_first
+                    pending_first = None
+                else:
+                    raw = await client_ws.receive_text()
+                    data = json.loads(raw)
 
-                    model = data.get("model", model)
-                    system_instruction = data.get("system_instruction", system_instruction)
-                    temperature = float(data.get("temperature", temperature))
-                    max_output_tokens = int(data.get("max_output_tokens", max_output_tokens))
-                    await client_ws.send_text(json.dumps({"type": "configured"}))
-                    continue
+                t = data.get("type")
 
-                if msg_type == "close":
+                if t == "close":
                     stop_event.set()
                     return
 
-                if msg_type == "text":
-                    text = data.get("text", "")
-                    if not isinstance(text, str) or not text.strip():
+                if t == "audio":
+                    # expects base64 PCM16 mono
+                    b64 = data.get("data")
+                    rate = data.get("rate", cfg["input_rate"])
+                    if not isinstance(b64, str) or not b64:
                         continue
+                    try:
+                        rate_i = int(rate)
+                    except Exception:
+                        rate_i = cfg["input_rate"]
 
-                    await ensure_gemini()
-
-                    # Send a single "turn" using clientContent.turns and turnComplete=true. :contentReference[oaicite:5]{index=5}
                     payload = {
-                        "clientContent": {
-                            "turns": [
-                                {
-                                    "role": "user",
-                                    "parts": [{"text": text}],
-                                }
-                            ],
-                            "turnComplete": True,
+                        "realtimeInput": {
+                            "audio": {
+                                "data": b64,
+                                "mimeType": f"audio/pcm;rate={rate_i}",
+                            }
                         }
                     }
                     await gemini_ws.send(json.dumps(payload))
                     continue
 
-                # Optional: raw passthrough (advanced users)
-                if msg_type == "live_raw":
-                    await ensure_gemini()
-                    payload = data.get("payload")
-                    if isinstance(payload, dict):
+                if t == "audio_end":
+                    # tell Gemini the input stream ended for this turn
+                    payload = {"realtimeInput": {"audioStreamEnd": True}}
+                    await gemini_ws.send(json.dumps(payload))
+                    continue
+
+                if t == "text":
+                    # Optional helper: send text as a turn (some native-audio sessions still accept it),
+                    # but for voice-first you should mainly use audio.
+                    text = data.get("text", "")
+                    if isinstance(text, str) and text.strip():
+                        payload = {
+                            "clientContent": {
+                                "turns": [{"role": "user", "parts": [{"text": text.strip()}]}],
+                                "turnComplete": True,
+                            }
+                        }
                         await gemini_ws.send(json.dumps(payload))
                     continue
 
-                await client_ws.send_text(
-                    json.dumps({"type": "error", "message": f"Unknown message type: {msg_type}"})
-                )
+                # Advanced passthrough
+                if t == "live_raw" and isinstance(data.get("payload"), dict):
+                    await gemini_ws.send(json.dumps(data["payload"]))
+                    continue
+
+                await client_ws.send_text(json.dumps({"type": "error", "message": f"Unknown message type: {t}"}))
 
         except WebSocketDisconnect:
             stop_event.set()
@@ -240,39 +303,40 @@ async def ws_proxy(client_ws: WebSocket):
                 pass
 
     async def forward_gemini_to_client():
-        """
-        Reads Gemini Live API server messages and forwards useful pieces to your client.
-        We extract text from serverContent.modelTurn.parts[].text when present. :contentReference[oaicite:6]{index=6}
-        """
         try:
-            await ensure_gemini()
-            await client_ws.send_text(json.dumps({"type": "ready"}))
-
             while not stop_event.is_set():
                 raw = await gemini_ws.recv()
                 msg = json.loads(raw)
 
-                # Optional debug passthrough:
-                await client_ws.send_text(json.dumps({"type": "gemini_raw", "message": msg}))
+                if DEBUG_GEMINI_RAW:
+                    await client_ws.send_text(json.dumps({"type": "gemini_raw", "message": msg}))
 
-                # The main streaming content arrives under "serverContent"
                 server_content = msg.get("serverContent")
                 if isinstance(server_content, dict):
-                    # modelTurn is Content (role+parts)
                     model_turn = server_content.get("modelTurn")
                     if isinstance(model_turn, dict):
-                        delta = _safe_get_text_from_content(model_turn)
-                        if delta:
-                            await client_ws.send_text(json.dumps({"type": "text_delta", "text": delta}))
+                        # text parts
+                        txt = _extract_text_parts(model_turn)
+                        if txt:
+                            await client_ws.send_text(json.dumps({"type": "text_delta", "text": txt}))
+
+                        # audio parts (inlineData)
+                        audios = _extract_inline_audio_parts(model_turn)
+                        for a in audios:
+                            await client_ws.send_text(
+                                json.dumps({"type": "audio_delta", "mime": a["mime"], "data": a["data"]})
+                            )
+
+                    # Some implementations also include transcription fields; pass through if present
+                    out_tx = server_content.get("outputTranscription")
+                    if isinstance(out_tx, dict) and isinstance(out_tx.get("text"), str):
+                        await client_ws.send_text(
+                            json.dumps({"type": "output_transcript_delta", "text": out_tx["text"]})
+                        )
 
-                    # When generationComplete true, we end the turn
                     if server_content.get("generationComplete") is True:
                         await client_ws.send_text(json.dumps({"type": "turn_complete"}))
 
-                # Tool calls (if you later enable tools in setup)
-                if "toolCall" in msg:
-                    await client_ws.send_text(json.dumps({"type": "tool_call", "toolCall": msg["toolCall"]}))
-
                 if "goAway" in msg:
                     await client_ws.send_text(json.dumps({"type": "go_away", "goAway": msg["goAway"]}))
 
@@ -284,7 +348,6 @@ async def ws_proxy(client_ws: WebSocket):
                 pass
 
     try:
-        # Run both directions
         await asyncio.gather(forward_client_to_gemini(), forward_gemini_to_client())
     finally:
         stop_event.set()