wehub-resource-sync
124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
import asyncio
|
|
import base64
|
|
import json
|
|
import logging
|
|
import os
|
|
|
|
from dotenv import load_dotenv
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import FileResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from gemini_live import GeminiLive
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Configuration
|
|
PROJECT_ID = os.getenv("PROJECT_ID", "your-gcp-project-id")
|
|
LOCATION = os.getenv("LOCATION", "us-central1")
|
|
MODEL = os.getenv("MODEL", "gemini-live-2.5-flash-native-audio")
|
|
|
|
# Initialize FastAPI
|
|
app = FastAPI()
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# Serve static files
|
|
app.mount("/static", StaticFiles(directory="frontend"), name="static")
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return FileResponse("frontend/index.html")
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
"""WebSocket endpoint for Gemini Live."""
|
|
await websocket.accept()
|
|
|
|
logger.info("WebSocket connection accepted")
|
|
|
|
audio_input_queue = asyncio.Queue()
|
|
video_input_queue = asyncio.Queue()
|
|
text_input_queue = asyncio.Queue()
|
|
|
|
async def audio_output_callback(data):
|
|
await websocket.send_bytes(data)
|
|
|
|
async def audio_interrupt_callback():
|
|
# The event queue handles the JSON message, but we might want to do something else here
|
|
pass
|
|
|
|
gemini_client = GeminiLive(
|
|
project_id=PROJECT_ID, location=LOCATION, model=MODEL, input_sample_rate=16000
|
|
)
|
|
|
|
async def receive_from_client():
|
|
try:
|
|
while True:
|
|
message = await websocket.receive()
|
|
|
|
if message.get("bytes"):
|
|
await audio_input_queue.put(message["bytes"])
|
|
elif message.get("text"):
|
|
text = message["text"]
|
|
try:
|
|
payload = json.loads(text)
|
|
if isinstance(payload, dict) and payload.get("type") == "image":
|
|
image_data = base64.b64decode(payload["data"])
|
|
await video_input_queue.put(image_data)
|
|
continue
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
await text_input_queue.put(text)
|
|
except WebSocketDisconnect:
|
|
logger.info("WebSocket disconnected")
|
|
except Exception as e:
|
|
logger.error(f"Error receiving from client: {e}")
|
|
|
|
receive_task = asyncio.create_task(receive_from_client())
|
|
|
|
async def run_session():
|
|
async for event in gemini_client.start_session(
|
|
audio_input_queue=audio_input_queue,
|
|
video_input_queue=video_input_queue,
|
|
text_input_queue=text_input_queue,
|
|
audio_output_callback=audio_output_callback,
|
|
audio_interrupt_callback=audio_interrupt_callback,
|
|
):
|
|
if event:
|
|
# Forward events (transcriptions, etc) to client
|
|
await websocket.send_json(event)
|
|
|
|
try:
|
|
await run_session()
|
|
except Exception as e:
|
|
logger.error(f"Error in Gemini session: {e}")
|
|
finally:
|
|
receive_task.cancel()
|
|
# Ensure websocket is closed if not already
|
|
try:
|
|
await websocket.close()
|
|
except:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
port = int(os.getenv("PORT", 8000))
|
|
uvicorn.run(app, host="localhost", port=port)
|