109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
from typing import Any
|
|
|
|
from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware
|
|
from langchain.agents import create_agent
|
|
from langchain.agents.middleware import ModelRetryMiddleware, TodoListMiddleware
|
|
from langchain.agents.middleware.types import AgentMiddleware
|
|
|
|
from yuxi.agents import BaseAgent, BaseState, load_chat_model
|
|
from yuxi.agents.backends import create_agent_filesystem_middleware
|
|
from yuxi.agents.buildin.chatbot.prompt import TODO_MID_PROMPT, build_prompt_with_context
|
|
from yuxi.agents.buildin.subagent.context import SubAgentContext
|
|
from yuxi.agents.context import prepare_agent_runtime_context
|
|
from yuxi.agents.middlewares import create_summary_middleware, save_attachments_to_fs
|
|
from yuxi.agents.middlewares.knowledge_base import KnowledgeBaseMiddleware
|
|
from yuxi.agents.middlewares.skills import SkillsMiddleware
|
|
from yuxi.agents.toolkits.service import resolve_configured_runtime_tools
|
|
|
|
|
|
_SUBAGENT_DISABLED_TOOLS = frozenset({"present_artifacts", "ask_user_question"})
|
|
|
|
|
|
def _tool_name(tool) -> str | None:
|
|
if isinstance(tool, dict):
|
|
name = tool.get("name")
|
|
else:
|
|
name = getattr(tool, "name", None)
|
|
return name if isinstance(name, str) else None
|
|
|
|
|
|
def _filter_disabled_tools(tools):
|
|
return [tool for tool in tools if _tool_name(tool) not in _SUBAGENT_DISABLED_TOOLS]
|
|
|
|
|
|
class _SubAgentToolFilterMiddleware(AgentMiddleware[Any, Any, Any]):
|
|
def wrap_model_call(self, request, handler):
|
|
return handler(request.override(tools=_filter_disabled_tools(request.tools or [])))
|
|
|
|
async def awrap_model_call(self, request, handler):
|
|
return await handler(request.override(tools=_filter_disabled_tools(request.tools or [])))
|
|
|
|
|
|
async def _build_middlewares(context):
|
|
summary_trigger_tokens = getattr(context, "summary_threshold", 100) * 1024
|
|
summary_middleware = create_summary_middleware(
|
|
model=load_chat_model(fully_specified_name=context.model),
|
|
trigger=("tokens", summary_trigger_tokens),
|
|
keep=("tokens", summary_trigger_tokens // 2),
|
|
trim_tokens_to_summarize=4000,
|
|
)
|
|
|
|
return [
|
|
create_agent_filesystem_middleware(
|
|
getattr(context, "tool_token_limit", 20) * 1024,
|
|
context=context,
|
|
),
|
|
save_attachments_to_fs,
|
|
KnowledgeBaseMiddleware(),
|
|
SkillsMiddleware(),
|
|
summary_middleware,
|
|
TodoListMiddleware(system_prompt=TODO_MID_PROMPT),
|
|
PatchToolCallsMiddleware(),
|
|
_SubAgentToolFilterMiddleware(),
|
|
ModelRetryMiddleware(),
|
|
]
|
|
|
|
|
|
class SubAgentBackend(BaseAgent):
|
|
name = "子智能体"
|
|
description = "用于被主智能体通过 task 工具调用的专用智能体后端。"
|
|
capabilities = ["file_upload", "files"]
|
|
context_schema = SubAgentContext
|
|
|
|
async def get_info(
|
|
self,
|
|
include_configurable_items: bool = True,
|
|
user_role: str | None = None,
|
|
db=None,
|
|
user=None,
|
|
):
|
|
info = await super().get_info(
|
|
include_configurable_items=include_configurable_items,
|
|
user_role=user_role,
|
|
db=db,
|
|
user=user,
|
|
)
|
|
tools_item = (info.get("configurable_items") or {}).get("tools")
|
|
if isinstance(tools_item, dict):
|
|
tools_item["options"] = [
|
|
option
|
|
for option in tools_item.get("options") or []
|
|
if option.get("key") not in _SUBAGENT_DISABLED_TOOLS
|
|
]
|
|
return info
|
|
|
|
async def get_graph(self, context=None, **kwargs):
|
|
context = await prepare_agent_runtime_context(
|
|
context or self.context_schema(),
|
|
context_schema=self.context_schema,
|
|
)
|
|
|
|
return create_agent(
|
|
model=load_chat_model(fully_specified_name=context.model),
|
|
tools=_filter_disabled_tools(await resolve_configured_runtime_tools(context)),
|
|
system_prompt=build_prompt_with_context(context),
|
|
middleware=await _build_middlewares(context),
|
|
state_schema=BaseState,
|
|
checkpointer=await self._get_checkpointer(),
|
|
)
|