life-echo/api/routers/websocket.py

"""
WebSocket 路由：实时对话通信
"""
import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Dict

from fastapi import WebSocket, WebSocketDisconnect, HTTPException, Query
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from agents import ConversationAgent, MemoryAgent
from agents.prompts import ConversationStage
from database import get_async_db
from database.models import Conversation, Segment
from database.models import User as UserModel
from services.auth_service import verify_token
from fastapi import HTTPException, status


class MessageType(str, Enum):
    """WebSocket 消息类型"""
    CONNECT = "connect"
    AUDIO_CHUNK = "audio_chunk"
    TEXT = "text"  # 文本消息
    TRANSCRIPT = "transcript"
    AGENT_RESPONSE = "agent_response"
    TTS_AUDIO = "tts_audio"
    END_CONVERSATION = "end_conversation"
    ERROR = "error"


# 连接管理
class ConnectionManager:
    """WebSocket 连接管理器"""
    
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.conversation_agents: Dict[str, ConversationAgent] = {}
        self.memory_agent = MemoryAgent()
    
    async def connect(self, websocket: WebSocket, conversation_id: str):
        """建立连接"""
        await websocket.accept()
        self.active_connections[conversation_id] = websocket
        self.conversation_agents[conversation_id] = ConversationAgent()
    
    def disconnect(self, conversation_id: str):
        """断开连接"""
        if conversation_id in self.active_connections:
            del self.active_connections[conversation_id]
        if conversation_id in self.conversation_agents:
            self.conversation_agents[conversation_id].clear_memory(conversation_id)
            del self.conversation_agents[conversation_id]
    
    async def send_message(self, conversation_id: str, message: dict):
        """发送消息"""
        if conversation_id in self.active_connections:
            websocket = self.active_connections[conversation_id]
            await websocket.send_json(message)
    
    async def receive_message(self, conversation_id: str) -> dict:
        """接收消息"""
        if conversation_id in self.active_connections:
            websocket = self.active_connections[conversation_id]
            return await websocket.receive_json()
        raise HTTPException(status_code=404, detail="Connection not found")


manager = ConnectionManager()


async def websocket_endpoint(
    websocket: WebSocket, 
    conversation_id: str, 
    token: str = Query(..., description="访问令牌")
):
    """
    WebSocket 端点：处理实时对话
    
    Args:
        websocket: WebSocket 连接
        conversation_id: 对话 ID
        token: 访问令牌（从查询参数获取）
    """
    # 验证JWT令牌
    payload = verify_token(token)
    if not payload or payload.get("type") != "access":
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无效的认证令牌")
        return
    
    user_id = payload.get("sub")
    if not user_id:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无效的令牌内容")
        return
    
    # 验证用户是否存在
    async for db in get_async_db():
        user = await db.get(UserModel, user_id)
        if not user:
            await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="用户不存在")
            return
        
        await manager.connect(websocket, conversation_id)
        
        try:
            # 发送连接确认
            await manager.send_message(conversation_id, {
                "type": MessageType.CONNECT,
                "conversation_id": conversation_id,
                "data": {"status": "connected"},
                "timestamp": datetime.now(timezone.utc).isoformat()
            })
            
            # 从数据库获取对话信息
            conversation = await db.get(Conversation, conversation_id)
            if not conversation:
                # 如果对话不存在，创建新对话
                conversation = Conversation(
                    id=conversation_id,
                    user_id=user_id,
                    started_at=datetime.now(timezone.utc),
                    status="active"
                )
                db.add(conversation)
                await db.commit()
            else:
                # 验证用户权限：只能访问自己的对话
                if conversation.user_id != user_id:
                    await manager.send_message(conversation_id, {
                        "type": MessageType.ERROR,
                        "data": {"message": "无权访问此对话"},
                        "timestamp": datetime.now(timezone.utc).isoformat()
                    })
                    await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无权访问此对话")
                    return
            
            
            current_stage = ConversationStage(conversation.conversation_stage) if conversation.conversation_stage else ConversationStage.CHILDHOOD
            
            # 主循环：处理消息
            while True:
                try:
                    message = await websocket.receive_json()
                    msg_type = message.get("type")
                    
                    if msg_type == MessageType.TEXT:
                        # 处理文本消息
                        text_message = message.get("data", {}).get("text", "")
                        
                        if text_message:
                            # 保存段落到数据库
                            segment = Segment(
                                id=str(uuid.uuid4()),
                                conversation_id=conversation_id,
                                transcript_text=text_message,
                                processed=False
                            )
                            db.add(segment)
                            await db.commit()
                            
                            # Agent 生成回应
                            current_stage = await process_user_message(
                                conversation_id=conversation_id,
                                user_message=text_message,
                                current_stage=current_stage,
                                conversation=conversation,
                                segment=segment,
                                db=db,
                                manager=manager
                            )
                    
                    elif msg_type == MessageType.END_CONVERSATION:
                        # 结束对话
                        conversation.status = "ended"
                        conversation.ended_at = datetime.now(timezone.utc)
                        await db.commit()
                        
                        # 触发整理 Agent
                        await process_conversation_segments(conversation_id, db)
                        
                        await manager.send_message(conversation_id, {
                            "type": MessageType.END_CONVERSATION,
                            "conversation_id": conversation_id,
                            "data": {"status": "ended"},
                            "timestamp": datetime.now(timezone.utc).isoformat()
                        })
                        break
                
                except Exception as e:
                    await manager.send_message(conversation_id, {
                        "type": MessageType.ERROR,
                        "data": {"message": str(e)},
                        "timestamp": datetime.now(timezone.utc).isoformat()
                    })
    
        except WebSocketDisconnect:
            manager.disconnect(conversation_id)
            break
        except Exception as e:
            manager.disconnect(conversation_id)
            raise


async def process_user_message(
    conversation_id: str,
    user_message: str,
    current_stage: ConversationStage,
    conversation: Conversation,
    segment: Segment,
    db: AsyncSession,
    manager: ConnectionManager
) -> ConversationStage:
    """
    处理用户消息，生成Agent回应
    
    Args:
        conversation_id: 对话ID
        user_message: 用户消息文本
        current_stage: 当前对话阶段
        conversation: 对话对象
        segment: 段落对象
        db: 数据库会话
        manager: 连接管理器
    
    Returns:
        更新后的对话阶段
    """
    agent = manager.conversation_agents.get(conversation_id)
    if agent:
        # 检测对话阶段
        detected_stage = agent.detect_stage(conversation_id, user_message)
        if detected_stage != current_stage:
            current_stage = detected_stage
            conversation.conversation_stage = current_stage.value
            await db.commit()
        
        # 获取已聊话题
        stmt_segments = select(Segment).where(
            Segment.conversation_id == conversation_id
        ).order_by(Segment.created_at)
        result_segments = await db.execute(stmt_segments)
        previous_segments = result_segments.scalars().all()
        covered_topics = [seg.topic_category for seg in previous_segments if seg.topic_category]
        
        # 生成回应
        response = agent.generate_response(
            conversation_id=conversation_id,
            user_message=user_message,
            current_stage=current_stage,
            covered_topics=covered_topics
        )
        
        # 更新段落的 Agent 回应
        segment.agent_response = response
        await db.commit()
        
        # 发送 Agent 回应（仅文字，不生成语音）
        await manager.send_message(conversation_id, {
            "type": MessageType.AGENT_RESPONSE,
            "conversation_id": conversation_id,
            "data": {"text": response},
            "timestamp": datetime.now(timezone.utc).isoformat()
        })
    
    return current_stage


async def process_conversation_segments(conversation_id: str, db: AsyncSession):
    """
    处理对话段落，生成章节
    
    Args:
        conversation_id: 对话 ID
        db: 数据库会话
    """
    # 获取所有未处理的段落
    stmt = select(Segment).where(
        Segment.conversation_id == conversation_id,
        Segment.processed == False
    )
    result = await db.execute(stmt)
    segments = result.scalars().all()
    
    if not segments:
        return
    
    # 准备段落数据
    segments_data = [
        {"transcript_text": seg.transcript_text}
        for seg in segments
    ]
    
    # 调用整理 Agent
    memory_agent = manager.memory_agent
    chapters_data = memory_agent.process_segments(segments_data)
    
    # 保存章节到数据库
    from database.models import Chapter as ChapterModel
    conversation = await db.get(Conversation, conversation_id)
    
    if not conversation:
        return
    
    for category, chapter_data in chapters_data.items():
        chapter = ChapterModel(
            id=str(uuid.uuid4()),
            user_id=conversation.user_id,
            title=chapter_data.get("title", f"章节-{category}"),
            content=chapter_data.get("content", ""),
            order_index=chapter_data.get("order_index", 999),
            status="completed",
            category=category,
            images=chapter_data.get("image_suggestions", [])
        )
        db.add(chapter)
    
    # 标记段落为已处理
    for seg in segments:
        seg.processed = True
    
    await db.commit()