Fix HF Space: FastAPI health check, Gemini WS bridge, deps & Docker CMD

Dockerfile

@@ -10,4 +10,4 @@ COPY . .
 ENV PORT=7860
 EXPOSE 7860
 
-CMD ["python", "app.py"]
+CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port ${PORT}"]

app.py

@@ -2,20 +2,30 @@
 import json
 import os
 import base64
-import websockets
+
+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.
-client = genai.Client(
-    http_options={"api_version": "v1alpha"}
-)
+client = genai.Client(http_options={"api_version": "v1alpha"})
+
+app = FastAPI()
+
+@app.get("/")
+def health():
+    # Hugging Face health check hits this as plain HTTP.
+    return {"status": "ok"}
+
+@app.websocket("/ws")
+async def gemini_ws_bridge(ws: WebSocket):
+    await ws.accept()
 
-async def gemini_session_handler(client_websocket: websockets.WebSocketServerProtocol):
     try:
-        config_message = await client_websocket.recv()
+        # First message must be config (same as your old code)
+        config_message = await ws.receive_text()
         config_data = json.loads(config_message)
 
         config = config_data.get("setup", {})
@@ -29,9 +39,11 @@ async def gemini_session_handler(client_websocket: websockets.WebSocketServerPro
         )
 
         async with client.aio.live.connect(model=MODEL, config=config) as session:
+
             async def send_to_gemini():
                 try:
-                    async for message in client_websocket:
+                    while True:
+                        message = await ws.receive_text()
                         data = json.loads(message)
 
                         if "realtime_input" in data:
@@ -40,6 +52,8 @@ async def gemini_session_handler(client_websocket: websockets.WebSocketServerPro
                                 payload = chunk.get("data")
                                 if mt in ("audio/pcm", "image/jpeg") and payload:
                                     await session.send({"mime_type": mt, "data": payload})
+                except WebSocketDisconnect:
+                    pass
                 except Exception as e:
                     print(f"send_to_gemini error: {e}")
 
@@ -53,36 +67,29 @@ async def gemini_session_handler(client_websocket: websockets.WebSocketServerPro
                         if model_turn:
                             for part in model_turn.parts:
                                 if getattr(part, "text", None):
-                                    await client_websocket.send(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 client_websocket.send(json.dumps({"audio": b64_audio}))
+                                    await ws.send_text(json.dumps({"audio": b64_audio}))
 
                         if response.server_content.turn_complete:
-                            # optional: notify turn complete
-                            await client_websocket.send(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}")
 
-            await asyncio.gather(
-                asyncio.create_task(send_to_gemini()),
-                asyncio.create_task(receive_from_gemini()),
-            )
+            await asyncio.gather(send_to_gemini(), receive_from_gemini())
 
+    except WebSocketDisconnect:
+        pass
     except Exception as e:
-        print(f"gemini_session_handler error: {e}")
+        print(f"gemini_ws_bridge error: {e}")
         try:
-            await client_websocket.send(json.dumps({"error": str(e)}))
+            await ws.send_text(json.dumps({"error": str(e)}))
+        except Exception:
+            pass
+        try:
+            await ws.close()
         except Exception:
             pass
-
-async def main():
-    port = int(os.environ.get("PORT", "7860"))
-    host = "0.0.0.0"
-    print(f"Starting WebSocket server on {host}:{port}")
-
-    async with websockets.serve(gemini_session_handler, host, port):
-        await asyncio.Future()  # run forever
-
-if __name__ == "__main__":
-    asyncio.run(main())

requirements.txt

@@ -1,2 +1,4 @@
 google-genai==0.3.0
 websockets==14.1
+fastapi==0.115.6
+uvicorn[standard]==0.32.1