zcbot/core/tool_registry.py

265 lines
12 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""声明式工具注册表(§3.4 / §7.5 #7;从 agent_builder.build_agent 析出,2026-07-23)。
背景:此前 build_agent 里 ~220 行工具注册,同一个「有 key/开关才注册」意图用了
四种写法(inline getenv / `*_configured()` 助手 / `Config.load()` 返 None / caps
开关),约定靠注释里的"沿用范式"口头传承。本模块把它收敛为一张表:
(组名, gate, factory)
- gate:零参 -> bool。env key / 渠道开关 / 能力开关全走这一种形态,False 则该组
工具对本 run 不存在(agent 看不到一个永远报错的工具 —— §7.5 #7 红线的代码强制)。
- factory:零参 -> list[Tool]。变体选择(如 seedance/vision 取 yaml 第一个 variant)
在 factory 里做,选不出返 [] 即自然不挂。
新增工具 = 表里加一行。ToolContext 装 build_agent 已解析好的上下文(路径 / 身份 /
媒体配置),factory 从 ctx 取参 —— 不在本模块重复读 env / yaml。
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional
from uuid import UUID
from tools.ask_user import AskUserTool
from tools.check_process import CheckProcessTool
from tools.documents import DocumentDownloadTool, DocumentListKbTool, DocumentSearchTool
from tools.fs import EditTool, GlobTool, GrepTool, ReadTool, WriteTool
from tools.gpt_image import GptImageTool
from tools.look_at_image import LookAtImageTool
from tools.materials_project import (
MaterialsProjectGetEntriesTool,
MaterialsProjectGetStructureTool,
MaterialsProjectSearchSummaryTool,
)
from tools.office_to_pdf import OfficeToPdfTool
from tools.read_document import ReadDocumentTool
from tools.rename_working_dir import RenameWorkingDirTool
from tools.run_python import RunPythonTool
from tools.schedule import (
ScheduleCancelTool, ScheduleCreateTool, ScheduleListTool, ScheduleUpdateTool,
)
from tools.seedance import SeedanceTool
from tools.seedream import SeedreamTool
from tools.send_email import SendEmailTool, smtp_configured
from tools.shell import ShellTool
from tools.skill_authoring import ForkSkillTool, SaveSkillTool
from tools.skill_tool import LoadSkillTool
from tools.task_progress import TaskProgressTool
from tools.transcribe_audio import TranscribeAudioTool
from tools.web_fetch import WebFetchTool
from tools.web_search import WebSearchTool
from tools.wechat_bot import WechatPushTool, wechat_push_available
from core.asr_lfasr import is_configured as lfasr_configured
from core.bocha_client import BochaConfig
from core.task_actions import DeferredTaskActions
@dataclass
class ToolContext:
"""build_agent 解析好的、工具装配需要的全部上下文。"""
tool_base: Path # fs/shell 类工具的 base_dir(cwd / task 目录)
ur_path: Path # user_root(输出渲染相对路径 + host-side 落点)
working_dir_path: Path # 该 task 的宿主工作目录绝对路径
task_id: str
uid: UUID
cfg: dict # config/agent.yaml(quotas 段)
caps: Any # ModelCapabilities(enable_run_python)
skills: Any # SkillRegistry
cancel_check: Optional[Callable[[], bool]]
scheduled_run: bool
deferred_actions: DeferredTaskActions
# 媒体(上游已 load 一次,避免重复读 yaml):
ark_cfg: Any # ArkConfig | None(豆包;None=ARK_API_KEY 缺)
img_provider: str # "doubao" / "unifyllm" / ""
img_key: str # 选中的 image variant key
img_cfg: Optional[dict] # 选中的 image variant 配置(None=不挂图像工具)
img_provider_cfg: Any # 该 provider 的 ArkConfig
video_variant: str # caller 指定的 video variant key(空=yaml 第一个)
office_to_pdf_available: bool # backend host 是否可调用 LibreOffice
def _env_set(name: str) -> Callable[[], bool]:
return lambda: bool(os.getenv(name, "").strip())
def _pick_variant(section: dict, preferred: str = "") -> tuple[str, Optional[dict]]:
"""从 yaml 段选 variant:preferred 命中优先,否则第一个 dict 条目;无 → ("", None)。"""
if preferred:
v = section.get(preferred)
if isinstance(v, dict):
return preferred, v
for key, v in section.items():
if isinstance(v, dict):
return key, v
return "", None
def build_tools(ctx: ToolContext) -> dict[str, Any]:
"""按注册表装配本 run 的工具集。gate=False / factory 返 [] 的组自然不挂。"""
base = dict(base_dir=ctx.tool_base, user_root=ctx.ur_path)
wd_base = dict(base_dir=ctx.working_dir_path, user_root=ctx.ur_path)
quotas = ctx.cfg.get("quotas") or {}
images_per_day = int(quotas.get("images_per_day", 0))
videos_per_day = int(quotas.get("videos_per_day", 0))
def _core() -> list:
# shell/run_python 带 task_id:background=true 的 bg proc 状态锚定
# `<user_root>/.zcbot_procs/<task_id>/`(DESIGN §8.12)。check_process 是它们的
# 配套查询/终止工具,host in-process(docker 模式下状态文件也在宿主侧)。
# web_fetch 无需 API key,始终可用。
return [
TaskProgressTool(**base),
AskUserTool(**base),
ReadTool(**base), WriteTool(**base), EditTool(**base),
GlobTool(**base), GrepTool(**base),
ShellTool(task_id=ctx.task_id, **base),
CheckProcessTool(task_id=ctx.task_id, **base),
WebFetchTool(**base),
]
def _task_actions() -> list:
return [
RenameWorkingDirTool(
ctx.deferred_actions,
working_dir=ctx.working_dir_path,
**base,
)
]
def _document_search() -> list:
return [
DocumentListKbTool(**base),
DocumentSearchTool(**base),
DocumentDownloadTool(working_dir=ctx.working_dir_path, **base),
]
def _materials_project() -> list:
return [
MaterialsProjectSearchSummaryTool(working_dir=ctx.working_dir_path, **base),
MaterialsProjectGetStructureTool(working_dir=ctx.working_dir_path, **base),
MaterialsProjectGetEntriesTool(working_dir=ctx.working_dir_path, **base),
]
def _load_skill() -> list:
# LoadSkillTool 返回头里的 dir 由 registry 按 skill.source 给容器内路径
# (内置 → /sandbox/skills,用户 → /workspace/.skills);host backend → host 绝对路径。
return [LoadSkillTool(registry=ctx.skills, **base)]
def _skill_authoring() -> list:
# 用户 skill 创作:恒挂(每个用户都能造自己的 skill)。host-side 直接写
# user_root/.skills —— 不走沙箱 fs(其 base_dir 锚 cwd / 容器 wd,够不到 .skills)。
d = ctx.ur_path / ".skills"
return [
SaveSkillTool(d, ctx.skills, **base),
ForkSkillTool(d, ctx.skills, **base),
]
def _schedules() -> list:
# 定时任务管理(§8.5)增删查。user_id 由 ctor 注入,不信模型传的 id。
return [
ScheduleCreateTool(ctx.uid, **base),
ScheduleListTool(ctx.uid, **base),
ScheduleUpdateTool(ctx.uid, **base),
ScheduleCancelTool(ctx.uid, **base),
]
def _send_email() -> list:
# base_dir 用 working_dir_path(宿主 task 目录):send_email 在宿主进程读附件,
# docker 下 agent 给的相对路径相对容器 workdir=task_dir,翻回宿主即 working_dir_path;
# tool 内 _resolve_user_file 再处理 /workspace 容器绝对路径。
return [SendEmailTool(**wd_base)]
def _wechat_push() -> list:
# base_dir 同 send_email(宿主进程读待发文件)。交互与定时 run 都可用
# (定时简报可主动推回用户微信,24h 窗口内)。
return [WechatPushTool(ctx.uid, task_id=ctx.task_id, **wd_base)]
def _run_python() -> list:
return [RunPythonTool(task_id=ctx.task_id, **base)]
def _office_to_pdf() -> list:
# LibreOffice 只装在 backend hostDocker agent 通过 typed tool 转换用户目录内
# 已有 Office 文件,不在沙盒 shell 里寻找/安装 soffice。
return [OfficeToPdfTool(**wd_base)]
def _image() -> list:
# 图像生成跨 provider 二选一(选择在 build_agent 上半部定好,与 system prompt
# 媒体段同源);本次 run 锁定该 variant,下一条消息可重选。
if ctx.img_cfg is None:
return []
cls_kwargs = dict(
image_variant_cfg=ctx.img_cfg, variant_key=ctx.img_key,
working_dir=ctx.working_dir_path, task_id=ctx.task_id, user_id=ctx.uid,
daily_limit=images_per_day, **base,
)
if ctx.img_provider == "doubao":
return [SeedreamTool(ark_cfg=ctx.img_provider_cfg, **cls_kwargs)]
return [GptImageTool(gw_cfg=ctx.img_provider_cfg, **cls_kwargs)]
def _video() -> list:
# 视频仍豆包独有。cancel_check:轮询期(典型 30-90s)响应用户停止按钮;
# 远端 cgt 任务无 cancel API,best-effort 不动远端。
key, v = _pick_variant(ctx.ark_cfg.raw.get("video") or {}, ctx.video_variant)
if v is None:
return []
return [SeedanceTool(
ark_cfg=ctx.ark_cfg, video_variant_cfg=v, variant_key=key,
working_dir=ctx.working_dir_path, task_id=ctx.task_id, user_id=ctx.uid,
cancel_check=ctx.cancel_check, daily_limit=videos_per_day, **base,
)]
def _vision() -> list:
# 看图 + 文档理解共用 vision variant(读图不分档,固定第一个)。
key, v = _pick_variant(ctx.ark_cfg.raw.get("vision") or {})
if v is None:
return []
kw = dict(
ark_cfg=ctx.ark_cfg, vision_variant_cfg=v, variant_key=key,
working_dir=ctx.working_dir_path, task_id=ctx.task_id, user_id=ctx.uid,
**base,
)
return [LookAtImageTool(**kw), ReadDocumentTool(**kw)]
def _transcribe() -> list:
# 录音文件转写(讯飞 LFASR,与 IAT 听写是两个服务两套 key)。
return [TranscribeAudioTool(
working_dir=ctx.working_dir_path, cancel_check=ctx.cancel_check, **base,
)]
def _web_search() -> list:
return [WebSearchTool(cfg=BochaConfig.load())]
# ── 注册表:(组名, gate, factory)。gate 判定统一零参 bool;新工具在此加行 ──
registry: list[tuple[str, Callable[[], bool], Callable[[], list]]] = [
("core", lambda: True, _core),
# 当前 working_dir 只能延迟到交互 run 收尾后改名;定时 run 不允许自行改目录。
("task_actions", lambda: not ctx.scheduled_run, _task_actions),
# Secret-bearing 域工具一律 host-side、仅对应 env 存在才注册(§7.5 #7):
# key 绝不进 run_python / 沙箱。
("document_search", _env_set("DOCUMENT_SEARCH_API_KEY"), _document_search),
("materials_project", _env_set("MP_API_KEY"), _materials_project),
("load_skill", lambda: bool(ctx.skills.skills), _load_skill),
("skill_authoring", lambda: True, _skill_authoring),
# 定时 run 内不挂 schedule_*(防任务造任务自我繁殖);仅交互对话可建/管 job。
("schedules", lambda: not ctx.scheduled_run, _schedules),
("send_email", smtp_configured, _send_email),
("wechat_push", wechat_push_available, _wechat_push),
("run_python", lambda: ctx.caps.enable_run_python, _run_python),
("office_to_pdf", lambda: ctx.office_to_pdf_available, _office_to_pdf),
("image", lambda: ctx.img_cfg is not None, _image),
("video", lambda: ctx.ark_cfg is not None, _video),
("vision", lambda: ctx.ark_cfg is not None, _vision),
("transcribe_audio", lfasr_configured, _transcribe),
("web_search", lambda: BochaConfig.load() is not None, _web_search),
]
tools: dict[str, Any] = {}
for _name, gate, factory in registry:
if gate():
for t in factory():
tools[t.name] = t
return tools