webrtc-latency-test/gateway/main.py

#!/usr/bin/env python3
"""
Gateway/Orquestrador de Microserviços
Integra Whisper, LLM, TTS e MuseTalk para conversação com avatar
"""
import asyncio
import time
import logging
from datetime import datetime
from typing import Dict, Any
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import httpx
import uvicorn

# Configuração de logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("gateway")

app = FastAPI(title="Avatar Conversation Gateway", version="1.0.0")

# URLs dos microserviços
WHISPER_URL = "http://localhost:5001"
LLM_URL = "http://localhost:5002"
TTS_URL = "http://localhost:5003"
MUSETALK_URL = "http://localhost:5004"

# Models
class ConversationRequest(BaseModel):
    audio_data: str  # Base64 encoded audio
    sample_rate: int = 16000
    conversation_id: str = "default"

class ConversationResponse(BaseModel):
    video_data: str  # Base64 encoded video
    transcript: str
    llm_response: str
    total_latency_ms: int
    latency_breakdown: Dict[str, int]
    timestamp: str

class LatencyMetrics(BaseModel):
    whisper_ms: int
    llm_ms: int
    tts_ms: int
    musetalk_ms: int
    total_ms: int
    timestamp: str

@app.get("/")
async def root():
    return {
        "service": "Avatar Conversation Gateway",
        "status": "running",
        "version": "1.0.0",
        "microservices": {
            "whisper": WHISPER_URL,
            "llm": LLM_URL,
            "tts": TTS_URL,
            "musetalk": MUSETALK_URL
        },
        "timestamp": datetime.now().isoformat()
    }

@app.get("/health")
async def health():
    """
    Verifica saúde de todos os microserviços
    """
    services_health = {}
    
    async with httpx.AsyncClient() as client:
        # Whisper
        try:
            resp = await client.get(f"{WHISPER_URL}/health", timeout=2.0)
            services_health["whisper"] = resp.json()
        except Exception as e:
            services_health["whisper"] = {"status": "unhealthy", "error": str(e)}
        
        # LLM
        try:
            resp = await client.get(f"{LLM_URL}/health", timeout=2.0)
            services_health["llm"] = resp.json()
        except Exception as e:
            services_health["llm"] = {"status": "unhealthy", "error": str(e)}
        
        # TTS
        try:
            resp = await client.get(f"{TTS_URL}/health", timeout=2.0)
            services_health["tts"] = resp.json()
        except Exception as e:
            services_health["tts"] = {"status": "unhealthy", "error": str(e)}
        
        # MuseTalk
        try:
            resp = await client.get(f"{MUSETALK_URL}/health", timeout=2.0)
            services_health["musetalk"] = resp.json()
        except Exception as e:
            services_health["musetalk"] = {"status": "unhealthy", "error": str(e)}
    
    all_healthy = all(
        s.get("status") == "healthy" 
        for s in services_health.values()
    )
    
    return {
        "status": "healthy" if all_healthy else "degraded",
        "services": services_health
    }

@app.post("/conversation", response_model=ConversationResponse)
async def conversation(request: ConversationRequest):
    """
    Pipeline completo: Audio -> Whisper -> LLM -> TTS -> MuseTalk -> Video
    Mede latência de cada etapa
    """
    start_time = time.time()
    latency_breakdown = {}
    
    logger.info(f"Starting conversation pipeline (conversation_id: {request.conversation_id})")
    
    async with httpx.AsyncClient() as client:
        # 1. WHISPER: Audio -> Text (STT)
        whisper_start = time.time()
        try:
            whisper_response = await client.post(
                f"{WHISPER_URL}/transcribe",
                json={
                    "audio_data": request.audio_data,
                    "sample_rate": request.sample_rate,
                    "language": "pt"
                },
                timeout=5.0
            )
            whisper_data = whisper_response.json()
            transcript = whisper_data["text"]
            latency_breakdown["whisper"] = int((time.time() - whisper_start) * 1000)
            logger.info(f"Whisper: '{transcript}' ({latency_breakdown['whisper']}ms)")
        except Exception as e:
            logger.error(f"Whisper error: {e}")
            raise HTTPException(status_code=500, detail=f"Whisper service error: {e}")
        
        # 2. LLM: Text -> Response
        llm_start = time.time()
        try:
            llm_response = await client.post(
                f"{LLM_URL}/generate",
                json={
                    "text": transcript,
                    "conversation_id": request.conversation_id,
                    "temperature": 0.7,
                    "max_tokens": 150
                },
                timeout=10.0
            )
            llm_data = llm_response.json()
            llm_text = llm_data["text"]
            latency_breakdown["llm"] = int((time.time() - llm_start) * 1000)
            logger.info(f"LLM: '{llm_text}' ({latency_breakdown['llm']}ms)")
        except Exception as e:
            logger.error(f"LLM error: {e}")
            raise HTTPException(status_code=500, detail=f"LLM service error: {e}")
        
        # 3. TTS: Text -> Audio
        tts_start = time.time()
        try:
            tts_response = await client.post(
                f"{TTS_URL}/synthesize",
                json={
                    "text": llm_text,
                    "voice_id": "pt-BR-male",
                    "speed": 1.0,
                    "sample_rate": 16000
                },
                timeout=10.0
            )
            tts_data = tts_response.json()
            audio_data = tts_data["audio_data"]
            latency_breakdown["tts"] = int((time.time() - tts_start) * 1000)
            logger.info(f"TTS: {tts_data['duration_seconds']:.2f}s audio ({latency_breakdown['tts']}ms)")
        except Exception as e:
            logger.error(f"TTS error: {e}")
            raise HTTPException(status_code=500, detail=f"TTS service error: {e}")
        
        # 4. MuseTalk: Audio -> Video
        musetalk_start = time.time()
        try:
            musetalk_response = await client.post(
                f"{MUSETALK_URL}/generate-video",
                json={
                    "audio_data": audio_data,
                    "sample_rate": 16000,
                    "avatar_id": "default"
                },
                timeout=15.0
            )
            musetalk_data = musetalk_response.json()
            video_data = musetalk_data["video_data"]
            latency_breakdown["musetalk"] = int((time.time() - musetalk_start) * 1000)
            logger.info(f"MuseTalk: {musetalk_data['fps']}fps video ({latency_breakdown['musetalk']}ms)")
        except Exception as e:
            logger.error(f"MuseTalk error: {e}")
            raise HTTPException(status_code=500, detail=f"MuseTalk service error: {e}")
    
    # Calcular latência total
    total_latency = int((time.time() - start_time) * 1000)
    
    logger.info(f"Pipeline complete! Total latency: {total_latency}ms")
    logger.info(f"Breakdown: {latency_breakdown}")
    
    return ConversationResponse(
        video_data=video_data,
        transcript=transcript,
        llm_response=llm_text,
        total_latency_ms=total_latency,
        latency_breakdown=latency_breakdown,
        timestamp=datetime.now().isoformat()
    )

@app.get("/metrics", response_model=LatencyMetrics)
async def get_metrics():
    """
    Executa um teste rápido e retorna métricas de latência
    """
    logger.info("Running latency test...")
    
    # Dados de teste
    test_request = ConversationRequest(
        audio_data="dGVzdCBhdWRpbw==",  # "test audio" em base64
        sample_rate=16000,
        conversation_id="test"
    )
    
    # Executar pipeline
    response = await conversation(test_request)
    
    return LatencyMetrics(
        whisper_ms=response.latency_breakdown["whisper"],
        llm_ms=response.latency_breakdown["llm"],
        tts_ms=response.latency_breakdown["tts"],
        musetalk_ms=response.latency_breakdown["musetalk"],
        total_ms=response.total_latency_ms,
        timestamp=response.timestamp
    )

if __name__ == "__main__":
    logger.info("Starting Gateway on port 8080")
    uvicorn.run(app, host="0.0.0.0", port=8080, log_level="info")