Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 813b4887ed |
230
main_v2.py
230
main_v2.py
@@ -777,103 +777,11 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
if not message.strip():
|
||||
continue
|
||||
|
||||
# 获取Agent配置
|
||||
# 1. 获取Agent配置
|
||||
agent_config = agent_service.get_agent_config(current_agent_id)
|
||||
agent_tools = agent_config.get('agent', {}).get('tools', [])
|
||||
|
||||
# 检查是否需要执行搜索
|
||||
search_context = None
|
||||
search_results_for_client = None # 用于发送给前端
|
||||
logger.info(f"检查搜索条件: agent_tools={agent_tools}, disabled_tools={disabled_tools}")
|
||||
|
||||
if 'search' in agent_tools and 'search' not in disabled_tools:
|
||||
# 只要启用了搜索工具且未禁用，就执行搜索
|
||||
logger.info("搜索条件满足，开始执行搜索")
|
||||
|
||||
# 执行搜索
|
||||
tool_service = ToolService(db)
|
||||
search_tool = tool_service.get_default_tool('search')
|
||||
logger.info(f"获取到搜索工具: {search_tool.name if search_tool else 'None'}")
|
||||
|
||||
if search_tool and search_tool.config.get('api_key'):
|
||||
# 执行搜索
|
||||
tool_service = ToolService(db)
|
||||
search_tool = tool_service.get_default_tool('search')
|
||||
|
||||
if search_tool and search_tool.config.get('api_key'):
|
||||
import httpx
|
||||
import time
|
||||
start_time = time.time()
|
||||
try:
|
||||
logger.info(f"执行搜索: query={message}")
|
||||
tavily_url = "https://api.tavily.com/search"
|
||||
config = search_tool.config
|
||||
payload = {
|
||||
"api_key": config.get('api_key'),
|
||||
"query": message,
|
||||
"max_results": config.get('max_results', 5),
|
||||
"search_depth": config.get('search_depth', 'basic')
|
||||
}
|
||||
|
||||
# 同步调用
|
||||
with httpx.Client(timeout=30) as client:
|
||||
resp = client.post(tavily_url, json=payload)
|
||||
search_result = resp.json()
|
||||
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
|
||||
if search_result.get("results"):
|
||||
# 构建搜索上下文（给LLM）
|
||||
search_context = "\n\n【搜索结果】\n"
|
||||
for i, r in enumerate(search_result["results"][:5], 1):
|
||||
search_context += f"{i}. {r.get('title', 'N/A')}\n {r.get('content', r.get('snippet', 'N/A'))[:200]}\n 来源: {r.get('url', 'N/A')}\n"
|
||||
logger.info(f"搜索完成: {len(search_result['results'])} 条结果")
|
||||
|
||||
# 构建搜索结果给前端展示
|
||||
search_results_for_client = [
|
||||
{
|
||||
"title": r.get('title', 'N/A'),
|
||||
"snippet": r.get('content', r.get('snippet', ''))[:150],
|
||||
"url": r.get('url', 'N/A')
|
||||
}
|
||||
for r in search_result["results"][:5]
|
||||
]
|
||||
|
||||
# 发送搜索结果给前端
|
||||
await websocket.send_json({
|
||||
"type": "search_results",
|
||||
"conversation_id": conversation_id,
|
||||
"results": search_results_for_client,
|
||||
"query": message
|
||||
})
|
||||
|
||||
# 更新统计和日志
|
||||
tool_service.increment_stats(search_tool.id, True)
|
||||
tool_service.log_usage({
|
||||
'tool_id': search_tool.id,
|
||||
'tool_type': 'search',
|
||||
'query': message,
|
||||
'success': True,
|
||||
'result_summary': f'{len(search_result["results"])} results',
|
||||
'conversation_id': conversation_id,
|
||||
'agent_id': current_agent_id,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
except Exception as e:
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.error(f"搜索失败: {e}")
|
||||
tool_service.increment_stats(search_tool.id, False)
|
||||
tool_service.log_usage({
|
||||
'tool_id': search_tool.id,
|
||||
'tool_type': 'search',
|
||||
'query': message,
|
||||
'success': False,
|
||||
'error_message': str(e),
|
||||
'conversation_id': conversation_id,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
|
||||
# 获取或创建会话
|
||||
# 2. 获取或创建会话（先有 conversation_id）
|
||||
if conversation_id:
|
||||
conversation = conv_service.get_conversation(conversation_id)
|
||||
else:
|
||||
@@ -884,7 +792,101 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
"conversation_id": conversation_id
|
||||
})
|
||||
|
||||
# 保存用户消息
|
||||
# 3. 广播用户消息（前端立即看到）
|
||||
await manager.send_to_user(MAIN_USER_ID, {
|
||||
"type": "user_message",
|
||||
"conversation_id": conversation_id,
|
||||
"message": {
|
||||
"id": None, # 临时，后面会保存
|
||||
"role": "user",
|
||||
"content": message,
|
||||
"source": "web",
|
||||
"created_at": datetime.utcnow().isoformat()
|
||||
}
|
||||
})
|
||||
|
||||
# 4. 执行搜索并发送搜索结果
|
||||
search_context = None
|
||||
logger.info(f"检查搜索条件: agent_tools={agent_tools}, disabled_tools={disabled_tools}")
|
||||
|
||||
if 'search' in agent_tools and 'search' not in disabled_tools:
|
||||
logger.info("搜索条件满足，开始执行搜索")
|
||||
|
||||
tool_service = ToolService(db)
|
||||
search_tool = tool_service.get_default_tool('search')
|
||||
logger.info(f"获取到搜索工具: {search_tool.name if search_tool else 'None'}")
|
||||
|
||||
if search_tool and search_tool.config.get('api_key'):
|
||||
import httpx
|
||||
import time
|
||||
start_time = time.time()
|
||||
try:
|
||||
logger.info(f"执行搜索: query={message}")
|
||||
tavily_url = "https://api.tavily.com/search"
|
||||
config = search_tool.config
|
||||
payload = {
|
||||
"api_key": config.get('api_key'),
|
||||
"query": message,
|
||||
"max_results": config.get('max_results', 5),
|
||||
"search_depth": config.get('search_depth', 'basic')
|
||||
}
|
||||
|
||||
with httpx.Client(timeout=30) as client:
|
||||
resp = client.post(tavily_url, json=payload)
|
||||
search_result = resp.json()
|
||||
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
|
||||
if search_result.get("results"):
|
||||
# 构建搜索上下文（给LLM）
|
||||
search_context = "\n\n【搜索结果】\n"
|
||||
for i, r in enumerate(search_result["results"][:5], 1):
|
||||
search_context += f"{i}. {r.get('title', 'N/A')}\n {r.get('content', r.get('snippet', 'N/A'))[:200]}\n 来源: {r.get('url', 'N/A')}\n"
|
||||
logger.info(f"搜索完成: {len(search_result['results'])} 条结果")
|
||||
|
||||
# 发送搜索结果给前端
|
||||
search_results_for_client = [
|
||||
{
|
||||
"title": r.get('title', 'N/A'),
|
||||
"snippet": r.get('content', r.get('snippet', ''))[:150],
|
||||
"url": r.get('url', 'N/A')
|
||||
}
|
||||
for r in search_result["results"][:5]
|
||||
]
|
||||
await websocket.send_json({
|
||||
"type": "search_results",
|
||||
"conversation_id": conversation_id,
|
||||
"results": search_results_for_client,
|
||||
"query": message
|
||||
})
|
||||
|
||||
# 更新统计和日志
|
||||
tool_service.increment_stats(search_tool.id, True)
|
||||
tool_service.log_usage({
|
||||
'tool_id': search_tool.id,
|
||||
'tool_type': 'search',
|
||||
'query': message,
|
||||
'success': True,
|
||||
'result_summary': f'{len(search_result["results"])} results',
|
||||
'conversation_id': conversation_id,
|
||||
'agent_id': current_agent_id,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
except Exception as e:
|
||||
duration_ms = int((time.time() - start_time) * 1000)
|
||||
logger.error(f"搜索失败: {e}")
|
||||
tool_service.increment_stats(search_tool.id, False)
|
||||
tool_service.log_usage({
|
||||
'tool_id': search_tool.id,
|
||||
'tool_type': 'search',
|
||||
'query': message,
|
||||
'success': False,
|
||||
'error_message': str(e),
|
||||
'conversation_id': conversation_id,
|
||||
'duration_ms': duration_ms
|
||||
})
|
||||
|
||||
# 5. 保存用户消息到数据库
|
||||
user_msg = conv_service.add_message(
|
||||
conversation_id=conversation.id,
|
||||
role='user',
|
||||
@@ -892,22 +894,16 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
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()
|
||||
}
|
||||
})
|
||||
# 6. 获取对话历史（包含刚保存的用户消息）
|
||||
history = conv_service.get_conversation_history(conversation_id, limit=agent_config['agent'].get('max_history', 20))
|
||||
|
||||
# 获取Agent配置
|
||||
agent_config = agent_service.get_agent_config(current_agent_id)
|
||||
# 7. 如果有搜索结果，添加到消息中
|
||||
if search_context:
|
||||
modified_system_prompt = agent_config['agent'].get('system_prompt', '') + "\n\n如果提供了搜索结果，请基于搜索结果回答用户问题，并注明信息来源。"
|
||||
agent_config['agent']['system_prompt'] = modified_system_prompt
|
||||
history.append({"role": "system", "content": f"以下是搜索到的相关信息，请参考这些内容回答用户问题：{search_context}"})
|
||||
|
||||
# 8. 调用LLM返回回复
|
||||
if not agent_config or not agent_config.get('provider'):
|
||||
await websocket.send_json({
|
||||
"type": "error",
|
||||
@@ -915,20 +911,7 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
})
|
||||
continue
|
||||
|
||||
# 获取对话历史
|
||||
history = conv_service.get_conversation_history(conversation_id, limit=agent_config['agent'].get('max_history', 20))
|
||||
|
||||
# 如果有搜索结果，添加到消息中
|
||||
if search_context:
|
||||
# 在系统提示中添加搜索结果说明
|
||||
modified_system_prompt = agent_config['agent'].get('system_prompt', '') + "\n\n如果提供了搜索结果，请基于搜索结果回答用户问题，并注明信息来源。"
|
||||
agent_config['agent']['system_prompt'] = modified_system_prompt
|
||||
# 将搜索结果作为系统消息添加到历史
|
||||
history.append({"role": "system", "content": f"以下是搜索到的相关信息，请参考这些内容回答用户问题：{search_context}"})
|
||||
|
||||
# 使用非流式调用LLM（简化版本，确保稳定）
|
||||
try:
|
||||
# 调用LLM（非流式）
|
||||
response, thinking_content = await llm_service.chat(
|
||||
messages=history,
|
||||
provider_config=agent_config['provider'],
|
||||
@@ -949,7 +932,7 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
model_used=agent_config['provider'].get('default_model')
|
||||
)
|
||||
|
||||
# 发送完整回复（包含思考内容）
|
||||
# 发送AI回复
|
||||
await websocket.send_json({
|
||||
"type": "assistant_message",
|
||||
"conversation_id": conversation_id,
|
||||
@@ -967,7 +950,6 @@ async def websocket_endpoint(websocket: WebSocket, user_id: str):
|
||||
|
||||
logger.info(f"AI回复已发送: conversation_id={conversation_id}")
|
||||
|
||||
# 启用发送按钮
|
||||
await websocket.send_json({
|
||||
"type": "stream_end",
|
||||
"conversation_id": conversation_id
|
||||
|
||||
@@ -280,9 +280,11 @@
|
||||
case 'stream_end': document.getElementById('sendBtn').disabled = false; break;
|
||||
case 'user_message':
|
||||
lastUserMessage = data.message.content; // 存储最后一条用户消息
|
||||
if (!isRegenerating) {
|
||||
// 如果是刚发送的消息，已经显示了，不再重复显示
|
||||
if (!isRegenerating && data.message.content !== lastSentMessage) {
|
||||
appendMessage('user', data.message.content);
|
||||
}
|
||||
lastSentMessage = null; // 清除标记
|
||||
// 注意：不要在这里重置 isRegenerating，要等 assistant_message 处理后再重置
|
||||
break;
|
||||
case 'assistant_message':
|
||||
@@ -756,6 +758,10 @@
|
||||
input.value = '';
|
||||
input.style.height = 'auto';
|
||||
|
||||
// 立即显示用户消息（不等后端广播）
|
||||
lastSentMessage = msg; // 记录最后发送的消息，避免重复显示
|
||||
appendMessage('user', msg);
|
||||
|
||||
// 获取工具禁用状态
|
||||
const enableSearch = document.getElementById('enableSearch').checked;
|
||||
const disabledTools = [];
|
||||
@@ -772,6 +778,8 @@
|
||||
}
|
||||
}
|
||||
|
||||
let lastSentMessage = null; // 记录最后发送的消息
|
||||
|
||||
function setupTextarea() {
|
||||
const textarea = document.getElementById('messageInput');
|
||||
textarea.addEventListener('keydown', e => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendMessage(); } });
|
||||
|
||||
Reference in New Issue
Block a user