From 90d31dba69eee48af55153c0cd827fd05213a90e Mon Sep 17 00:00:00 2001 From: hubian <908234780@qq.com> Date: Sun, 12 Apr 2026 16:49:39 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20WebSocket=E6=95=B0=E6=8D=AE=E5=BA=93?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E9=97=AE=E9=A2=98=E4=BF=AE=E5=A4=8D=20-=20?= =?UTF-8?q?=E6=AF=8F=E6=AC=A1=E6=B6=88=E6=81=AF=E5=A4=84=E7=90=86=E4=BD=BF?= =?UTF-8?q?=E7=94=A8=E6=96=B0=E4=BC=9A=E8=AF=9D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main_v2.py | 303 +++++++++++++++++++++++++++-------------------------- 1 file changed, 155 insertions(+), 148 deletions(-) diff --git a/main_v2.py b/main_v2.py index 60bb908..4986ee9 100644 --- a/main_v2.py +++ b/main_v2.py @@ -511,181 +511,188 @@ async def delete_conversation(conversation_id: str, db: Session = Depends(get_db # ==================== WebSocket路由 ==================== @app.websocket("/ws/{user_id}") -async def websocket_endpoint(websocket: WebSocket, user_id: str, db: Session = Depends(get_db)): +async def websocket_endpoint(websocket: WebSocket, user_id: str): """WebSocket连接 - 实时对话""" actual_user_id = MAIN_USER_ID await manager.connect(websocket, actual_user_id) - conv_service = ConversationService(db) - user = conv_service.get_or_create_user(MAIN_USER_ID, display_name="主用户", user_type='web') - # 获取默认Agent配置 - agent_service = AgentService(db) - default_agent = agent_service.get_default_agent() + # 初始化时获取默认Agent ID + db = SessionLocal() + try: + agent_service = AgentService(db) + default_agent = agent_service.get_default_agent() + default_agent_id = default_agent.id if default_agent else None + finally: + db.close() current_conversation_id = None - current_agent_id = default_agent.id if default_agent else None + current_agent_id = default_agent_id try: while True: data = await websocket.receive_json() action = data.get("action") - if action == "select_conversation": - current_conversation_id = data.get("conversation_id") - conversation = conv_service.get_conversation(current_conversation_id) + # 每次消息处理时创建新的数据库会话，处理完后关闭 + try: + db = SessionLocal() + conv_service = ConversationService(db) + agent_service = AgentService(db) + user = conv_service.get_or_create_user(MAIN_USER_ID, display_name="主用户", user_type='web') - if conversation: - messages = conv_service.get_messages(conversation.id) - await websocket.send_json({ - "type": "history", - "conversation_id": current_conversation_id, - "messages": [ - { - "role": m.role, - "content": m.content, - "thinking_content": m.thinking_content, - "source": m.source, - "created_at": m.created_at.isoformat() - } - for m in messages - ] - }) - - elif action == "switch_agent": - # 切换Agent - new_agent_id = data.get("agent_id") - agent = agent_service.get_agent(new_agent_id) - if agent and agent.is_active: - current_agent_id = new_agent_id - await websocket.send_json({ - "type": "agent_switched", - "agent_id": current_agent_id, - "agent_name": agent.display_name or agent.name - }) - - elif action == "chat": - message = data.get("message", "") - conversation_id = data.get("conversation_id") - enable_thinking = data.get("enable_thinking", True) # 可临时关闭思考 - agent_id_override = data.get("agent_id") # 前端可以指定agent - - # 如果前端指定了agent，使用它 - if agent_id_override: - agent = agent_service.get_agent(agent_id_override) - if agent and agent.is_active: - current_agent_id = agent_id_override - - if not message.strip(): - continue - - # 获取或创建会话 - if conversation_id: - conversation = conv_service.get_conversation(conversation_id) - else: - conversation = conv_service.create_conversation(user.id) - conversation_id = conversation.conversation_id - await websocket.send_json({ - "type": "conversation_created", - "conversation_id": conversation_id - }) - - # 保存用户消息 - user_msg = conv_service.add_message( - conversation_id=conversation.id, - role='user', - content=message, - source='web' - ) - - # 广播用户消息 - await manager.send_to_user(MAIN_USER_ID, { - "type": "user_message", - "conversation_id": conversation_id, - "message": { - "id": user_msg.id, - "role": "user", - "content": message, - "source": "web", - "created_at": user_msg.created_at.isoformat() - } - }) - - # 获取Agent配置 - agent_config = agent_service.get_agent_config(current_agent_id) - - if not agent_config or not agent_config.get('provider'): - await websocket.send_json({ - "type": "error", - "message": "Agent配置不完整" - }) - continue - - # 获取对话历史 - history = conv_service.get_conversation_history(conversation_id, limit=agent_config['agent'].get('max_history', 20)) - - # 调用LLM - try: - # 发送"正在思考"状态 - if agent_config['agent'].get('enable_thinking') and enable_thinking: + if action == "select_conversation": + current_conversation_id = data.get("conversation_id") + conversation = conv_service.get_conversation(current_conversation_id) + + if conversation: + messages = conv_service.get_messages(conversation.id) await websocket.send_json({ - "type": "thinking_start", + "type": "history", + "conversation_id": current_conversation_id, + "messages": [ + { + "role": m.role, + "content": m.content, + "thinking_content": m.thinking_content, + "source": m.source, + "created_at": m.created_at.isoformat() + } + for m in messages + ] + }) + + elif action == "switch_agent": + # 切换Agent + new_agent_id = data.get("agent_id") + agent = agent_service.get_agent(new_agent_id) + if agent and agent.is_active: + current_agent_id = new_agent_id + await websocket.send_json({ + "type": "agent_switched", + "agent_id": current_agent_id, + "agent_name": agent.display_name or agent.name + }) + + elif action == "chat": + message = data.get("message", "") + conversation_id = data.get("conversation_id") + enable_thinking = data.get("enable_thinking", True) + agent_id_override = data.get("agent_id") + + if agent_id_override: + agent = agent_service.get_agent(agent_id_override) + if agent and agent.is_active: + current_agent_id = agent_id_override + + if not message.strip(): + continue + + # 获取或创建会话 + if conversation_id: + conversation = conv_service.get_conversation(conversation_id) + else: + conversation = conv_service.create_conversation(user.id) + conversation_id = conversation.conversation_id + await websocket.send_json({ + "type": "conversation_created", "conversation_id": conversation_id }) - response, thinking_content = await llm_service.chat( - messages=history, - provider_config=agent_config['provider'], - agent_config=agent_config['agent'], - enable_thinking=enable_thinking - ) - - # 发送思考内容 - if thinking_content: - await websocket.send_json({ - "type": "thinking_content", - "conversation_id": conversation_id, - "content": thinking_content - }) - - # 发送思考结束 - await websocket.send_json({ - "type": "thinking_end", - "conversation_id": conversation_id - }) - - # 保存AI回复 - assistant_msg = conv_service.add_message( + # 保存用户消息 + user_msg = conv_service.add_message( conversation_id=conversation.id, - role='assistant', - content=response, - source='web', - thinking_content=thinking_content, - agent_id=current_agent_id, - model_used=agent_config['provider'].get('default_model') + role='user', + content=message, + source='web' ) - # 广播AI回复 + # 广播用户消息 await manager.send_to_user(MAIN_USER_ID, { - "type": "assistant_message", + "type": "user_message", "conversation_id": conversation_id, "message": { - "id": assistant_msg.id, - "role": "assistant", - "content": response, - "thinking_content": thinking_content, + "id": user_msg.id, + "role": "user", + "content": message, "source": "web", - "agent_id": current_agent_id, - "agent_name": agent_config['agent'].get('display_name'), - "created_at": assistant_msg.created_at.isoformat() + "created_at": user_msg.created_at.isoformat() } }) - except Exception as e: - logger.error(f"LLM调用失败: {e}") - await websocket.send_json({ - "type": "error", - "message": f"AI服务暂时不可用: {str(e)}" - }) + # 获取Agent配置 + agent_config = agent_service.get_agent_config(current_agent_id) + + if not agent_config or not agent_config.get('provider'): + await websocket.send_json({ + "type": "error", + "message": "Agent配置不完整" + }) + continue + + # 获取对话历史 + history = conv_service.get_conversation_history(conversation_id, limit=agent_config['agent'].get('max_history', 20)) + + # 调用LLM + try: + if agent_config['agent'].get('enable_thinking') and enable_thinking: + await websocket.send_json({ + "type": "thinking_start", + "conversation_id": conversation_id + }) + + response, thinking_content = await llm_service.chat( + messages=history, + provider_config=agent_config['provider'], + agent_config=agent_config['agent'], + enable_thinking=enable_thinking + ) + + if thinking_content: + await websocket.send_json({ + "type": "thinking_content", + "conversation_id": conversation_id, + "content": thinking_content + }) + + await websocket.send_json({ + "type": "thinking_end", + "conversation_id": conversation_id + }) + + assistant_msg = conv_service.add_message( + conversation_id=conversation.id, + role='assistant', + content=response, + source='web', + thinking_content=thinking_content, + agent_id=current_agent_id, + model_used=agent_config['provider'].get('default_model') + ) + + await manager.send_to_user(MAIN_USER_ID, { + "type": "assistant_message", + "conversation_id": conversation_id, + "message": { + "id": assistant_msg.id, + "role": "assistant", + "content": response, + "thinking_content": thinking_content, + "source": "web", + "agent_id": current_agent_id, + "agent_name": agent_config['agent'].get('display_name'), + "created_at": assistant_msg.created_at.isoformat() + } + }) + + except Exception as e: + logger.error(f"LLM调用失败: {e}") + await websocket.send_json({ + "type": "error", + "message": f"AI服务暂时不可用: {str(e)}" + }) + + finally: + db.close() except WebSocketDisconnect: manager.disconnect(websocket, user_id)