425 lines
17 KiB
Python
425 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import os
|
||
from abc import abstractmethod
|
||
from contextlib import suppress
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from langchain_core.messages import ToolMessage
|
||
from langgraph.checkpoint.memory import InMemorySaver
|
||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
|
||
from langgraph.graph.state import CompiledStateGraph
|
||
from langgraph.types import Command
|
||
|
||
from yuxi import config as sys_config
|
||
from yuxi.agents.context import BaseContext, resolve_agent_resource_options
|
||
from yuxi.storage.postgres.manager import pg_manager
|
||
from yuxi.utils import logger
|
||
from yuxi.utils.subagent_thread_utils import make_child_thread_id
|
||
|
||
|
||
def _json_safe(value: Any) -> Any:
|
||
if value is None or isinstance(value, str | int | float | bool):
|
||
return value
|
||
if isinstance(value, dict):
|
||
return {str(key): _json_safe(child) for key, child in value.items()}
|
||
if isinstance(value, list | tuple):
|
||
return [_json_safe(child) for child in value]
|
||
if hasattr(value, "model_dump"):
|
||
return _json_safe(value.model_dump())
|
||
return str(value)
|
||
|
||
|
||
def _normalize_tool_event_data(data: Any) -> Any:
|
||
"""规整 tools 流事件：write_todos / task 等返回 Command 的工具，其 tool-finished
|
||
output 是 Command 对象，_json_safe 只能退化成 repr 字符串，前端无法关联结果。
|
||
这里从 Command.update["messages"] 取出真正的 ToolMessage，使其与普通工具一致。"""
|
||
if not isinstance(data, dict) or data.get("event") != "tool-finished":
|
||
return data
|
||
output = data.get("output")
|
||
if not isinstance(output, Command):
|
||
return data
|
||
update = output.update if isinstance(output.update, dict) else {}
|
||
messages = update.get("messages")
|
||
if not isinstance(messages, list):
|
||
return data
|
||
tool_call_id = data.get("tool_call_id")
|
||
tool_message = next(
|
||
(m for m in messages if isinstance(m, ToolMessage) and m.tool_call_id == tool_call_id),
|
||
next((m for m in messages if isinstance(m, ToolMessage)), None),
|
||
)
|
||
if tool_message is None:
|
||
return data
|
||
return {**data, "output": tool_message}
|
||
|
||
|
||
def _metadata_thread_id(value: Any) -> str | None:
|
||
if not isinstance(value, dict):
|
||
return None
|
||
for key in ("thread_id", "subagent_thread_id"):
|
||
thread_id = value.get(key)
|
||
if isinstance(thread_id, str) and thread_id.strip():
|
||
return thread_id.strip()
|
||
for key in ("metadata", "configurable", "config"):
|
||
thread_id = _metadata_thread_id(value.get(key))
|
||
if thread_id:
|
||
return thread_id
|
||
return None
|
||
|
||
|
||
def _subagent_route_for_namespace(
|
||
routes: dict[tuple[str, ...], dict[str, str]], namespace: list[str]
|
||
) -> dict[str, str] | None:
|
||
ns = tuple(namespace)
|
||
for path, route in sorted(routes.items(), key=lambda item: len(item[0]), reverse=True):
|
||
if ns[: len(path)] == path:
|
||
return route
|
||
return None
|
||
|
||
|
||
async def _collect_subagent_routes(run, parent_thread_id: str, routes: dict[tuple[str, ...], dict[str, str]]) -> None:
|
||
subagents = getattr(run, "subagents", None)
|
||
if subagents is None:
|
||
return
|
||
|
||
try:
|
||
async for subagent in subagents:
|
||
path = tuple(getattr(subagent, "path", ()) or ())
|
||
subagent_type = getattr(subagent, "name", None) or getattr(subagent, "graph_name", None)
|
||
cause = getattr(subagent, "cause", None)
|
||
tool_call_id = (
|
||
cause.get("tool_call_id") if isinstance(cause, dict) else getattr(subagent, "trigger_call_id", None)
|
||
)
|
||
state = getattr(subagent, "state", None)
|
||
metadata = getattr(subagent, "metadata", None)
|
||
thread_id = _metadata_thread_id(metadata) or _metadata_thread_id(state)
|
||
if not thread_id and isinstance(subagent_type, str) and isinstance(tool_call_id, str) and tool_call_id:
|
||
thread_id = make_child_thread_id(parent_thread_id, subagent_type, tool_call_id)
|
||
if path and isinstance(subagent_type, str) and isinstance(tool_call_id, str) and tool_call_id and thread_id:
|
||
routes[path] = {
|
||
"thread_id": thread_id,
|
||
"parent_thread_id": parent_thread_id,
|
||
"subagent_type": subagent_type,
|
||
"tool_call_id": tool_call_id,
|
||
}
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as exc:
|
||
logger.debug(f"collect subagent stream routes failed: {exc}")
|
||
|
||
|
||
class BaseAgent:
|
||
"""
|
||
定义一个基础 Agent 供 各类 graph 继承
|
||
"""
|
||
|
||
name = "base_agent"
|
||
description = "base_agent"
|
||
capabilities: list[str] = [] # 智能体能力列表，如 ["file_upload", "web_search"] 等
|
||
context_schema: type[BaseContext] = BaseContext # 智能体上下文 schema
|
||
|
||
def __init__(self, **kwargs):
|
||
self.graph = None # will be covered by get_graph
|
||
self.checkpointer = None
|
||
self._async_conn = None
|
||
self.workdir = Path(sys_config.save_dir) / "agents" / self.module_name
|
||
self.workdir.mkdir(parents=True, exist_ok=True)
|
||
|
||
@property
|
||
def module_name(self) -> str:
|
||
"""Get the module name of the agent class."""
|
||
return self.__class__.__module__.split(".")[-2]
|
||
|
||
@property
|
||
def id(self) -> str:
|
||
"""Get the agent's class name."""
|
||
return self.__class__.__name__
|
||
|
||
async def get_info(
|
||
self,
|
||
include_configurable_items: bool = True,
|
||
user_role: str | None = None,
|
||
db=None,
|
||
user=None,
|
||
):
|
||
# metadata 固定在代码中，由各 Agent 的类属性提供
|
||
metadata = self.load_metadata()
|
||
configurable_items = {}
|
||
if include_configurable_items:
|
||
configurable_items = self.context_schema.get_configurable_items(user_role=user_role)
|
||
if db is not None and user is not None:
|
||
resource_fields = {
|
||
item["kind"]
|
||
for item in configurable_items.values()
|
||
if item.get("kind") in {"tools", "knowledges", "mcps", "skills", "subagents"}
|
||
}
|
||
resource_options = await resolve_agent_resource_options(resource_fields, db=db, user=user)
|
||
for item in configurable_items.values():
|
||
if item.get("kind") in resource_options:
|
||
item["options"] = resource_options[item["kind"]]
|
||
|
||
# Merge metadata with class attributes, metadata takes precedence
|
||
return {
|
||
"id": self.id,
|
||
"name": getattr(self, "name", "Unknown"),
|
||
"description": getattr(self, "description", "Unknown"),
|
||
"metadata": metadata,
|
||
"configurable_items": configurable_items,
|
||
"capabilities": getattr(self, "capabilities", []), # 智能体能力列表
|
||
}
|
||
|
||
async def get_config(self):
|
||
return self.context_schema()
|
||
|
||
async def stream_values(self, messages: list[str], input_context=None, **kwargs):
|
||
context = self.context_schema()
|
||
context.update_from_dict(input_context or {})
|
||
graph = await self.get_graph(context=context)
|
||
for event in graph.astream({"messages": messages}, stream_mode="values", context=context):
|
||
yield event["messages"]
|
||
|
||
async def stream_messages(self, messages: list[str], input_context=None, **kwargs):
|
||
context = self.context_schema()
|
||
context.update_from_dict(input_context or {})
|
||
graph = await self.get_graph(context=context)
|
||
logger.debug(f"stream_messages: {context=}")
|
||
|
||
# 构建配置：LangGraph 会自动从 checkpointer 恢复 state
|
||
input_config = {
|
||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||
"recursion_limit": 300,
|
||
}
|
||
|
||
# langfuse metadata and callbacks integration
|
||
if callbacks := kwargs.get("callbacks"):
|
||
input_config["callbacks"] = list(callbacks)
|
||
if metadata := kwargs.get("metadata"):
|
||
input_config["metadata"] = dict(metadata)
|
||
if tags := kwargs.get("tags"):
|
||
input_config["tags"] = list(tags)
|
||
|
||
async for msg, metadata in graph.astream(
|
||
{"messages": messages},
|
||
stream_mode="messages",
|
||
context=context,
|
||
config=input_config,
|
||
):
|
||
yield msg, metadata
|
||
|
||
async def _stream_input_with_state(self, graph_input, input_context=None, **kwargs):
|
||
context = self.context_schema()
|
||
context.update_from_dict(input_context or {})
|
||
graph = await self.get_graph(context=context)
|
||
logger.debug(f"stream_with_state: {context=}")
|
||
|
||
input_config = {
|
||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||
"recursion_limit": 300,
|
||
}
|
||
|
||
if callbacks := kwargs.get("callbacks"):
|
||
input_config["callbacks"] = list(callbacks)
|
||
if metadata := kwargs.get("metadata"):
|
||
input_config["metadata"] = dict(metadata)
|
||
if tags := kwargs.get("tags"):
|
||
input_config["tags"] = list(tags)
|
||
|
||
run = await graph.astream_events(
|
||
graph_input,
|
||
context=context,
|
||
config=input_config,
|
||
version="v3",
|
||
)
|
||
subagent_routes: dict[tuple[str, ...], dict[str, str]] = {}
|
||
route_task = asyncio.create_task(_collect_subagent_routes(run, context.thread_id, subagent_routes))
|
||
try:
|
||
async for event in run:
|
||
params = event.get("params") or {}
|
||
namespace = list(params.get("namespace") or [])
|
||
method = event.get("method")
|
||
data = params.get("data")
|
||
subagent_route = _subagent_route_for_namespace(subagent_routes, namespace)
|
||
|
||
if method == "messages":
|
||
msg, metadata = data
|
||
metadata = dict(metadata or {})
|
||
actual_thread_id = (
|
||
_metadata_thread_id(metadata) or _metadata_thread_id(params) or _metadata_thread_id(data)
|
||
)
|
||
metadata["namespace"] = namespace
|
||
metadata["stream_event"] = {"method": method, "namespace": namespace}
|
||
if subagent_route:
|
||
metadata.update(subagent_route)
|
||
if actual_thread_id:
|
||
metadata["thread_id"] = actual_thread_id
|
||
yield "messages", (msg, metadata)
|
||
elif method == "values" and not namespace:
|
||
yield "values", data
|
||
elif method in {"tasks", "tools", "lifecycle"}:
|
||
if method == "tools":
|
||
data = _normalize_tool_event_data(data)
|
||
event_payload = {
|
||
"method": method,
|
||
"namespace": namespace,
|
||
"data": _json_safe(data),
|
||
}
|
||
actual_thread_id = _metadata_thread_id(params) or _metadata_thread_id(data)
|
||
if subagent_route:
|
||
event_payload.update(subagent_route)
|
||
if actual_thread_id:
|
||
event_payload["thread_id"] = actual_thread_id
|
||
yield "stream_event", event_payload
|
||
finally:
|
||
route_task.cancel()
|
||
with suppress(asyncio.CancelledError):
|
||
await route_task
|
||
|
||
async def stream_messages_with_state(self, messages: list[str], input_context=None, **kwargs):
|
||
async for event in self._stream_input_with_state({"messages": messages}, input_context, **kwargs):
|
||
yield event
|
||
|
||
async def stream_resume_with_state(self, resume_input, input_context=None, **kwargs):
|
||
async for event in self._stream_input_with_state(resume_input, input_context, **kwargs):
|
||
yield event
|
||
|
||
async def invoke_messages(self, messages: list[str], input_context=None, **kwargs):
|
||
context = self.context_schema()
|
||
context.update_from_dict(input_context or {})
|
||
graph = await self.get_graph(context=context)
|
||
logger.debug(f"invoke_messages: {context}")
|
||
|
||
# 构建配置
|
||
input_config = {
|
||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||
"recursion_limit": 100,
|
||
}
|
||
|
||
# langfuse metadata and callbacks integration
|
||
if callbacks := kwargs.get("callbacks"):
|
||
input_config["callbacks"] = list(callbacks)
|
||
if metadata := kwargs.get("metadata"):
|
||
input_config["metadata"] = dict(metadata)
|
||
if tags := kwargs.get("tags"):
|
||
input_config["tags"] = list(tags)
|
||
|
||
msg = await graph.ainvoke(
|
||
{"messages": messages},
|
||
context=context,
|
||
config=input_config,
|
||
)
|
||
return msg
|
||
|
||
async def check_checkpointer(self):
|
||
app = await self.get_graph()
|
||
if not hasattr(app, "checkpointer") or app.checkpointer is None:
|
||
return False
|
||
return True
|
||
|
||
async def get_history(self, uid, thread_id) -> list[dict]:
|
||
"""获取历史消息"""
|
||
try:
|
||
app = await self.get_graph()
|
||
|
||
if not await self.check_checkpointer():
|
||
return []
|
||
|
||
config = {"configurable": {"thread_id": thread_id, "uid": uid}}
|
||
state = await app.aget_state(config)
|
||
|
||
result = []
|
||
if state:
|
||
messages = state.values.get("messages", [])
|
||
for msg in messages:
|
||
if hasattr(msg, "model_dump"):
|
||
msg_dict = msg.model_dump() # 转换成字典
|
||
else:
|
||
msg_dict = dict(msg) if hasattr(msg, "__dict__") else {"content": str(msg)}
|
||
result.append(msg_dict)
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取智能体 {self.name} 历史消息出错: {e}")
|
||
return []
|
||
|
||
def reload_graph(self):
|
||
"""重置 graph 缓存，强制下次调用 get_graph 时重新构建"""
|
||
self.graph = None
|
||
logger.info(f"{self.name} graph 缓存已清空，将在下次调用时重新构建")
|
||
|
||
@abstractmethod
|
||
async def get_graph(self, **kwargs) -> CompiledStateGraph:
|
||
"""
|
||
获取并编译对话图实例。
|
||
必须确保在编译时设置 checkpointer，否则将无法获取历史记录。
|
||
例如: graph = workflow.compile(checkpointer=sqlite_checkpointer)
|
||
"""
|
||
pass
|
||
|
||
async def _get_checkpointer(self):
|
||
if self.checkpointer is not None:
|
||
return self.checkpointer
|
||
|
||
checkpointer = None
|
||
backend = os.getenv("LANGGRAPH_CHECKPOINTER_BACKEND", "sqlite").strip().lower()
|
||
|
||
if backend == "postgres":
|
||
checkpointer = await self._create_postgres_checkpointer()
|
||
|
||
if checkpointer is None:
|
||
try:
|
||
checkpointer = AsyncSqliteSaver(await self.get_async_conn())
|
||
except Exception as e:
|
||
logger.error(f"构建 sqlite checkpointer 失败: {e}, 尝试使用内存存储")
|
||
checkpointer = InMemorySaver()
|
||
|
||
self.checkpointer = checkpointer
|
||
return self.checkpointer
|
||
|
||
async def _create_postgres_checkpointer(self):
|
||
postgres_url = os.getenv("POSTGRES_URL")
|
||
if not postgres_url:
|
||
logger.warning("POSTGRES_URL 未配置，无法启用 postgres checkpointer，回退 sqlite")
|
||
return None
|
||
|
||
try:
|
||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # type: ignore
|
||
except Exception as e:
|
||
logger.warning(f"langgraph postgres checkpointer 不可用，回退 sqlite: {e}")
|
||
return None
|
||
|
||
try:
|
||
saver = AsyncPostgresSaver(pg_manager.langgraph_pool)
|
||
|
||
logger.info(f"{self.name} 使用 postgres checkpointer")
|
||
return saver
|
||
except Exception as e:
|
||
logger.warning(f"初始化 postgres checkpointer 失败，回退 sqlite: {e}")
|
||
return None
|
||
|
||
async def get_async_conn(self) -> aiosqlite.Connection:
|
||
"""获取异步数据库连接"""
|
||
if self._async_conn is not None:
|
||
return self._async_conn
|
||
|
||
conn = await aiosqlite.connect(os.path.join(self.workdir, "aio_history.db"))
|
||
# Patch: langgraph's AsyncSqliteSaver expects is_alive() method which aiosqlite may not have
|
||
if not hasattr(conn, "is_alive"):
|
||
conn.is_alive = lambda: True
|
||
self._async_conn = conn
|
||
return self._async_conn
|
||
|
||
async def get_aio_memory(self) -> AsyncSqliteSaver:
|
||
"""获取异步存储实例"""
|
||
return AsyncSqliteSaver(await self.get_async_conn())
|
||
|
||
def load_metadata(self) -> dict:
|
||
"""Load metadata from agent class attribute."""
|
||
metadata = getattr(self, "metadata", {})
|
||
if isinstance(metadata, dict):
|
||
return metadata
|
||
logger.warning(f"Agent {self.module_name} metadata is not a dict, fallback to empty metadata")
|
||
return {}
|