Added new endpoints for the gemini live feature

2026-04-02 14:50:03 +08:00
parent 622cf89211
commit cfec9c9bf5
4 changed files with 340 additions and 1 deletions
@@ -0,0 +1,85 @@
+from fastapi import APIRouter, Depends, Request, HTTPException
+from app.api.deps import get_current_module
+from app.models.module import Module
+from sqlalchemy.orm import Session
+from app.core.database import get_db
+from app.core.limiter import limiter
+from app.core.config import settings
+from pydantic import BaseModel
+import httpx
+
+router = APIRouter()
+
+GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"
+GEMINI_LIVE_WSS_URI = (
+    "wss://generativelanguage.googleapis.com/ws/"
+    "google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent"
+)
+
+
+class LiveTokenRequest(BaseModel):
+    model: str = "gemini-2.0-flash-live-001"
+    system_instruction: str
+    voice_name: str | None = "Puck"
+
+
+@router.post("/live-token")
+@limiter.limit(settings.RATE_LIMIT)
+async def generate_live_token(
+    request: Request,
+    body: LiveTokenRequest,
+    module: Module = Depends(get_current_module),
+    db: Session = Depends(get_db),
+):
+    if not settings.GOOGLE_API_KEY or settings.GOOGLE_API_KEY == "your-google-api-key":
+        return {
+            "token": "mock-ephemeral-token",
+            "expires_at": "2026-04-02T02:00:00Z",
+            "websocket_uri": GEMINI_LIVE_WSS_URI,
+        }
+
+    payload = {
+        "systemInstruction": {
+            "parts": [{"text": body.system_instruction}]
+        },
+        "generationConfig": {
+            "responseModalities": ["AUDIO", "TEXT"],
+        },
+    }
+
+    if body.voice_name:
+        payload["generationConfig"]["speechConfig"] = {
+            "voiceConfig": {
+                "prebuiltVoiceConfig": {"voiceName": body.voice_name}
+            }
+        }
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        response = await client.post(
+            f"{GEMINI_API_BASE}/models/{body.model}:generateEphemeralToken",
+            json=payload,
+            params={"key": settings.GOOGLE_API_KEY},
+        )
+
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=502,
+            detail=f"Gemini token API error {response.status_code}: {response.text}",
+        )
+
+    data = response.json()
+
+    # Response name is "ephemeralTokens/{token_value}"
+    token_value = data.get("name", "").split("/")[-1]
+    expires_at = data.get("expireTime", "")
+
+    if module:
+        module.ingress_tokens += 1
+        module.total_tokens += 1
+        db.commit()
+
+    return {
+        "token": token_value,
+        "expires_at": expires_at,
+        "websocket_uri": GEMINI_LIVE_WSS_URI,
+    }