This commit is contained in:
@@ -1,198 +1,88 @@
|
||||
"""
|
||||
AI Agent 服务类 - 完全简化版本！
|
||||
按照指南实现，不用 stream_mode="messages" 避免重复 token！
|
||||
AI Agent 服务类
|
||||
"""
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
from typing import AsyncGenerator, Dict, Any, Optional, Tuple
|
||||
from typing import AsyncGenerator, Dict, Any
|
||||
|
||||
# LangGraph 序列化器（修复 checkpoint 反序列化警告）
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
# 本地模块
|
||||
from backend.app.model_services import get_cached_chat_services
|
||||
from backend.app.main_graph.main_graph_builder import build_agent_graph
|
||||
from backend.app.logger import debug, info, warning, error
|
||||
from backend.app.main_graph.state import AgentState
|
||||
from .stream_context import set_stream_queue
|
||||
from backend.app.logger import info
|
||||
from backend.app.memory.mem0_client import Mem0Client
|
||||
|
||||
from .service_config import ServiceConfig
|
||||
from .stream_handler import run_graph_stream
|
||||
|
||||
|
||||
class AIAgentService:
|
||||
def __init__(self, checkpointer):
|
||||
self.checkpointer = checkpointer
|
||||
self.graph = None
|
||||
self.chat_services = None
|
||||
# Mem0 客户端
|
||||
self.config: ServiceConfig = None
|
||||
self.mem0_client = None
|
||||
|
||||
async def initialize(self):
|
||||
# 0. 初始化 Mem0 客户端
|
||||
from ..memory.mem0_client import Mem0Client
|
||||
async def initialize(self) -> "AIAgentService":
|
||||
"""初始化 Agent 服务"""
|
||||
self.mem0_client = Mem0Client()
|
||||
|
||||
# 1. 获取缓存的模型字典
|
||||
|
||||
self.chat_services = get_cached_chat_services()
|
||||
info(f"✅ 加载了 {len(self.chat_services)} 个可用模型: {list(self.chat_services.keys())}")
|
||||
|
||||
# 2. 构建图
|
||||
info(f"🔄 构建 Agent 图...")
|
||||
|
||||
graph_builder = build_agent_graph(
|
||||
chat_services=self.chat_services,
|
||||
mem0_client=self.mem0_client
|
||||
)
|
||||
|
||||
# 编译图
|
||||
self.graph = graph_builder.compile(checkpointer=self.checkpointer)
|
||||
|
||||
self.config = ServiceConfig(self.chat_services)
|
||||
info(f"✅ Agent 图初始化完成")
|
||||
|
||||
|
||||
return self
|
||||
|
||||
def _resolve_model(self, model: str) -> str:
|
||||
"""
|
||||
解析并验证模型名称，不可用时回退到第一个可用模型
|
||||
|
||||
Args:
|
||||
model: 目标模型名称
|
||||
|
||||
Returns:
|
||||
实际使用的模型名称
|
||||
"""
|
||||
if not model or model not in self.chat_services:
|
||||
fallback = next(iter(self.chat_services.keys()))
|
||||
warning(f"模型 '{model}' 不可用，回退到 '{fallback}'")
|
||||
return fallback
|
||||
return model
|
||||
|
||||
def _build_invocation(
|
||||
def _resolve_and_build(
|
||||
self, message: str, thread_id: str, model: str, user_id: str
|
||||
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
|
||||
"""
|
||||
构建图调用所需的 config 和 input_state
|
||||
|
||||
Args:
|
||||
message: 用户消息
|
||||
thread_id: 会话 ID
|
||||
model: 模型名称
|
||||
user_id: 用户 ID
|
||||
|
||||
Returns:
|
||||
(config, input_state) 元组
|
||||
"""
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
},
|
||||
"metadata": {"user_id": user_id}
|
||||
}
|
||||
|
||||
input_state = {
|
||||
"messages": [HumanMessage(content=message)],
|
||||
"user_id": user_id,
|
||||
}
|
||||
return config, input_state
|
||||
):
|
||||
"""解析模型并构建调用参数"""
|
||||
resolved_model = self.config.resolve_model(model)
|
||||
return resolved_model, self.config.build_invocation(
|
||||
message, thread_id, resolved_model, user_id
|
||||
)
|
||||
|
||||
async def process_message(
|
||||
self, message: str, thread_id: str, model: str = "", user_id: str = "default_user"
|
||||
) -> dict:
|
||||
"""处理用户消息，返回包含回复、token统计和耗时的字典"""
|
||||
# 解析模型名称
|
||||
resolved_model = self._resolve_model(model)
|
||||
|
||||
# 构建调用参数
|
||||
config, input_state = self._build_invocation(message, thread_id, resolved_model, user_id)
|
||||
resolved_model, (config, input_state) = self._resolve_and_build(
|
||||
message, thread_id, model, user_id
|
||||
)
|
||||
|
||||
result = await self.graph.ainvoke(input_state, config=config)
|
||||
|
||||
# 优先使用 final_reply（finalize 节点返回）
|
||||
reply = result.get("final_reply", "")
|
||||
if not reply and result.get("messages"):
|
||||
reply = result["messages"][-1].content
|
||||
|
||||
token_usage = result.get("last_token_usage", {})
|
||||
elapsed_time = result.get("last_elapsed_time", 0.0)
|
||||
|
||||
# 获取元数据
|
||||
metadata = result.get("metadata", {})
|
||||
|
||||
return {
|
||||
"reply": reply,
|
||||
"token_usage": token_usage,
|
||||
"elapsed_time": elapsed_time,
|
||||
"token_usage": result.get("last_token_usage", {}),
|
||||
"elapsed_time": result.get("last_elapsed_time", 0.0),
|
||||
"model_used": resolved_model,
|
||||
"metadata": metadata
|
||||
"metadata": result.get("metadata", {}),
|
||||
}
|
||||
|
||||
async def process_message_stream(
|
||||
self, message: str, thread_id: str, model: str = "", user_id: str = "default_user"
|
||||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||||
"""流式处理消息 - 完全简化！"""
|
||||
# 解析模型名称
|
||||
resolved_model = self._resolve_model(model)
|
||||
|
||||
# 构建调用参数
|
||||
config, input_state = self._build_invocation(message, thread_id, resolved_model, user_id)
|
||||
"""流式处理消息"""
|
||||
resolved_model, (config, input_state) = self._resolve_and_build(
|
||||
message, thread_id, model, user_id
|
||||
)
|
||||
|
||||
info(f"🚀 开始执行 Agent 图，指定模型: {resolved_model}")
|
||||
actual_model_used = resolved_model
|
||||
|
||||
# 创建 token 队列
|
||||
queue = asyncio.Queue()
|
||||
set_stream_queue(queue) # 设置上下文变量
|
||||
|
||||
async def run_graph():
|
||||
"""后台任务：运行 graph，流式事件都从 agent 节点内部发送！"""
|
||||
try:
|
||||
info(f"📡 开始调用 graph.astream()...")
|
||||
|
||||
# 注意：只用 stream_mode=["updates"]，不要 "messages"！避免重复 token！
|
||||
async for _ in self.graph.astream(
|
||||
input_state,
|
||||
config=config,
|
||||
stream_mode=["updates"],
|
||||
version="v2",
|
||||
subgraphs=True
|
||||
):
|
||||
# 流式事件都从 agent.py 节点内部通过队列发送了
|
||||
# 这里不需要再发送任何事件
|
||||
pass
|
||||
except Exception as e:
|
||||
error(f"❌ 执行图时出错: {e}")
|
||||
import traceback
|
||||
error(f"📋 堆栈: {traceback.format_exc()}")
|
||||
await queue.put({"type": "error", "message": str(e)})
|
||||
finally:
|
||||
await queue.put(None) # 结束哨兵
|
||||
|
||||
# 启动后台任务
|
||||
bg_task = asyncio.create_task(run_graph())
|
||||
|
||||
try:
|
||||
while True:
|
||||
event = await queue.get()
|
||||
if event is None:
|
||||
break
|
||||
async for event in run_graph_stream(self.graph, input_state, config):
|
||||
if event.get("type") != "done":
|
||||
yield event
|
||||
|
||||
except GeneratorExit:
|
||||
# 客户端断开连接，取消后台任务
|
||||
info("⚠️ GeneratorExit，取消后台任务")
|
||||
bg_task.cancel()
|
||||
raise
|
||||
finally:
|
||||
# 保证任务被清理
|
||||
if not bg_task.done():
|
||||
info("⏹️ 清理后台任务")
|
||||
bg_task.cancel()
|
||||
try:
|
||||
await bg_task
|
||||
except asyncio.CancelledError:
|
||||
info("✅ 后台任务已取消")
|
||||
|
||||
# 发送结束事件，保证前端平稳关闭
|
||||
yield {
|
||||
"type": "done",
|
||||
"model_used": actual_model_used
|
||||
}
|
||||
else:
|
||||
yield {**event, "model_used": resolved_model}
|
||||
|
||||
Reference in New Issue
Block a user