247 lines
10 KiB
Python
247 lines
10 KiB
Python
"""cron 工具：给 Agent 提供“定时任务管理”能力。
|
||
|
||
这个工具是 LLM 在对话中可调用的 function tool，主要负责三件事：
|
||
1. `add`：创建一个定时任务（周期/cron/一次性）；
|
||
2. `list`：列出现有任务；
|
||
3. `remove`：删除指定任务。
|
||
|
||
设计定位说明：
|
||
- 本工具只做“任务管理面”，不直接负责“定时器循环”；
|
||
- 真正的调度与执行由 `CronService` 统一负责（start/stop/on_job）；
|
||
- 工具层通过 `set_context(channel, chat_id)` 注入当前会话路由，
|
||
从而让定时任务在触发后把结果回投到正确会话。
|
||
"""
|
||
|
||
from typing import Any
|
||
|
||
from nanobot.agent.tools.base import Tool
|
||
from nanobot.cron.service import CronService
|
||
from nanobot.cron.types import CronSchedule
|
||
|
||
|
||
class CronTool(Tool):
|
||
"""对话可调用的 cron 管理工具。
|
||
|
||
调用来源：
|
||
- 主 agent 在工具调用回合中发起 `cron(...)`。
|
||
|
||
关键约束：
|
||
- action 仅支持 `add/list/remove` 三种；
|
||
- `add` 必须带 message，并且必须先注入 session 上下文（channel/chat_id）；
|
||
- 时间相关参数三选一：`every_seconds` / `cron_expr` / `at`。
|
||
"""
|
||
|
||
def __init__(self, cron_service: CronService):
|
||
# 持有同一个 CronService 实例，保证：
|
||
# 1) CLI 命令与 agent 工具看到同一份 jobs.json；
|
||
# 2) 任务状态（next_run、enabled）在进程内一致。
|
||
self._cron = cron_service
|
||
# 路由上下文由 AgentLoop 每轮注入。
|
||
# 任务触发时将按该路由把结果投递回原会话。
|
||
self._channel = ""
|
||
self._chat_id = ""
|
||
self._session_key = ""
|
||
|
||
def set_context(self, channel: str, chat_id: str, session_key: str | None = None) -> None:
|
||
"""设置当前会话路由上下文。
|
||
|
||
为什么需要它：
|
||
- 用户在 A 会话里让 agent“每天提醒我”，
|
||
任务未来触发时应回到 A，而不是误发到其他会话。
|
||
- 因此 channel/chat_id 不依赖模型每次显式传参，
|
||
而是由运行时在调用前预注入默认目标。
|
||
"""
|
||
self._channel = channel
|
||
self._chat_id = chat_id
|
||
self._session_key = session_key or f"{channel}:{chat_id}"
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
# 暴露给模型的工具名。模型会以 `cron(...)` 发起 function call。
|
||
return "cron"
|
||
|
||
@property
|
||
def description(self) -> str:
|
||
# 给模型看的简要能力描述，尽量短而明确。
|
||
return "Schedule reminders and recurring tasks. Actions: add, list, remove. Use mode=reminder or task."
|
||
|
||
@property
|
||
def parameters(self) -> dict[str, Any]:
|
||
# OpenAI function schema：
|
||
# - 定义参数结构与类型；
|
||
# - 由 ToolRegistry 在调用前做基础参数校验。
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"action": {
|
||
"type": "string",
|
||
"enum": ["add", "list", "remove"],
|
||
"description": "Action to perform"
|
||
},
|
||
"message": {
|
||
"type": "string",
|
||
# add 时的任务文本：
|
||
# - 既可做“纯提醒文案”，也可做“交给 agent 执行的提示”。
|
||
"description": "Reminder message (for add)"
|
||
},
|
||
"mode": {
|
||
"type": "string",
|
||
"enum": ["reminder", "task"],
|
||
"description": "Execution mode: reminder sends message directly; task re-enters agent"
|
||
},
|
||
"every_seconds": {
|
||
"type": "integer",
|
||
# 固定间隔调度（单位秒），内部会转换为毫秒。
|
||
"description": "Interval in seconds (for recurring tasks)"
|
||
},
|
||
"cron_expr": {
|
||
"type": "string",
|
||
# 标准 cron 表达式（5 段），例如每天 9 点：0 9 * * *
|
||
"description": "Cron expression like '0 9 * * *' (for scheduled tasks)"
|
||
},
|
||
"tz": {
|
||
"type": "string",
|
||
# 仅与 cron_expr 搭配使用的 IANA 时区。
|
||
"description": "IANA timezone for cron expressions (e.g. 'America/Vancouver')"
|
||
},
|
||
"at": {
|
||
"type": "string",
|
||
# 一次性触发时间，ISO 格式（本地/带偏移都可由 fromisoformat 解析）。
|
||
"description": "ISO datetime for one-time execution (e.g. '2026-02-12T10:30:00')"
|
||
},
|
||
"job_id": {
|
||
"type": "string",
|
||
"description": "Job ID (for remove)"
|
||
}
|
||
},
|
||
"required": ["action"]
|
||
}
|
||
|
||
async def execute(
|
||
self,
|
||
action: str,
|
||
message: str = "",
|
||
mode: str | None = None,
|
||
every_seconds: int | None = None,
|
||
cron_expr: str | None = None,
|
||
tz: str | None = None,
|
||
at: str | None = None,
|
||
job_id: str | None = None,
|
||
**kwargs: Any
|
||
) -> str:
|
||
"""工具主入口：按 action 分发到具体处理函数。
|
||
|
||
注意：
|
||
- 这里不直接抛异常给上层；尽量返回可读错误字符串。
|
||
- 真正未捕获异常（如非法日期解析）会被 ToolRegistry 包装成 Error 文本。
|
||
"""
|
||
# add：创建任务（并立即持久化），返回任务 ID。
|
||
if action == "add":
|
||
return self._add_job(message, mode, every_seconds, cron_expr, tz, at)
|
||
# list：只读取并格式化输出，不改状态。
|
||
elif action == "list":
|
||
return self._list_jobs()
|
||
# remove：按 ID 删除任务并重置调度器。
|
||
elif action == "remove":
|
||
return self._remove_job(job_id)
|
||
# schema 已限制枚举，这里是兜底防御。
|
||
return f"Unknown action: {action}"
|
||
|
||
def _add_job(
|
||
self,
|
||
message: str,
|
||
mode: str | None,
|
||
every_seconds: int | None,
|
||
cron_expr: str | None,
|
||
tz: str | None,
|
||
at: str | None,
|
||
) -> str:
|
||
"""创建任务并写入 CronService。
|
||
|
||
参数优先级（互斥选择）：
|
||
1. `every_seconds` -> 固定间隔任务
|
||
2. `cron_expr` -> cron 表达式任务
|
||
3. `at` -> 一次性任务（执行后自动删除）
|
||
"""
|
||
# message 是 add 的必填语义字段：没有内容就无法定义“要做什么”。
|
||
if not message:
|
||
return "Error: message is required for add"
|
||
# channel/chat_id 由 AgentLoop 注入；
|
||
# 若缺失，说明当前调用上下文不完整，无法保证结果回投目标正确。
|
||
if not self._channel or not self._chat_id:
|
||
return "Error: no session context (channel/chat_id)"
|
||
# 时区仅对 cron 表达式有意义；避免用户误把 tz 用在 every/at 上。
|
||
if tz and not cron_expr:
|
||
return "Error: tz can only be used with cron_expr"
|
||
# 尽早校验时区，提前给出明确错误，避免把非法数据写入存储。
|
||
if tz:
|
||
from zoneinfo import ZoneInfo
|
||
try:
|
||
ZoneInfo(tz)
|
||
except (KeyError, Exception):
|
||
return f"Error: unknown timezone '{tz}'"
|
||
|
||
# mode 缺省时默认按“提醒”处理：
|
||
# - 与 cron skill 的说明一致；
|
||
# - 避免把原始建任务指令再次送回 agent，造成任务自复制。
|
||
normalized_mode = (mode or "reminder").strip().lower()
|
||
if normalized_mode not in {"reminder", "task"}:
|
||
return "Error: mode must be 'reminder' or 'task'"
|
||
payload_kind = "system_event" if normalized_mode == "reminder" else "agent_turn"
|
||
|
||
# 构建调度对象：
|
||
# - CronService 内部统一使用毫秒时间戳；
|
||
# - `at` 任务默认 delete_after_run=True，执行一次后自动移除。
|
||
delete_after = False
|
||
if every_seconds:
|
||
schedule = CronSchedule(kind="every", every_ms=every_seconds * 1000)
|
||
elif cron_expr:
|
||
schedule = CronSchedule(kind="cron", expr=cron_expr, tz=tz)
|
||
elif at:
|
||
from datetime import datetime
|
||
# fromisoformat 解析失败会抛 ValueError，
|
||
# 该异常会由 ToolRegistry 统一转换为错误字符串返回给模型。
|
||
dt = datetime.fromisoformat(at)
|
||
at_ms = int(dt.timestamp() * 1000)
|
||
schedule = CronSchedule(kind="at", at_ms=at_ms)
|
||
delete_after = True
|
||
else:
|
||
return "Error: either every_seconds, cron_expr, or at is required"
|
||
|
||
# 创建任务并持久化：
|
||
# - name 使用 message 前 30 字符做简短标题，便于列表展示；
|
||
# - deliver=True：任务触发后默认向当前会话投递结果；
|
||
# - channel/to 使用注入上下文，确保消息路由一致。
|
||
job = self._cron.add_job(
|
||
name=message[:30],
|
||
schedule=schedule,
|
||
message=message,
|
||
payload_kind=payload_kind,
|
||
session_key=self._session_key or None,
|
||
deliver=True,
|
||
channel=self._channel,
|
||
to=self._chat_id,
|
||
delete_after_run=delete_after,
|
||
)
|
||
# 返回简明确认文本，便于模型后续引用 job_id 做删除或说明。
|
||
return f"Created {normalized_mode} job '{job.name}' (id: {job.id})"
|
||
|
||
def _list_jobs(self) -> str:
|
||
"""列出当前可见任务（默认仅启用任务）。"""
|
||
jobs = self._cron.list_jobs()
|
||
if not jobs:
|
||
return "No scheduled jobs."
|
||
# 输出格式保持轻量，避免把过多状态塞给模型。
|
||
# 详细状态（next_run/last_error）可在 CLI 的 `nanobot cron list` 查看。
|
||
lines = [f"- {j.name} (id: {j.id}, {j.schedule.kind})" for j in jobs]
|
||
return "Scheduled jobs:\n" + "\n".join(lines)
|
||
|
||
def _remove_job(self, job_id: str | None) -> str:
|
||
"""按 ID 删除任务。"""
|
||
if not job_id:
|
||
return "Error: job_id is required for remove"
|
||
# remove_job 返回 bool，工具层负责转换成对话友好的文案。
|
||
if self._cron.remove_job(job_id):
|
||
return f"Removed job {job_id}"
|
||
return f"Job {job_id} not found"
|