325 lines
13 KiB
Python
325 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Annotated, Any
|
|
|
|
from deepagents import SubagentTransformer
|
|
from deepagents.middleware._utils import append_to_system_message
|
|
from langchain.agents.middleware.types import AgentMiddleware, ContextT, ModelRequest, ModelResponse, ResponseT
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
from langchain_core.tools import StructuredTool
|
|
from langgraph.prebuilt.tool_node import ToolRuntime
|
|
from langgraph.types import Command
|
|
|
|
from yuxi.agents.subagent_thread import make_child_thread_id
|
|
from yuxi.repositories.agent_repository import SUB_AGENT_BACKEND_ID, AgentRepository
|
|
from yuxi.repositories.user_repository import UserRepository
|
|
from yuxi.storage.postgres.manager import pg_manager
|
|
from yuxi.storage.postgres.models_business import Agent
|
|
from yuxi.utils.datetime_utils import utc_isoformat
|
|
|
|
_EXCLUDED_STATE_KEYS = {
|
|
"messages",
|
|
"todos",
|
|
"structured_response",
|
|
"skills_metadata",
|
|
"skills_load_errors",
|
|
"activated_skills",
|
|
"memory_contents",
|
|
}
|
|
|
|
TASK_SYSTEM_PROMPT = """## `task`（子智能体任务工具）
|
|
|
|
你可以使用 `task` 工具把复杂、独立的子任务交给已配置的子智能体处理。子智能体只返回最终结果，你看不到它的中间步骤。
|
|
|
|
使用原则：
|
|
- 任务足够复杂、可以独立完成、或需要隔离上下文时使用。
|
|
- 多个互不依赖的子任务可以并行调用多个 `task`。
|
|
- 简单问题或少量直接工具调用不要委派。
|
|
- 调用时必须选择下方可用的 `subagent_type`，并在 `description` 中写清目标、上下文和期望输出。
|
|
- 不要通过 shell、curl、HTTP API 或命令行间接调用子智能体；需要子智能体时必须使用 `task` 工具。
|
|
|
|
Available subagent types:
|
|
|
|
{available_agents}"""
|
|
|
|
TASK_TOOL_DESCRIPTION = """Launch a configured Yuxi subagent to handle an isolated task.
|
|
|
|
Available subagent types:
|
|
{available_agents}
|
|
|
|
Use `subagent_type` to select one available subagent and put the full task brief in `description`.
|
|
Do not call subagents through shell, curl, HTTP APIs, or command-line indirection."""
|
|
|
|
|
|
TASK_DESCRIPTION_ARG = "需要子智能体独立完成的任务描述，包含必要上下文和期望输出。"
|
|
SUBAGENT_TYPE_ARG = "要调用的子智能体标识，必须是工具描述中列出的可用类型之一。"
|
|
|
|
|
|
def _get_agent_backend(backend_id: str):
|
|
from yuxi.agents.buildin import agent_manager
|
|
|
|
return agent_manager.get_agent(backend_id)
|
|
|
|
|
|
def _final_assistant_text(messages: list[Any]) -> str:
|
|
for message in reversed(messages):
|
|
if isinstance(message, AIMessage):
|
|
text = message.text.rstrip() if message.text else ""
|
|
if text:
|
|
return text
|
|
return "子智能体已完成任务，但没有返回文本结果。"
|
|
|
|
|
|
def _result_artifacts(result: dict[str, Any]) -> list[str]:
|
|
artifacts = result.get("artifacts")
|
|
return list(artifacts) if isinstance(artifacts, list) else []
|
|
|
|
|
|
def _preview_text(text: str, limit: int = 500) -> str:
|
|
return text if len(text) <= limit else f"{text[:limit]}..."
|
|
|
|
|
|
def _completed_tool_response(result: dict[str, Any], tool_call_id: str, subagent_run: dict[str, Any]) -> Command:
|
|
final_text = _final_assistant_text(result.get("messages") or [])
|
|
artifacts = _result_artifacts(result)
|
|
subagent_run = {
|
|
**subagent_run,
|
|
"status": "completed",
|
|
"completed_at": utc_isoformat(),
|
|
"result_preview": _preview_text(final_text),
|
|
"error": None,
|
|
"artifacts": artifacts,
|
|
}
|
|
update: dict[str, Any] = {"messages": [ToolMessage(final_text, tool_call_id=tool_call_id)]}
|
|
if artifacts:
|
|
update["artifacts"] = artifacts
|
|
update["subagent_runs"] = [subagent_run]
|
|
return Command(update=update)
|
|
|
|
|
|
def _failed_tool_response(error: Exception, tool_call_id: str, subagent_run: dict[str, Any]) -> Command:
|
|
error_text = str(error)
|
|
message = f"子智能体 {subagent_run['subagent_type']} 调用失败：{error_text}"
|
|
update = {
|
|
"messages": [ToolMessage(message, tool_call_id=tool_call_id)],
|
|
"subagent_runs": [
|
|
{
|
|
**subagent_run,
|
|
"status": "failed",
|
|
"completed_at": utc_isoformat(),
|
|
"result_preview": _preview_text(message),
|
|
"error": error_text,
|
|
"artifacts": [],
|
|
}
|
|
],
|
|
}
|
|
return Command(update=update)
|
|
|
|
|
|
def _state_for_child(
|
|
description: str,
|
|
runtime: ToolRuntime,
|
|
*,
|
|
parent_thread_id: str,
|
|
file_thread_id: str,
|
|
skills_thread_id: str,
|
|
) -> dict[str, Any]:
|
|
state = {key: value for key, value in runtime.state.items() if key not in _EXCLUDED_STATE_KEYS}
|
|
state.update(
|
|
{
|
|
"parent_thread_id": parent_thread_id,
|
|
"file_thread_id": file_thread_id,
|
|
"skills_thread_id": skills_thread_id,
|
|
}
|
|
)
|
|
state["messages"] = [HumanMessage(content=description)]
|
|
return state
|
|
|
|
|
|
def _child_config(
|
|
runtime: ToolRuntime,
|
|
*,
|
|
child_thread_id: str,
|
|
uid: str,
|
|
parent_thread_id: str,
|
|
file_thread_id: str,
|
|
skills_thread_id: str,
|
|
) -> dict:
|
|
parent_config = runtime.config or {}
|
|
config: dict[str, Any] = {}
|
|
if "callbacks" in parent_config:
|
|
config["callbacks"] = parent_config["callbacks"]
|
|
if "tags" in parent_config:
|
|
config["tags"] = parent_config["tags"]
|
|
parent_configurable = (
|
|
parent_config.get("configurable") if isinstance(parent_config.get("configurable"), dict) else {}
|
|
)
|
|
parent_configurable = {
|
|
key: value
|
|
for key, value in parent_configurable.items()
|
|
if not str(key).startswith(("checkpoint_", "__pregel_"))
|
|
}
|
|
config["configurable"] = {
|
|
**parent_configurable,
|
|
"thread_id": child_thread_id,
|
|
"uid": uid,
|
|
"parent_thread_id": parent_thread_id,
|
|
"file_thread_id": file_thread_id,
|
|
"skills_thread_id": skills_thread_id,
|
|
"ls_agent_type": "subagent",
|
|
}
|
|
config["recursion_limit"] = parent_config.get("recursion_limit", 300)
|
|
return config
|
|
|
|
|
|
class YuxiSubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
|
|
def __init__(self, *, parent_context, subagents: list[Agent]) -> None:
|
|
super().__init__()
|
|
self.parent_context = parent_context
|
|
self.subagents = {agent.slug: agent for agent in subagents}
|
|
available_agents = "\n".join(f"- {agent.slug}: {agent.description or agent.name}" for agent in subagents)
|
|
self.system_prompt = TASK_SYSTEM_PROMPT.format(available_agents=available_agents)
|
|
self.tools = [self._build_task_tool(available_agents)]
|
|
self.subagent_names = frozenset(self.subagents)
|
|
self.transformers = [
|
|
lambda scope: SubagentTransformer(scope, subagent_names=self.subagent_names)
|
|
]
|
|
|
|
def _build_task_tool(self, available_agents: str) -> StructuredTool:
|
|
def task(
|
|
description: Annotated[str, TASK_DESCRIPTION_ARG],
|
|
subagent_type: Annotated[str, SUBAGENT_TYPE_ARG],
|
|
runtime: ToolRuntime,
|
|
) -> str:
|
|
return "task 工具仅支持异步调用"
|
|
|
|
async def atask(
|
|
description: Annotated[str, TASK_DESCRIPTION_ARG],
|
|
subagent_type: Annotated[str, SUBAGENT_TYPE_ARG],
|
|
runtime: ToolRuntime,
|
|
) -> str | Command:
|
|
if subagent_type not in self.subagents:
|
|
allowed = ", ".join(f"`{slug}`" for slug in self.subagents)
|
|
return f"无法调用子智能体 {subagent_type}，可用子智能体只有：{allowed}"
|
|
if not runtime.tool_call_id:
|
|
raise ValueError("Tool call ID is required for subagent invocation")
|
|
|
|
parent_thread_id = str(
|
|
getattr(self.parent_context, "parent_thread_id", None) or self.parent_context.thread_id
|
|
)
|
|
file_thread_id = str(getattr(self.parent_context, "file_thread_id", None) or parent_thread_id)
|
|
uid = str(getattr(self.parent_context, "uid", "") or "").strip()
|
|
if not uid:
|
|
return "无法调用子智能体：当前运行时缺少 uid"
|
|
|
|
agent = self.subagents[subagent_type]
|
|
backend = _get_agent_backend(agent.backend_id)
|
|
if not backend or agent.backend_id != SUB_AGENT_BACKEND_ID:
|
|
return f"无法调用子智能体 {subagent_type}：后端配置无效"
|
|
|
|
child_thread_id = make_child_thread_id(parent_thread_id, agent.slug, runtime.tool_call_id)
|
|
subagent_run = {
|
|
"id": runtime.tool_call_id,
|
|
"subagent_type": subagent_type,
|
|
"subagent_name": agent.name,
|
|
"child_thread_id": child_thread_id,
|
|
"description": description,
|
|
"created_at": utc_isoformat(),
|
|
}
|
|
child_context = backend.context_schema()
|
|
config_context = (agent.config_json or {}).get("context") if isinstance(agent.config_json, dict) else None
|
|
if isinstance(config_context, dict):
|
|
child_context.update_from_dict(config_context)
|
|
child_context.uid = uid
|
|
child_context.thread_id = child_thread_id
|
|
child_context.parent_thread_id = parent_thread_id
|
|
child_context.file_thread_id = file_thread_id
|
|
child_context.skills_thread_id = child_thread_id
|
|
child_context.is_subagent_runtime = True
|
|
|
|
try:
|
|
graph = await backend.get_graph(context=child_context)
|
|
result = await graph.ainvoke(
|
|
_state_for_child(
|
|
description,
|
|
runtime,
|
|
parent_thread_id=parent_thread_id,
|
|
file_thread_id=file_thread_id,
|
|
skills_thread_id=child_thread_id,
|
|
),
|
|
config=_child_config(
|
|
runtime,
|
|
child_thread_id=child_thread_id,
|
|
uid=uid,
|
|
parent_thread_id=parent_thread_id,
|
|
file_thread_id=file_thread_id,
|
|
skills_thread_id=child_thread_id,
|
|
),
|
|
context=child_context,
|
|
)
|
|
except Exception as exc:
|
|
return _failed_tool_response(exc, runtime.tool_call_id, subagent_run)
|
|
return _completed_tool_response(result, runtime.tool_call_id, subagent_run)
|
|
|
|
return StructuredTool.from_function(
|
|
name="task",
|
|
func=task,
|
|
coroutine=atask,
|
|
description=TASK_TOOL_DESCRIPTION.format(available_agents=available_agents),
|
|
infer_schema=True,
|
|
)
|
|
|
|
def wrap_model_call(
|
|
self,
|
|
request: ModelRequest[ContextT],
|
|
handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],
|
|
) -> ModelResponse[ResponseT]:
|
|
return handler(
|
|
request.override(system_message=append_to_system_message(request.system_message, self.system_prompt))
|
|
)
|
|
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest[ContextT],
|
|
handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],
|
|
) -> ModelResponse[ResponseT]:
|
|
return await handler(
|
|
request.override(system_message=append_to_system_message(request.system_message, self.system_prompt))
|
|
)
|
|
|
|
|
|
async def create_subagent_task_middleware(parent_context) -> YuxiSubAgentMiddleware | None:
|
|
selected_slugs = [
|
|
str(slug).strip() for slug in (getattr(parent_context, "subagents", None) or []) if str(slug).strip()
|
|
]
|
|
uid = str(getattr(parent_context, "uid", "") or "").strip()
|
|
if not uid:
|
|
return None
|
|
|
|
async with pg_manager.get_async_session_context() as db:
|
|
user = await UserRepository().get_by_uid_with_db(db, uid)
|
|
if user is None:
|
|
return None
|
|
repo = AgentRepository(db)
|
|
if selected_slugs:
|
|
subagents: list[Agent] = []
|
|
seen: set[str] = set()
|
|
for slug in selected_slugs:
|
|
if slug in seen:
|
|
continue
|
|
seen.add(slug)
|
|
agent = await repo.get_visible_subagent_by_slug(slug=slug, user=user)
|
|
if agent and agent.backend_id == SUB_AGENT_BACKEND_ID:
|
|
subagents.append(agent)
|
|
else:
|
|
subagents = [
|
|
agent
|
|
for agent in await repo.list_visible_subagents(user=user)
|
|
if agent.backend_id == SUB_AGENT_BACKEND_ID
|
|
]
|
|
|
|
if not subagents:
|
|
return None
|
|
return YuxiSubAgentMiddleware(parent_context=parent_context, subagents=subagents)
|