app.py

# Version: Text Bot Support v2.0 - Enhanced routing for text clients
import os
import logging
import time
from datetime import datetime
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from websocket_handler import handle_websocket_connection
from enhanced_websocket_handler import handle_enhanced_websocket_connection
from conversational_websocket_handler import handle_conversational_websocket
from hybrid_llm_service import HybridLLMService
from voice_service import VoiceService
from groq_voice_service import groq_voice_service  # Import the new Groq voice service
from rag_service import search_documents_async
from lancedb_service import LanceDBService
from scenario_analysis_service import ScenarioAnalysisService
from evidence_pack_export import export_evidence_pack_pdf, export_evidence_pack_csv
from groq_websocket_handler import groq_websocket_handler
import config
from dotenv import load_dotenv
import json
import base64

# MCP and Authentication imports
from fastapi import Depends
from pydantic import BaseModel
from typing import Optional
from auth import get_current_user

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Get configuration
config_dict = {
    "ALLOWED_ORIGINS": config.ALLOWED_ORIGINS,
    "ENABLE_VOICE_FEATURES": config.ENABLE_VOICE_FEATURES
}

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler"""
    # Startup
    logger.info("🚀 Starting Voice Bot Application...")
    
    # Setup sample documents if database is empty
    try:
        from setup_documents import setup_sample_documents
        await setup_sample_documents()
    except Exception as e:
        logger.warning(f"⚠️ Could not setup sample documents: {e}")
    
    logger.info("✅ Application started successfully")
    yield
    # Shutdown (if needed)
    logger.info("🛑 Shutting down Voice Bot Application...")

# Create FastAPI application
app = FastAPI(
    title="Voice Bot Government Assistant",
    description="AI-powered voice assistant for government policies and services",
    version="1.0.0",
    lifespan=lifespan
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=config.ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize services (lazy loading for HF Spaces)
llm_service = None
voice_service = None
lancedb_service = None
scenario_service = None

def get_llm_service():
    global llm_service
    if llm_service is None:
        llm_service = HybridLLMService()
    return llm_service

def get_voice_service():
    global voice_service
    if voice_service is None:
        voice_service = VoiceService()
    return voice_service

def get_lancedb_service():
    global lancedb_service
    if lancedb_service is None:
        lancedb_service = LanceDBService()
    return lancedb_service

def get_scenario_service():
    global scenario_service
    if scenario_service is None:
        scenario_service = ScenarioAnalysisService()
    return scenario_service

# Evidence Pack Export endpoint
@app.api_route("/export_evidence_pack", methods=["GET", "POST", "OPTIONS", "HEAD"])
async def export_evidence_pack(request: Request, format: str = "pdf"):
    """Export evidence pack in PDF or CSV format"""
    try:
        # Handle CORS preflight and HEAD requests
        if request.method == "OPTIONS":
            return JSONResponse({"status": "ok"}, status_code=200)
        
        if request.method == "HEAD":
            # For HEAD requests, return headers without body
            return JSONResponse(
                {"status": "ok"}, 
                status_code=200,
                headers={"Content-Type": "application/pdf" if format.lower() == "pdf" else "text/csv"}
            )
        
        # Handle both GET and POST requests
        if request.method == "POST":
            try:
                data = await request.json()
                # Format can come from query params or request body
                format = request.query_params.get("format", data.get("format", "pdf"))
            except Exception:
                # If JSON parsing fails, use query params
                data = {
                    "query": request.query_params.get("query", ""),
                    "format": format,
                    "timestamp": datetime.now().isoformat()
                }
        else:  # GET request
            # For GET requests, we need some default data structure
            data = {
                "query": request.query_params.get("query", ""),
                "format": format,
                "timestamp": datetime.now().isoformat(),
                "message": "Sample evidence pack export"
            }
        
        if format.lower() == "pdf":
            file_path = export_evidence_pack_pdf(data)
            return FileResponse(
                file_path, 
                media_type="application/pdf", 
                filename="evidence_pack.pdf",
                headers={"Content-Disposition": "attachment; filename=evidence_pack.pdf"}
            )
        elif format.lower() == "csv":
            file_path = export_evidence_pack_csv(data)
            return FileResponse(
                file_path, 
                media_type="text/csv", 
                filename="evidence_pack.csv",
                headers={"Content-Disposition": "attachment; filename=evidence_pack.csv"}
            )
        else:
            return JSONResponse({"error": "Invalid format. Use 'pdf' or 'csv'"}, status_code=400)
    except Exception as e:
        logger.error(f"Export evidence pack error: {str(e)}")
        return JSONResponse({"error": f"Export failed: {str(e)}"}, status_code=500)

# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "voice-bot-api",
        "timestamp": datetime.now().isoformat(),
        "version": "1.0.0"
    }

# Root endpoint
@app.get("/")
async def root():
    """Root endpoint with service information"""
    return {
        "message": "Voice Bot Government Assistant API",
        "status": "running",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "chat": "/chat",
            "websocket": "/ws",
            "websocket_stream": "/ws/stream",
            "export_evidence_pack": "/export_evidence_pack",
            "docs": "/docs"
        }
    }

# Chat endpoint
@app.post("/chat")
async def chat_endpoint(request: dict):
    """Text-based chat endpoint"""
    try:
        message = request.get("message", "")
        if not message:
            raise HTTPException(status_code=400, detail="Message is required")
        
        llm = get_llm_service()
        response = await llm.get_response(message)
        
        return {
            "response": response,
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# WebSocket endpoints
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time communication"""
    await handle_enhanced_websocket_connection(websocket)

@app.websocket("/ws/stream")
async def websocket_stream_endpoint(websocket: WebSocket):
    """
    Enhanced WebSocket endpoint compatible with friend's frontend format
    Handles both JSON audio format and friend's JSON+Binary format
    """
    # Accept connection and get session ID
    session_id = await groq_websocket_handler.connect(websocket)
    
    # Send connection successful message (friend's format)
    await websocket.send_json({"type": "connection_successful"})
    logger.info("✅ WebSocket connection established")
    
    try:
        while True:
            try:
                # Try to receive JSON message first
                message_text = await websocket.receive_text()
                message = json.loads(message_text)
                
                # Check if this is friend's format with lang field
                if "lang" in message and "type" not in message:
                    logger.info(f"📱 Received friend's format: {message}")
                    
                    # This is friend's format - expect binary audio next
                    try:
                        # Receive the binary audio data
                        audio_bytes = await websocket.receive_bytes()
                        logger.info(f"🎤 Received {len(audio_bytes)} bytes of audio data")
                        
                        # Convert to base64 for internal processing
                        audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
                        
                        # Create standard message format for internal processing
                        standard_message = {
                            "type": "audio_data",
                            "language": "en" if message.get("lang", "").lower().startswith("eng") else message.get("lang", "en"),
                            "audio_data": audio_base64,
                            "user_id": message.get("user_id")
                        }
                        
                        # Process using standard handler
                        await groq_websocket_handler.handle_stream_message(websocket, session_id, standard_message)
                        
                    except Exception as audio_error:
                        logger.error(f"❌ Error receiving binary audio: {audio_error}")
                        await websocket.send_json({
                            "type": "error", 
                            "message": "Failed to receive audio data"
                        })
                        
                else:
                    # This is standard format - process normally
                    await groq_websocket_handler.handle_stream_message(websocket, session_id, message)
                    
            except json.JSONDecodeError:
                await websocket.send_json({
                    "type": "error",
                    "message": "Invalid JSON message"
                })
                continue
                
    except Exception as e:
        logger.error(f"❌ WebSocket stream error: {e}")
    finally:
        await groq_websocket_handler.disconnect(session_id)

@app.websocket("/ws/conversational")
async def websocket_conversational_endpoint(websocket: WebSocket):
    """
    Enhanced Conversational WebSocket endpoint with session memory and personalization
    Based on friend's conversational implementation for better user experience
    """
    await handle_conversational_websocket(websocket)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)