fix: WebSocket数据库会话问题修复 - 每次消息处理使用新会话
This commit is contained in:
303
main_v2.py
303
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)
|
||||
|
||||
Reference in New Issue
Block a user