from __future__ import annotations import asyncio import hashlib import uuid 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.context import build_agent_input_context from yuxi.repositories.agent_repository import SUB_AGENT_BACKEND_ID, AgentRepository from yuxi.repositories.agent_run_repository import AgentRunRepository 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 from yuxi.utils.subagent_thread_utils import make_child_thread_id _CHILD_STATE_INHERIT_KEYS: frozenset[str] = frozenset() _TERMINAL_RUN_STATUSES = {"completed", "failed", "cancelled", "interrupted"} TASK_SYSTEM_PROMPT = """## `task`（子智能体任务工具） 你可以使用 `task` 工具把复杂、独立的子任务交给已配置的子智能体处理。子智能体只返回最终结果，你看不到它的中间步骤。 工具结果会包含子智能体线程 ID，后续需要继续同一个子任务时，把该 ID 作为 `thread_id` 传回 `task`。 使用原则： - 任务足够复杂、可以独立完成、或需要隔离上下文时使用。 - 多个互不依赖的子任务可以并行调用多个 `task`。 - 继续既有子智能体任务时传入之前结果中的 `thread_id`；新任务不要填写 `thread_id`。 - 不要并行调用同一个 `thread_id`，避免多个续跑请求同时写入同一子线程。 - 简单问题或少量直接工具调用不要委派。 - 调用时必须选择下方可用的 `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`. Omit `thread_id` for a new task. To continue a previous subagent task, pass the child thread ID returned by that prior task result as `thread_id`. Do not call subagents through shell, curl, HTTP APIs, or command-line indirection.""" TASK_DESCRIPTION_ARG = "需要子智能体独立完成的任务描述，包含必要上下文和期望输出。" SUBAGENT_TYPE_ARG = "要调用的子智能体标识，必须是工具描述中列出的可用类型之一。" THREAD_ID_ARG = "可选。要继续的既有子智能体线程 ID，必须来自之前 task 工具结果；新任务不要填写。" 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 _tool_result_with_thread_id(child_thread_id: str, content: str) -> str: return f"> 子智能体线程 ID: {child_thread_id}\n\n---\n\n{content}" def _new_child_thread_id( requested_thread_id: str | None, *, parent_thread_id: str, agent_slug: str, tool_call_id: str, ) -> tuple[str, bool]: requested_thread_id = str(requested_thread_id or "").strip() if requested_thread_id: return requested_thread_id, True return make_child_thread_id(parent_thread_id, agent_slug, tool_call_id), False def _subagent_request_id(parent_run_id: str, child_thread_id: str, tool_call_id: str, agent_slug: str) -> str: digest = hashlib.sha256(f"{parent_run_id}:{child_thread_id}:{tool_call_id}:{agent_slug}".encode()).hexdigest() return f"subagent:{digest[:48]}" def _with_run_payload(subagent_run: dict[str, Any], run) -> dict[str, Any]: if not run: return subagent_run return {**subagent_run, **_agent_run_state_payload(run)} 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, } tool_result = _tool_result_with_thread_id(subagent_run["child_thread_id"], final_text) update: dict[str, Any] = {"messages": [ToolMessage(tool_result, tool_call_id=tool_call_id)]} if artifacts: update["artifacts"] = artifacts update["subagent_runs"] = [subagent_run] return Command(update=update) def _reused_run_response(run, tool_call_id: str, subagent_run: dict[str, Any]) -> Command: status = str(getattr(run, "status", "") or "unknown") if status == "completed": message = "子智能体任务已完成，未重复执行。" elif status in _TERMINAL_RUN_STATUSES: error_message = str(getattr(run, "error_message", "") or "") message = f"子智能体任务已结束，状态：{status}。{error_message}".strip() else: message = f"子智能体任务已存在，当前状态：{status}，未重复提交。" subagent_run = { **subagent_run, **_agent_run_state_payload(run), "status": status, "result_preview": _preview_text(message), } tool_message = ToolMessage( _tool_result_with_thread_id(subagent_run["child_thread_id"], message), tool_call_id=tool_call_id, ) return Command(update={"messages": [tool_message], "subagent_runs": [subagent_run]}) def _agent_run_state_payload(run) -> dict[str, Any]: payload = { "run_id": run.id, "status": run.status, "parent_agent_run_id": run.parent_agent_run_id, "created_at": utc_isoformat(run.created_at) if run.created_at else None, "completed_at": utc_isoformat(run.finished_at) if run.finished_at else None, "error": run.error_message, } return {key: value for key, value in payload.items() if value is not None} 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}" tool_result = _tool_result_with_thread_id(subagent_run["child_thread_id"], message) update = { "messages": [ToolMessage(tool_result, 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, continuing: bool = False, ) -> dict[str, Any]: state = {} if continuing else {key: runtime.state[key] for key in _CHILD_STATE_INHERIT_KEYS if key in runtime.state} 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, subagent_type: str, run_id: str | None = None, request_id: str | None = None, ) -> 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, "subagent_type": subagent_type, "subagent_thread_id": child_thread_id, "subagent_tool_call_id": runtime.tool_call_id, "run_id": run_id, "request_id": request_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)] async def _create_subagent_run( self, *, child_thread_id: str, description: str, subagent_type: str, agent: Agent, uid: str, parent_thread_id: str, tool_call_id: str, continuing: bool, ): parent_run_id = str(getattr(self.parent_context, "run_id", "") or "").strip() if not parent_run_id: raise ValueError("当前运行时缺少父运行 ID，无法记录子智能体运行") async with pg_manager.get_async_session_context() as db: repo = AgentRunRepository(db) parent_run = await repo.get_run_for_user(parent_run_id, uid) if not parent_run: raise ValueError("父运行任务不存在") if continuing: previous = await repo.get_latest_subagent_run_by_thread_for_user(child_thread_id, uid) if not previous or previous.conversation_id != parent_run.conversation_id: raise ValueError( f"无法继续子智能体线程 {child_thread_id}：当前对话中没有找到对应的子智能体运行记录" ) if previous.agent_id != subagent_type: raise ValueError( f"无法继续子智能体线程 {child_thread_id}：该线程属于子智能体 {previous.agent_id or '未知'}" ) request_id = _subagent_request_id(parent_run_id, child_thread_id, tool_call_id, agent.slug) existing = await repo.get_run_by_request_id(request_id) if existing: return existing, False run = await repo.create_run( run_id=str(uuid.uuid4()), thread_id=child_thread_id, agent_id=subagent_type, uid=uid, request_id=request_id, conversation_id=parent_run.conversation_id, parent_agent_run_id=parent_run.id, run_type="subagent", checkpoint_thread_id=child_thread_id, input_payload={ "description": description, "tool_call_id": tool_call_id, "subagent_type": subagent_type, "subagent_name": agent.name, "parent_thread_id": parent_thread_id, "child_thread_id": child_thread_id, "continuing": continuing, }, ) return await repo.mark_running(run.id), True async def _set_subagent_run_status( self, run_id: str, status: str, *, error_type: str | None = None, error_message: str | None = None, ): async with pg_manager.get_async_session_context() as db: return await AgentRunRepository(db).set_terminal_status( run_id, status=status, error_type=error_type, error_message=error_message, ) 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, thread_id: Annotated[str | None, THREAD_ID_ARG] = None, ) -> str: return "task 工具仅支持异步调用" async def atask( description: Annotated[str, TASK_DESCRIPTION_ARG], subagent_type: Annotated[str, SUBAGENT_TYPE_ARG], runtime: ToolRuntime, thread_id: Annotated[str | None, THREAD_ID_ARG] = None, ) -> 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, continuing = _new_child_thread_id( thread_id, parent_thread_id=parent_thread_id, agent_slug=agent.slug, tool_call_id=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(), } try: run, is_new_run = await self._create_subagent_run( child_thread_id=child_thread_id, description=description, subagent_type=subagent_type, agent=agent, uid=uid, parent_thread_id=parent_thread_id, tool_call_id=runtime.tool_call_id, continuing=continuing, ) except ValueError as exc: return str(exc) subagent_run = _with_run_payload(subagent_run, run) if not is_new_run: return _reused_run_response(run, runtime.tool_call_id, subagent_run) child_context = backend.context_schema() config_context = (agent.config_json or {}).get("context") if isinstance(agent.config_json, dict) else None child_input_context = await build_agent_input_context( config_context if isinstance(config_context, dict) else {}, thread_id=child_thread_id, uid=uid, run_id=run.id, request_id=run.request_id, ) child_context.update_from_dict(child_input_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.run_id = run.id child_context.request_id = run.request_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, continuing=continuing, ), 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, subagent_type=subagent_type, run_id=run.id, request_id=run.request_id, ), context=child_context, ) except asyncio.CancelledError: await self._set_subagent_run_status(run.id, "cancelled") raise except Exception as exc: failed_run = await self._set_subagent_run_status( run.id, "failed", error_type=type(exc).__name__, error_message=str(exc), ) return _failed_tool_response(exc, runtime.tool_call_id, _with_run_payload(subagent_run, failed_run)) completed_run = await self._set_subagent_run_status(run.id, "completed") return _completed_tool_response( result, runtime.tool_call_id, _with_run_payload(subagent_run, completed_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)