""" WebSocket 路由：实时对话通信 支持异步 Agent 调用和 Redis 会话存储 """ import asyncio import logging import uuid from dataclasses import dataclass, field from datetime import datetime, timezone from enum import Enum from typing import Dict, List, Optional, Set, Tuple from fastapi import WebSocket, WebSocketDisconnect, HTTPException, status from starlette.websockets import WebSocketState from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from agents import ConversationAgent, MemoryAgent from agents.memoir_processor import BackgroundTaskRunner 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 services.memoir_state_service import get_or_create_state from services import asr_service logger = logging.getLogger(__name__) LEGACY_VOICE_SESSION_ID = "legacy" class MessageType(str, Enum): """WebSocket 消息类型""" CONNECT = "connect" AUDIO_CHUNK = "audio_chunk" AUDIO_SEGMENT = "audio_segment" # 分段语音消息（长语音持续上传） AUDIO_MESSAGE = "audio_message" # 完整音频消息（类似微信语音） TRANSCRIBE_ONLY = "transcribe_only" # 仅转写，不落库、不触发 Agent，用于「转文字」发送 TEXT = "text" # 文本消息 TRANSCRIPT = "transcript" # 语音转文字结果 AGENT_RESPONSE = "agent_response" TTS_AUDIO = "tts_audio" END_CONVERSATION = "end_conversation" MEMOIR_UPDATE = "memoir_update" ERROR = "error" # 连接管理 class ConnectionManager: """WebSocket 连接管理器""" def __init__(self): self.active_connections: Dict[str, WebSocket] = {} self.segment_states: Dict[Tuple[str, str], "SegmentStreamState"] = {} # ConversationAgent 现在是无状态的（会话存储在 Redis），可以复用 self.conversation_agent = ConversationAgent() self.memory_agent = MemoryAgent() self.background_runner = BackgroundTaskRunner() async def connect(self, websocket: WebSocket, conversation_id: str): """建立连接""" await websocket.accept() self.active_connections[conversation_id] = websocket async def disconnect(self, conversation_id: str): """断开连接""" if conversation_id in self.active_connections: del self.active_connections[conversation_id] stale_keys = [ key for key, state in self.segment_states.items() if key[0] == conversation_id and not state.active_tasks ] for key in stale_keys: self.segment_states.pop(key, None) # 清除 Redis 中的会话记忆（可选，也可以保留用于恢复） # await self.conversation_agent.clear_memory(conversation_id) def get_or_create_segment_state( self, conversation_id: str, voice_session_id: str, ) -> "SegmentStreamState": state_key = (conversation_id, voice_session_id) if state_key not in self.segment_states: self.segment_states[state_key] = SegmentStreamState() return self.segment_states[state_key] def register_segment_task( self, conversation_id: str, voice_session_id: str, task: asyncio.Task, ) -> None: state_key = (conversation_id, voice_session_id) state = self.get_or_create_segment_state(conversation_id, voice_session_id) state.active_tasks.add(task) def _cleanup(done_task: asyncio.Task) -> None: state.active_tasks.discard(done_task) if not state.active_tasks and conversation_id not in self.active_connections: self.segment_states.pop(state_key, None) if done_task.cancelled(): return exc = done_task.exception() if exc: logger.error( "分段处理任务异常 " f"(conversation_id={conversation_id}, voice_session_id={voice_session_id}): {exc}", exc_info=True, ) task.add_done_callback(_cleanup) async def send_message(self, conversation_id: str, message: dict): """发送消息""" if conversation_id in self.active_connections: websocket = self.active_connections[conversation_id] try: # 尝试发送消息，如果连接已关闭会抛出异常 await websocket.send_json(message) except (RuntimeError, Exception) as e: logger.warning(f"发送消息失败 (conversation_id={conversation_id}): {e}") # 如果发送失败，从连接列表中移除 if conversation_id in self.active_connections: del self.active_connections[conversation_id] 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() @dataclass class SegmentStreamState: """会话内分段处理状态（用于并行 ASR + 有序聚合）""" lock: asyncio.Lock = field(default_factory=asyncio.Lock) pending_indices: Set[int] = field(default_factory=set) processed_indices: Set[int] = field(default_factory=set) buffered_transcripts: Dict[int, Tuple[str, Segment]] = field(default_factory=dict) consumed_index: int = -1 active_tasks: Set[asyncio.Task] = field(default_factory=set) def _normalize_voice_session_id(voice_session_id: Optional[str]) -> str: if voice_session_id: return str(voice_session_id) return LEGACY_VOICE_SESSION_ID def _voice_session_id_from_client_segment_id(client_segment_id: Optional[str]) -> Optional[str]: if not client_segment_id: return None session_id, separator, _ = client_segment_id.rpartition("-") if separator and session_id: return session_id return None def _build_segment_audio_url(voice_session_id: str, segment_index: int) -> str: """构建分段语音的幂等标识（conversation_id + voice_session_id + segment_index）。""" return f"audio-segment:{voice_session_id}:{segment_index}" def _extract_segment_scope(audio_url: Optional[str]) -> Optional[Tuple[str, int]]: """从 audio_url 中解析 voice_session_id 与 segment_index。兼容旧格式 audio-segment:{index}。""" prefix = "audio-segment:" if not audio_url or not audio_url.startswith(prefix): return None payload = audio_url[len(prefix):] voice_session_id_raw, separator, segment_index_raw = payload.rpartition(":") try: if separator: return (_normalize_voice_session_id(voice_session_id_raw), int(segment_index_raw)) return (LEGACY_VOICE_SESSION_ID, int(payload)) except ValueError: return None def _is_transcribe_failure(transcript_text: Optional[str]) -> bool: if not transcript_text: return True return transcript_text.startswith("转写失败") async def _find_existing_segment_by_index( db: AsyncSession, conversation_id: str, voice_session_id: str, segment_index: int, ) -> Optional[Segment]: """ 按 conversation + voice_session_id + segment_index 查找已落库分段。 说明：测试桩的 execute() 不会真正执行 where，所以这里做一次 Python 侧过滤，兼容真实 DB 和单测桩。 """ segment_audio_url = _build_segment_audio_url(voice_session_id, segment_index) stmt = select(Segment).where( Segment.conversation_id == conversation_id, Segment.audio_url == segment_audio_url, ).order_by(Segment.created_at.desc()) result = await db.execute(stmt) candidates = result.scalars().all() for item in candidates: if item.conversation_id == conversation_id and item.audio_url == segment_audio_url: return item return None async def _get_persisted_contiguous_segment_index( db: AsyncSession, conversation_id: str, voice_session_id: str, ) -> int: """读取数据库中当前 voice session 已连续落库的最大 segment_index，用于重连恢复。""" stmt = select(Segment).where(Segment.conversation_id == conversation_id) result = await db.execute(stmt) candidates = result.scalars().all() persisted_indices: Set[int] = set() for item in candidates: if item.conversation_id != conversation_id: continue segment_scope = _extract_segment_scope(item.audio_url) if not segment_scope: continue item_voice_session_id, item_index = segment_scope if item_voice_session_id != voice_session_id: continue persisted_indices.add(item_index) contiguous_index = -1 while contiguous_index + 1 in persisted_indices: contiguous_index += 1 return contiguous_index async def _send_segment_transition_feedback( conversation_id: str, segment_index: int, manager: ConnectionManager, ) -> None: """ASR 处理中先给陪伴式过渡反馈，避免用户感知卡住。""" await manager.send_message(conversation_id, { "type": MessageType.AGENT_RESPONSE, "conversation_id": conversation_id, "data": { "text": "我在认真听，你继续说，我会边听边整理重点。", "transition": True, "segment_index": segment_index, }, "timestamp": datetime.now(timezone.utc).isoformat(), }) async def _process_audio_segment_async( conversation_id: str, user_id: str, voice_session_id: str, segment_index: int, audio_base64: str, audio_duration: int, is_last: bool, manager: ConnectionManager, ) -> None: """分段语音的异步处理：并行 ASR + 幂等落库 + 有序聚合触发 Agent。""" state = manager.get_or_create_segment_state(conversation_id, voice_session_id) try: # 每个分段任务使用独立 DB Session，避免与主循环共享同一 AsyncSession 导致并发冲突。 async for db in get_async_db(): conversation = await db.get(Conversation, conversation_id) user = await db.get(UserModel, user_id) if not conversation: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "对话不存在，分段处理已取消"}, "timestamp": datetime.now(timezone.utc).isoformat(), }) return if not user: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "用户不存在，分段处理已取消"}, "timestamp": datetime.now(timezone.utc).isoformat(), }) return async with state.lock: should_prime_state = ( state.consumed_index < 0 and not state.processed_indices and not state.buffered_transcripts ) if should_prime_state: persisted_contiguous_index = await _get_persisted_contiguous_segment_index( db=db, conversation_id=conversation_id, voice_session_id=voice_session_id, ) if persisted_contiguous_index >= 0: async with state.lock: state.consumed_index = max(state.consumed_index, persisted_contiguous_index) transcript_text = await asr_service.transcribe(audio_base64) await manager.send_message(conversation_id, { "type": MessageType.TRANSCRIPT, "conversation_id": conversation_id, "data": { "text": transcript_text or "", "audio_duration": audio_duration, "voice_session_id": voice_session_id, "segment_index": segment_index, "is_last": is_last, }, "timestamp": datetime.now(timezone.utc).isoformat(), }) if _is_transcribe_failure(transcript_text): await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": { "message": f"分段 {segment_index} 转写失败，请重试该片段", "segment_index": segment_index, }, "timestamp": datetime.now(timezone.utc).isoformat(), }) return existing_segment = await _find_existing_segment_by_index( db=db, conversation_id=conversation_id, voice_session_id=voice_session_id, segment_index=segment_index, ) if existing_segment: # 该分段已成功入库，视为重传：不重复入库、不重复触发 Agent。 async with state.lock: state.processed_indices.add(segment_index) logger.info( "分段已存在，按幂等处理跳过: " f"conversation_id={conversation_id}, voice_session_id={voice_session_id}, segment_index={segment_index}" ) return else: segment = Segment( id=str(uuid.uuid4()), conversation_id=conversation_id, transcript_text=transcript_text or "", audio_url=_build_segment_audio_url(voice_session_id, segment_index), processed=False, ) db.add(segment) await db.commit() await db.refresh(segment) await manager.background_runner.queue_message(conversation.user_id, segment.id) ready_segments: List[Tuple[int, str, Segment]] = [] async with state.lock: state.processed_indices.add(segment_index) state.buffered_transcripts[segment_index] = (transcript_text or "", segment) next_index = state.consumed_index + 1 while next_index in state.buffered_transcripts: text, seg = state.buffered_transcripts.pop(next_index) ready_segments.append((next_index, text, seg)) state.consumed_index = next_index next_index += 1 # 仅当前缀分段连续时才触发 Agent，保证增量上下文顺序正确。 for _, ordered_text, ordered_segment in ready_segments: await process_user_message( conversation_id=conversation_id, user_message=ordered_text, conversation=conversation, segment=ordered_segment, db=db, manager=manager, user=user, ) break except Exception as e: logger.error( f"处理语音分段失败: conversation_id={conversation_id}, segment_index={segment_index}, error={e}", exc_info=True, ) await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": { "message": f"分段处理失败: {str(e)}", "segment_index": segment_index, }, "timestamp": datetime.now(timezone.utc).isoformat(), }) finally: async with state.lock: state.pending_indices.discard(segment_index) async def websocket_endpoint( websocket: WebSocket, conversation_id: str ): """ WebSocket 端点：处理实时对话 Args: websocket: WebSocket 连接 conversation_id: 对话 ID """ # 从查询参数获取token token = websocket.query_params.get("token") if not token: await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="缺少访问令牌") return # 验证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: try: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "无权访问此对话"}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception: pass # 如果发送失败，直接关闭连接 await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="无权访问此对话") return # 首次连接时检查资料完整性，发送资料收集开场白 missing_profile = _get_missing_profile_fields(user) if missing_profile: try: greetings = await manager.conversation_agent.generate_profile_greeting( conversation_id=conversation_id, missing_fields=missing_profile, nickname=user.nickname or "", ) import asyncio as _asyncio_greet for i, text in enumerate(greetings): await manager.send_message(conversation_id, { "type": MessageType.AGENT_RESPONSE, "conversation_id": conversation_id, "data": {"text": text, "index": i, "total": len(greetings)}, "timestamp": datetime.now(timezone.utc).isoformat() }) if i < len(greetings) - 1: await _asyncio_greet.sleep(0.5) except Exception as e: logger.error(f"发送资料收集开场白失败: {e}", exc_info=True) # 主循环：处理消息 while True: try: if websocket.application_state != WebSocketState.CONNECTED: logger.info(f"WebSocket 已非连接状态，退出循环: conversation_id={conversation_id}") break 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: # 校验对话轮数配额 from routers.quota import get_segment_count, check_can_send_message seg_count = await get_segment_count(user_id, db) can_send, quota_msg = check_can_send_message(user.subscription_type, seg_count) if not can_send: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": quota_msg, "code": "QUOTA_EXCEEDED"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue # 保存段落到数据库 segment = Segment( id=str(uuid.uuid4()), conversation_id=conversation_id, transcript_text=text_message, processed=False ) db.add(segment) await db.commit() await db.refresh(segment) await manager.background_runner.queue_message(conversation.user_id, segment.id) # Agent 生成回应 await process_user_message( conversation_id=conversation_id, user_message=text_message, conversation=conversation, segment=segment, db=db, manager=manager, user=user, ) elif msg_type == MessageType.AUDIO_SEGMENT: # 处理分段语音消息（长语音持续上传） data = message.get("data", {}) audio_base64 = data.get("audio_base64", "") segment_index_raw = data.get("segment_index") voice_session_id = _normalize_voice_session_id( data.get("voice_session_id") or _voice_session_id_from_client_segment_id(data.get("client_segment_id")) ) is_last = bool(data.get("is_last", False)) audio_duration = int(data.get("duration", 0) or 0) if not audio_base64: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "缺少 audio_base64"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue if segment_index_raw is None: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "缺少 segment_index"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue try: segment_index = int(segment_index_raw) except (TypeError, ValueError): await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "segment_index 必须为整数"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue if segment_index < 0: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "segment_index 不能为负数"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue # 校验对话轮数配额（分段也计入对话轮次） from routers.quota import get_segment_count, check_can_send_message seg_count = await get_segment_count(user_id, db) can_send, quota_msg = check_can_send_message(user.subscription_type, seg_count) if not can_send: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": quota_msg, "code": "QUOTA_EXCEEDED"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue segment_state = manager.get_or_create_segment_state( conversation_id, voice_session_id, ) should_process = False async with segment_state.lock: already_seen = ( segment_index in segment_state.pending_indices or segment_index in segment_state.processed_indices or segment_index <= segment_state.consumed_index ) if not already_seen: segment_state.pending_indices.add(segment_index) should_process = True if not should_process: logger.info( "收到重复分段，跳过处理: " f"conversation_id={conversation_id}, voice_session_id={voice_session_id}, segment_index={segment_index}" ) continue # 先发过渡反馈，减少“等待空白”体感 await _send_segment_transition_feedback( conversation_id=conversation_id, segment_index=segment_index, manager=manager, ) task = asyncio.create_task( _process_audio_segment_async( conversation_id=conversation_id, user_id=user_id, voice_session_id=voice_session_id, segment_index=segment_index, audio_base64=audio_base64, audio_duration=audio_duration, is_last=is_last, manager=manager, ) ) manager.register_segment_task(conversation_id, voice_session_id, task) elif msg_type == MessageType.AUDIO_MESSAGE: # 处理完整音频消息（类似微信语音） data = message.get("data", {}) audio_base64 = data.get("audio_base64", "") audio_duration = data.get("duration", 0) if audio_base64: # 校验对话轮数配额 from routers.quota import get_segment_count, check_can_send_message seg_count = await get_segment_count(user_id, db) can_send, quota_msg = check_can_send_message(user.subscription_type, seg_count) if not can_send: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": quota_msg, "code": "QUOTA_EXCEEDED"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue logger.info(f"收到音频消息，时长: {audio_duration}s") try: # 1. ASR 转写 transcript_text = await asr_service.transcribe(audio_base64) logger.info(f"ASR 转写结果: {transcript_text}") # 2. 发送转写结果给客户端 await manager.send_message(conversation_id, { "type": MessageType.TRANSCRIPT, "conversation_id": conversation_id, "data": { "text": transcript_text, "audio_duration": audio_duration }, "timestamp": datetime.now(timezone.utc).isoformat() }) # 3. 保存段落到数据库（包含转写文本和音频信息） segment = Segment( id=str(uuid.uuid4()), conversation_id=conversation_id, transcript_text=transcript_text, audio_url=f"audio:{audio_duration}s", # 简化存储，标记为音频消息 processed=False ) db.add(segment) await db.commit() await db.refresh(segment) await manager.background_runner.queue_message(conversation.user_id, segment.id) # 4. Agent 生成回应（基于转写文本） if transcript_text and not transcript_text.startswith("转写失败"): await process_user_message( conversation_id=conversation_id, user_message=transcript_text, conversation=conversation, segment=segment, db=db, manager=manager, user=user, ) else: # 转写失败，发送错误消息 await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "语音转写失败，请重试或使用文字输入"}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception as e: logger.error(f"处理音频消息失败: {e}", exc_info=True) await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": f"处理音频消息失败: {str(e)}"}, "timestamp": datetime.now(timezone.utc).isoformat() }) elif msg_type == MessageType.TRANSCRIBE_ONLY: # 仅转写：不落库、不触发 Agent，用于客户端「转文字」后发文本 data = message.get("data", {}) audio_base64 = data.get("audio_base64", "") if not audio_base64: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": "缺少 audio_base64"}, "timestamp": datetime.now(timezone.utc).isoformat() }) continue try: transcript_text = await asr_service.transcribe(audio_base64) await manager.send_message(conversation_id, { "type": MessageType.TRANSCRIPT, "conversation_id": conversation_id, "data": {"text": transcript_text or ""}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception as e: logger.error(f"仅转写失败: {e}", exc_info=True) await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": f"转写失败: {str(e)}"}, "timestamp": datetime.now(timezone.utc).isoformat() }) 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 RuntimeError as e: # 检查是否是断开连接或未连接状态（如 accept 前/后连接被关闭） error_msg = str(e) if ( "disconnect" in error_msg.lower() or "Cannot call \"receive\"" in error_msg or "accept" in error_msg.lower() and "not connected" in error_msg.lower() ): logger.info(f"WebSocket 连接已断开或未就绪: conversation_id={conversation_id}, error={error_msg}") break else: logger.error(f"处理消息时发生 RuntimeError: {e}", exc_info=True) # 只在连接仍然活跃时发送错误消息 if conversation_id in manager.active_connections: try: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": str(e)}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception as send_error: logger.warning(f"发送错误消息失败: {send_error}") break except WebSocketDisconnect: logger.info(f"WebSocket 断开连接: conversation_id={conversation_id}") break except Exception as e: logger.error(f"处理消息时发生错误: {e}", exc_info=True) # 只在连接仍然活跃时发送错误消息 if conversation_id in manager.active_connections: try: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": str(e)}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception as send_error: logger.warning(f"发送错误消息失败: {send_error}") break except WebSocketDisconnect: logger.info(f"WebSocket 断开连接: conversation_id={conversation_id}") await manager.disconnect(conversation_id) except Exception as e: logger.error(f"WebSocket 端点发生错误: {e}", exc_info=True) await manager.disconnect(conversation_id) finally: # 确保清理连接 await manager.disconnect(conversation_id) def _get_missing_profile_fields(user: UserModel) -> list: """检查用户缺失的资料字段""" from agents.prompts.profile_prompts import get_missing_profile_fields return get_missing_profile_fields( birth_year=user.birth_year, birth_place=user.birth_place, grew_up_place=user.grew_up_place, occupation=user.occupation, ) def _get_filled_profile_fields(user: UserModel) -> dict: """获取用户已有的资料字段（中文展示）""" from agents.prompts.profile_prompts import PROFILE_FIELD_NAMES filled = {} if user.birth_year: filled["birth_year"] = str(user.birth_year) if user.birth_place: filled["birth_place"] = user.birth_place if user.grew_up_place: filled["grew_up_place"] = user.grew_up_place if user.occupation: filled["occupation"] = user.occupation return filled async def _apply_extracted_profile(user: UserModel, extracted: dict, db: AsyncSession): """将提取到的资料信息保存到用户模型""" changed = False if "birth_year" in extracted and not user.birth_year: user.birth_year = extracted["birth_year"] changed = True if "birth_place" in extracted and not user.birth_place: user.birth_place = extracted["birth_place"] changed = True if "grew_up_place" in extracted and not user.grew_up_place: user.grew_up_place = extracted["grew_up_place"] changed = True if "occupation" in extracted and not user.occupation: user.occupation = extracted["occupation"] changed = True if changed: await db.commit() await db.refresh(user) async def process_user_message( conversation_id: str, user_message: str, conversation: Conversation, segment: Segment, db: AsyncSession, manager: ConnectionManager, user: UserModel = None, ) -> None: """ 处理用户消息，生成Agent回应（异步版本） 支持资料收集模式和正式访谈模式 """ import asyncio as _asyncio agent = manager.conversation_agent # --- 资料收集模式 --- if user: missing = _get_missing_profile_fields(user) if missing: try: extracted = await agent.extract_profile_from_message(user_message, missing) if extracted: await _apply_extracted_profile(user, extracted, db) remaining = _get_missing_profile_fields(user) filled = _get_filled_profile_fields(user) responses = await agent.generate_profile_followup( conversation_id=conversation_id, user_message=user_message, missing_fields=remaining, filled_fields=filled, nickname=user.nickname or "", ) segment.agent_response = "\n\n".join(responses) await db.commit() for i, response_text in enumerate(responses): await manager.send_message(conversation_id, { "type": MessageType.AGENT_RESPONSE, "conversation_id": conversation_id, "data": {"text": response_text, "index": i, "total": len(responses)}, "timestamp": datetime.now(timezone.utc).isoformat() }) if i < len(responses) - 1: await _asyncio.sleep(0.5) return except Exception as e: logger.error(f"资料收集处理失败: {e}", exc_info=True) # --- 正式访谈模式 --- state = await get_or_create_state(conversation.user_id, db) if conversation.conversation_stage != state.current_stage: conversation.conversation_stage = state.current_stage 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] # 构建用户资料上下文 user_profile_context = "" if user: from agents.prompts.profile_prompts import format_user_profile_context user_profile_context = format_user_profile_context( birth_year=user.birth_year, birth_place=user.birth_place, grew_up_place=user.grew_up_place, occupation=user.occupation, ) try: responses = await agent.generate_response_with_state( conversation_id=conversation_id, user_message=user_message, memoir_state=state, user_profile_context=user_profile_context, ) segment.agent_response = "\n\n".join(responses) await db.commit() for i, response_text in enumerate(responses): await manager.send_message(conversation_id, { "type": MessageType.AGENT_RESPONSE, "conversation_id": conversation_id, "data": {"text": response_text, "index": i, "total": len(responses)}, "timestamp": datetime.now(timezone.utc).isoformat() }) if i < len(responses) - 1: await _asyncio.sleep(0.5) except Exception as e: logger.error(f"处理用户消息失败: {e}", exc_info=True) if conversation_id in manager.active_connections: try: await manager.send_message(conversation_id, { "type": MessageType.ERROR, "data": {"message": f"生成回应失败: {str(e)}"}, "timestamp": datetime.now(timezone.utc).isoformat() }) except Exception as send_error: logger.warning(f"发送错误消息失败: {send_error}") async def process_conversation_segments(conversation_id: str, db: AsyncSession): """ 处理对话段落，生成章节（对话结束时调用） 注意：大部分处理已通过 Celery 任务增量完成 这里立即提交所有待处理的段落到 Celery Args: conversation_id: 对话 ID db: 数据库会话 """ # 获取对话信息 conversation = await db.get(Conversation, conversation_id) if not conversation: return # 获取所有未处理的段落 stmt = select(Segment).where( Segment.conversation_id == conversation_id, Segment.processed == False ) result = await db.execute(stmt) segments = result.scalars().all() if not segments: # 没有未处理的段落，直接 flush 待处理任务 await manager.background_runner.flush_pending(conversation.user_id) return # 免费版仅允许 1 个章节整理，提交前校验 from database.models import User as UserModel from routers.quota import get_chapter_count, check_can_submit_organize user = await db.get(UserModel, conversation.user_id) if user: chapter_count = await get_chapter_count(user.id, db) can_submit, _ = check_can_submit_organize(user.subscription_type, chapter_count) if not can_submit: logger.info( f"用户 {user.id} 章节配额已用尽，跳过提交整理任务: conversation_id={conversation_id}" ) await manager.background_runner.flush_pending(conversation.user_id) return # 将未处理的段落直接提交到 Celery（不通过去抖） segment_ids = [seg.id for seg in segments] try: from tasks.memoir_tasks import process_memoir_segments process_memoir_segments.delay(conversation.user_id, segment_ids) logger.info(f"对话结束，提交 Celery 任务: conversation_id={conversation_id}, segments={len(segment_ids)}") except Exception as e: logger.error(f"提交 Celery 任务失败: {e}") # 同时 flush 任何待处理的任务 await manager.background_runner.flush_pending(conversation.user_id)