Final stable FastAPI Gemini WS bridge for HF

app.py

@@ -2,29 +2,41 @@
 import json
 import os
 import base64
+import logging
+
+# ---- Logging: reduce HF probe noise ----
+logging.basicConfig(level=logging.INFO)
+logging.getLogger("websockets").setLevel(logging.ERROR)
+logging.getLogger("websockets.server").setLevel(logging.ERROR)
+logging.getLogger("uvicorn.error").setLevel(logging.INFO)
+logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
+
+print("APP STARTED ✅ FastAPI + Gemini WebSocket Bridge")
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from google import genai
 
 MODEL = os.environ.get("MODEL", "gemini-2.0-flash-exp")
 
-# IMPORTANT: don't hardcode keys in code.
-# Set GOOGLE_API_KEY as a Hugging Face Secret.
+# IMPORTANT: Set GOOGLE_API_KEY as a Hugging Face Secret
 client = genai.Client(http_options={"api_version": "v1alpha"})
 
 app = FastAPI()
 
+
+# ---- HTTP health check (required by Hugging Face) ----
 @app.get("/")
 def health():
-    # Hugging Face health check hits this as plain HTTP.
     return {"status": "ok"}
 
+
+# ---- WebSocket endpoint ----
 @app.websocket("/ws")
 async def gemini_ws_bridge(ws: WebSocket):
     await ws.accept()
 
     try:
-        # First message must be config (same as your old code)
+        # First message must be setup/config
         config_message = await ws.receive_text()
         config_data = json.loads(config_message)
 
@@ -51,11 +63,13 @@ async def gemini_ws_bridge(ws: WebSocket):
                                 mt = chunk.get("mime_type")
                                 payload = chunk.get("data")
                                 if mt in ("audio/pcm", "image/jpeg") and payload:
-                                    await session.send({"mime_type": mt, "data": payload})
+                                    await session.send(
+                                        {"mime_type": mt, "data": payload}
+                                    )
                 except WebSocketDisconnect:
                     pass
                 except Exception as e:
-                    print(f"send_to_gemini error: {e}")
+                    logging.error(f"send_to_gemini error: {e}")
 
             async def receive_from_gemini():
                 try:
@@ -67,24 +81,32 @@ async def gemini_ws_bridge(ws: WebSocket):
                         if model_turn:
                             for part in model_turn.parts:
                                 if getattr(part, "text", None):
-                                    await ws.send_text(json.dumps({"text": part.text}))
+                                    await ws.send_text(
+                                        json.dumps({"text": part.text})
+                                    )
                                 elif getattr(part, "inline_data", None):
-                                    b64_audio = base64.b64encode(part.inline_data.data).decode("utf-8")
-                                    await ws.send_text(json.dumps({"audio": b64_audio}))
+                                    b64_audio = base64.b64encode(
+                                        part.inline_data.data
+                                    ).decode("utf-8")
+                                    await ws.send_text(
+                                        json.dumps({"audio": b64_audio})
+                                    )
 
                         if response.server_content.turn_complete:
-                            await ws.send_text(json.dumps({"turn_complete": True}))
+                            await ws.send_text(
+                                json.dumps({"turn_complete": True})
+                            )
                 except WebSocketDisconnect:
                     pass
                 except Exception as e:
-                    print(f"receive_from_gemini error: {e}")
+                    logging.error(f"receive_from_gemini error: {e}")
 
             await asyncio.gather(send_to_gemini(), receive_from_gemini())
 
     except WebSocketDisconnect:
         pass
     except Exception as e:
-        print(f"gemini_ws_bridge error: {e}")
+        logging.error(f"gemini_ws_bridge error: {e}")
         try:
             await ws.send_text(json.dumps({"error": str(e)}))
         except Exception: