server_webrtc.py

"""
FastAPI Server with WebRTC Signaling for Voice Agent

This server provides:
- WebRTC signaling endpoint (/offer)
- Health check endpoint (/health)
- Cognito JWT token validation
"""
import asyncio
import os
from contextlib import asynccontextmanager
from typing import Optional

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Request, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from loguru import logger
from pydantic import BaseModel

from pipecat.transports.smallwebrtc.connection import SmallWebRTCConnection

load_dotenv(override=True)

# Import bot
from bot_webrtc import run_voice_agent


class OfferRequest(BaseModel):
    """WebRTC offer request."""
    sdp: str
    type: str = "offer"


class OfferResponse(BaseModel):
    """WebRTC answer response."""
    sdp: str
    type: str = "answer"


# Track active connections
active_connections: dict = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler."""
    logger.info("Starting Voice Agent Server")
    yield
    logger.info("Shutting down Voice Agent Server")
    # Cancel all active connections
    for conn_id, task in active_connections.items():
        task.cancel()
        logger.info(f"Cancelled connection: {conn_id}")


app = FastAPI(
    title="AWS Voice Agent API",
    description="Voice Agent with WebRTC and LangSmith Tracing",
    version="1.0.0",
    lifespan=lifespan
)

# CORS configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Optional: Cognito JWT validation
async def validate_token(request: Request) -> Optional[dict]:
    """Validate Cognito JWT token from Authorization header."""
    auth_header = request.headers.get("Authorization")
    
    # Skip validation if Cognito is not configured
    if not os.getenv("COGNITO_USER_POOL_ID"):
        return {"sub": "anonymous"}
    
    if not auth_header or not auth_header.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid authorization header")
    
    token = auth_header.split(" ")[1]
    
    # TODO: Implement proper JWT validation with Cognito
    # For now, just check if token exists
    # In production, use python-jose or cognitojwt libraries
    if not token:
        raise HTTPException(status_code=401, detail="Invalid token")
    
    return {"sub": "user", "token": token}


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy", "service": "voice-agent"}


@app.get("/")
async def root():
    """Root endpoint."""
    return {"message": "AWS Voice Agent API", "version": "1.0.0"}


@app.post("/offer", response_model=OfferResponse)
async def webrtc_offer(
    request: OfferRequest,
    # user: dict = Depends(validate_token)  # Uncomment to enable auth
):
    """
    WebRTC signaling endpoint.
    
    Receives SDP offer from client, creates a WebRTC connection,
    spawns a voice agent, and returns SDP answer.
    """
    try:
        logger.info("Received WebRTC offer")
        
        # Create WebRTC connection
        connection = SmallWebRTCConnection()
        
        # Set remote description (offer) and get local description (answer)
        answer_sdp = await connection.create_answer(request.sdp)
        
        # Start the voice agent in background
        connection_id = id(connection)
        
        async def run_bot():
            try:
                await run_voice_agent(connection)
            except Exception as e:
                logger.error(f"Bot error: {e}")
            finally:
                if connection_id in active_connections:
                    del active_connections[connection_id]
        
        task = asyncio.create_task(run_bot())
        active_connections[connection_id] = task
        
        logger.info(f"Started voice agent, connection ID: {connection_id}")
        
        return OfferResponse(
            sdp=answer_sdp,
            type="answer"
        )
        
    except Exception as e:
        logger.error(f"Error handling WebRTC offer: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/status")
async def get_status():
    """Get server status."""
    return {
        "active_connections": len(active_connections),
        "tracing_enabled": bool(os.getenv("LANGSMITH_API_KEY")),
        "region": os.getenv("AWS_REGION", "us-east-1"),
    }


if __name__ == "__main__":
    import uvicorn
    
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 7860))
    
    logger.info(f"Starting server on {host}:{port}")
    
    uvicorn.run(
        "server_webrtc:app",
        host=host,
        port=port,
        reload=True,
    )