Compare commits
22 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| f0789d6bbc | |||
| e05233fb4f | |||
| 7fa143b5b0 | |||
| b573638bf8 | |||
| af997aa5c5 | |||
| b1feaee976 | |||
| 87f9f4a7d8 | |||
| a1f1032000 | |||
| 7d6a345a7d | |||
| c6f157aa97 | |||
| 6e87f59fab | |||
| 066d2fe44d | |||
| 23935a1a28 | |||
| 08c4f79313 | |||
| 186f69c87a | |||
| 8c90fd5641 | |||
| 248ac4e471 | |||
| 51cc8161f1 | |||
| a62fe929c1 | |||
| 6adeb9b371 | |||
| 051fd5c1c8 | |||
| 90d31dba69 |
340
main_v2.py
340
main_v2.py
@@ -60,8 +60,13 @@ class ConnectionManager:
|
||||
for connection in self.active_connections[user_id]:
|
||||
try:
|
||||
await connection.send_json(message)
|
||||
except:
|
||||
pass
|
||||
logger.info(f"发送消息到用户 {user_id}: {message.get('type', 'unknown')}")
|
||||
except Exception as e:
|
||||
logger.error(f"发送消息失败: {e}")
|
||||
|
||||
async def ping(self, user_id: str):
|
||||
"""发送心跳ping"""
|
||||
await self.send_to_user(user_id, {"type": "ping"})
|
||||
|
||||
|
||||
manager = ConnectionManager()
|
||||
@@ -511,181 +516,210 @@ 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)
|
||||
|
||||
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:
|
||||
data = await websocket.receive_json()
|
||||
except Exception as json_err:
|
||||
logger.error(f"JSON解析错误: {json_err}")
|
||||
# 如果连接已断开，退出循环
|
||||
if "disconnect" in str(json_err).lower() or "closed" in str(json_err).lower():
|
||||
logger.info("WebSocket已断开，退出循环")
|
||||
break
|
||||
try:
|
||||
# 发送"正在思考"状态
|
||||
if agent_config['agent'].get('enable_thinking') and enable_thinking:
|
||||
text_data = await websocket.receive_text()
|
||||
if text_data.strip():
|
||||
data = json.loads(text_data)
|
||||
else:
|
||||
continue
|
||||
except Exception as text_err:
|
||||
logger.error(f"文本消息解析错误: {text_err}")
|
||||
if "disconnect" in str(text_err).lower() or "closed" in str(text_err).lower():
|
||||
logger.info("WebSocket已断开，退出循环")
|
||||
break
|
||||
continue
|
||||
|
||||
action = data.get("action")
|
||||
logger.info(f"WebSocket收到消息: action={action}")
|
||||
|
||||
# 每次消息处理时创建新的数据库会话，处理完后关闭
|
||||
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 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)
|
||||
|
||||
# 获取对话使用的Agent ID
|
||||
conv_agent_id = conversation.current_agent_id
|
||||
|
||||
await websocket.send_json({
|
||||
"type": "thinking_start",
|
||||
"type": "history",
|
||||
"conversation_id": current_conversation_id,
|
||||
"agent_id": conv_agent_id, # 返回对话的Agent ID
|
||||
"messages": [
|
||||
{
|
||||
"role": m.role,
|
||||
"content": m.content,
|
||||
"thinking_content": m.thinking_content,
|
||||
"agent_id": m.agent_id, # 每条消息的Agent ID
|
||||
"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:
|
||||
# 调用LLM（非流式）
|
||||
response, thinking_content = await llm_service.chat(
|
||||
messages=history,
|
||||
provider_config=agent_config['provider'],
|
||||
agent_config=agent_config['agent'],
|
||||
enable_thinking=enable_thinking
|
||||
)
|
||||
|
||||
logger.info(f"LLM响应: response长度={len(response)}, thinking长度={len(thinking_content) if thinking_content else 0}")
|
||||
|
||||
# 保存AI回复
|
||||
assistant_msg = conv_service.add_message(
|
||||
conversation_id=conversation.id,
|
||||
role='assistant',
|
||||
content=response,
|
||||
source='web',
|
||||
thinking_content=thinking_content if thinking_content else None,
|
||||
agent_id=current_agent_id,
|
||||
model_used=agent_config['provider'].get('default_model')
|
||||
)
|
||||
|
||||
# 发送完整回复（包含思考内容）
|
||||
await websocket.send_json({
|
||||
"type": "assistant_message",
|
||||
"conversation_id": conversation_id,
|
||||
"message": {
|
||||
"id": assistant_msg.id,
|
||||
"role": "assistant",
|
||||
"content": response,
|
||||
"thinking_content": thinking_content if thinking_content else None,
|
||||
"source": "web",
|
||||
"agent_id": current_agent_id,
|
||||
"agent_name": agent_config['agent'].get('display_name'),
|
||||
"created_at": assistant_msg.created_at.isoformat()
|
||||
}
|
||||
})
|
||||
|
||||
logger.info(f"AI回复已发送: conversation_id={conversation_id}")
|
||||
|
||||
# 启用发送按钮
|
||||
await websocket.send_json({
|
||||
"type": "stream_end",
|
||||
"conversation_id": conversation_id
|
||||
})
|
||||
|
||||
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)
|
||||
|
||||
@@ -128,6 +128,8 @@ class LLMService:
|
||||
# 处理思考功能
|
||||
if enable_thinking and agent_config.get('enable_thinking', True):
|
||||
thinking_prompt = agent_config.get('thinking_prompt')
|
||||
thinking_prefix = agent_config.get('thinking_prefix', '')
|
||||
thinking_suffix = agent_config.get('thinking_suffix', '')
|
||||
|
||||
if supports_thinking and thinking_model:
|
||||
# 使用专门的思考模型
|
||||
@@ -139,22 +141,15 @@ class LLMService:
|
||||
thinking_result = await self._call_api(
|
||||
api_base, api_key, thinking_model, thinking_messages,
|
||||
max_tokens=min(max_tokens, 1000),
|
||||
temperature=0.3 # 思考时降低温度
|
||||
temperature=0.3
|
||||
)
|
||||
thinking_content = thinking_result
|
||||
except Exception as e:
|
||||
logger.warning(f"思考模型调用失败: {e}")
|
||||
|
||||
elif supports_thinking:
|
||||
# Provider支持思考但无单独模型，尝试在回复中获取思考部分
|
||||
pass # 在回复解析时处理
|
||||
|
||||
elif thinking_prompt:
|
||||
# Provider不支持思考，但Agent配置了思考提示词
|
||||
# 将思考提示词添加到系统提示
|
||||
enhanced_system = f"{system_prompt}\n\n在回答之前，请先思考问题。思考过程请用{agent_config.get('thinking_prefix', '')}和{agent_config.get('thinking_suffix', '')}包裹，然后再给出正式回答。"
|
||||
if thinking_prompt:
|
||||
enhanced_system += f"\n思考指导：{thinking_prompt}"
|
||||
# Agent配置了思考提示词，添加到系统提示中
|
||||
enhanced_system = f"{system_prompt}\n\n{thinking_prompt}"
|
||||
final_messages[0] = {"role": "system", "content": enhanced_system}
|
||||
|
||||
# 调用主模型
|
||||
@@ -165,19 +160,37 @@ class LLMService:
|
||||
temperature=temperature
|
||||
)
|
||||
|
||||
# 尝试从回复中提取思考内容
|
||||
if enable_thinking and not supports_thinking:
|
||||
# 尝试从回复中提取思考内容（支持DeepSeek R1、GLM等模型的思考模式）
|
||||
if enable_thinking and agent_config.get('enable_thinking', True):
|
||||
thinking_prefix = agent_config.get('thinking_prefix', '')
|
||||
thinking_suffix = agent_config.get('thinking_suffix', '')
|
||||
|
||||
if thinking_prefix and thinking_suffix:
|
||||
# 提取思考部分
|
||||
pattern = f"{re.escape(thinking_prefix)}(.*?)?{re.escape(thinking_suffix)}"
|
||||
match = re.search(pattern, response, re.DOTALL)
|
||||
if match:
|
||||
thinking_content = match.group(1).strip()
|
||||
# 移除思考部分，只保留回复
|
||||
response = re.sub(pattern, '', response, flags=re.DOTALL).strip()
|
||||
# 如果没有配置前缀后缀，使用常见的思考标记
|
||||
if not thinking_prefix:
|
||||
# 尝试常见的思考标记
|
||||
common_thinking_markers = [
|
||||
('<think>', '</think>'),
|
||||
('【思考】', '【回答】'),
|
||||
('Thought:', 'Answer:'),
|
||||
('思考：', '回答：'),
|
||||
]
|
||||
for prefix, suffix in common_thinking_markers:
|
||||
if prefix in response and suffix in response:
|
||||
thinking_prefix = prefix
|
||||
thinking_suffix = suffix
|
||||
break
|
||||
|
||||
# 提取思考部分
|
||||
if thinking_prefix and thinking_suffix and thinking_prefix in response:
|
||||
try:
|
||||
start_idx = response.find(thinking_prefix)
|
||||
end_idx = response.find(thinking_suffix, start_idx)
|
||||
if end_idx > start_idx:
|
||||
thinking_content = response[start_idx + len(thinking_prefix):end_idx].strip()
|
||||
# 移除思考部分，只保留回复
|
||||
response = response[end_idx + len(thinking_suffix):].strip()
|
||||
except Exception as e:
|
||||
logger.warning(f"提取思考内容失败: {e}")
|
||||
|
||||
return response, thinking_content
|
||||
|
||||
@@ -268,8 +281,7 @@ class LLMService:
|
||||
|
||||
thinking_prefix = agent_config.get('thinking_prefix', '')
|
||||
thinking_suffix = agent_config.get('thinking_suffix', '')
|
||||
in_thinking = False
|
||||
thinking_buffer = ""
|
||||
buffer = "" # 用于累积和检测思考部分
|
||||
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
async with client.stream("POST", url, headers=headers, json=payload) as response:
|
||||
@@ -284,35 +296,46 @@ class LLMService:
|
||||
delta = data['choices'][0].get('delta', {})
|
||||
if 'content' in delta:
|
||||
text = delta['content']
|
||||
buffer += text
|
||||
|
||||
# 检测思考部分
|
||||
if thinking_prefix and thinking_suffix:
|
||||
for char in text:
|
||||
if in_thinking:
|
||||
thinking_buffer += char
|
||||
# 检查是否结束思考
|
||||
if thinking_buffer.endswith(thinking_suffix):
|
||||
thinking_content = thinking_buffer[:-len(thinking_suffix)]
|
||||
yield {"type": "thinking", "text": thinking_content}
|
||||
in_thinking = False
|
||||
thinking_buffer = ""
|
||||
# 检测思考部分（简化逻辑）
|
||||
if thinking_prefix and thinking_suffix and thinking_prefix in buffer:
|
||||
# 尝试解析思考部分
|
||||
try:
|
||||
start_idx = buffer.find(thinking_prefix)
|
||||
if start_idx >= 0:
|
||||
# 找到思考开始，继续找结束
|
||||
end_idx = buffer.find(thinking_suffix, start_idx)
|
||||
if end_idx > start_idx:
|
||||
# 思考部分完整，发送思考然后发送内容
|
||||
thinking = buffer[start_idx + len(thinking_prefix):end_idx]
|
||||
yield {"type": "thinking", "text": thinking}
|
||||
# 发送思考后的内容
|
||||
remaining = buffer[end_idx + len(thinking_suffix):]
|
||||
if remaining:
|
||||
yield {"type": "content", "text": remaining}
|
||||
buffer = ""
|
||||
else:
|
||||
# 检查是否接近结束
|
||||
suffix_len = len(thinking_suffix)
|
||||
if len(thinking_buffer) >= suffix_len:
|
||||
yield {"type": "thinking", "text": thinking_buffer[-suffix_len:]}
|
||||
# 思考部分还没结束，先发送之前的内容
|
||||
if start_idx > 0:
|
||||
yield {"type": "content", "text": buffer[:start_idx]}
|
||||
# 等待更多数据来完成思考部分
|
||||
buffer = buffer[start_idx:]
|
||||
else:
|
||||
if char == thinking_prefix[0]:
|
||||
# 可能开始思考
|
||||
thinking_buffer = char
|
||||
if len(thinking_prefix) == 1:
|
||||
in_thinking = True
|
||||
else:
|
||||
yield {"type": "content", "text": char}
|
||||
# 没有思考标记，直接发送内容
|
||||
yield {"type": "content", "text": text}
|
||||
buffer = ""
|
||||
except:
|
||||
yield {"type": "content", "text": text}
|
||||
else:
|
||||
# 没有思考标记配置，直接发送内容
|
||||
yield {"type": "content", "text": text}
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# 处理剩余buffer
|
||||
if buffer:
|
||||
yield {"type": "content", "text": buffer}
|
||||
|
||||
|
||||
# 全局实例
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user