Asfmq
f885c0a4a8
- 服务层拆分：删除 api/helpers.rs，新增 citation/note/session/pipeline/paper/vision 独立服务模块
- Agent 工具精简：paper_content+paper_outline 合并为 paper.rs，图片分析逻辑下沉至 services/vision
- 并发模型升级：std::sync::{Mutex,RwLock} → tokio::sync::{Mutex,RwLock}，消除 async
上下文中的阻塞风险
- 客户端加固：HTTP 客户端统一超时配置、ADS 429 / arXiv 503 自动重试、构造函数返回 Result
- 启动安全：全局 panic hook 日志化、空密码拒绝启动、向量表维度不匹配需显式确认
- CLI 扩展：构建完整 AppState 复用服务层，新增 Content/Outline/Citations/Search/Process 子命令
- 前端：移动端汉堡菜单、侧栏滑出面板、引用星系触屏手势（单指拖拽/双指缩放）
757 lines
26 KiB
Rust
757 lines
26 KiB
Rust
// src/agent/tools/mod.rs
|
||
//
|
||
// 科研智能体工具集定义与实现。
|
||
// 每个工具遵循 AgentTool trait，向大模型声明 JSON Schema 参数定义，
|
||
// 并在 execute 中调用已有的服务层完成实际业务操作。
|
||
//
|
||
// 按功能域拆分为子模块：
|
||
// filesystem/ — 文件 I/O（read、grep、glob、bash、write、edit）
|
||
// astro/ — 天文科研（文献搜索/下载/解析、RAG、天体查询、笔记）
|
||
// team.rs — 团队协作
|
||
// todo.rs — 任务规划
|
||
// compress.rs — 手动上下文压缩
|
||
|
||
use async_trait::async_trait;
|
||
use serde_json::json;
|
||
use std::sync::Arc;
|
||
use tokio::sync::RwLock;
|
||
|
||
use crate::agent::runtime::file_cache::FileStateCache;
|
||
use crate::agent::skills::SkillRegistry;
|
||
use crate::api::AppState;
|
||
use crate::clients::llm::ToolDefinition;
|
||
|
||
pub mod ask_user;
|
||
pub mod astro;
|
||
mod background;
|
||
mod compress;
|
||
mod filesystem;
|
||
pub mod memory;
|
||
pub mod persist;
|
||
pub mod search_history;
|
||
mod skill;
|
||
pub mod subagent;
|
||
mod team;
|
||
mod todo;
|
||
|
||
pub use ask_user::AskUserTool;
|
||
pub use astro::analyze_image::AnalyzeImageTool;
|
||
pub use astro::research::citation_network::GetCitationNetworkTool;
|
||
pub use astro::research::library_search::SearchLocalLibraryTool;
|
||
pub use astro::research::metadata::GetPaperMetadataTool;
|
||
pub use astro::research::note::SaveNoteTool;
|
||
pub use astro::research::paper::{GetPaperContentTool, GetPaperOutlineTool};
|
||
pub use astro::research::rag::RagSearchTool;
|
||
pub use astro::research::target::QueryTargetTool;
|
||
pub use astro::system::process::ProcessPaperTool;
|
||
pub use astro::system::search::SearchPapersTool;
|
||
pub use background::{BgTaskCheckTool, BgTaskRunTool};
|
||
pub use compress::CompressTool;
|
||
pub use filesystem::{
|
||
FileEditTool, FileWriteTool, GlobFilesTool, GrepFilesTool, ReadFileTool, RunBashTool,
|
||
};
|
||
pub use search_history::SearchHistoryTool;
|
||
pub use skill::LoadSkillTool;
|
||
pub use subagent::SubAgentTool;
|
||
pub use team::{CheckTeamInboxTool, SendTeammateMessageTool, SpawnTeammateTool, TeamBroadcastTool};
|
||
pub use todo::persist_tasks;
|
||
pub use todo::TodoWriteTool;
|
||
|
||
/// 工具执行上下文，封装全局共享状态
|
||
#[derive(Clone)]
|
||
pub struct ToolContext {
|
||
pub app_state: Arc<AppState>,
|
||
/// 当前会话 ID（用于工具将数据关联到正确的 session）
|
||
pub session_id: String,
|
||
/// 静默模式：子代理运行时为 true，跳过用户权限提示
|
||
pub silent: bool,
|
||
/// 文件状态缓存（跨工具调用共享，用于 Read 去重）
|
||
pub read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||
/// 当前 SSE 发送通道（工具可通过此通道向用户推送中间事件）
|
||
pub sse_tx: Option<tokio::sync::mpsc::UnboundedSender<crate::agent::runtime::AgentStreamEvent>>,
|
||
/// 是否启用 LLM 思考模式（继承自父代理配置）
|
||
pub enable_thinking: bool,
|
||
/// 附加允许目录（扩展文件沙箱范围）
|
||
pub additional_allowed_dirs: Vec<String>,
|
||
/// 当前工具调用 ID（工具可通过此 ID 将流式输出关联到自身）
|
||
pub tool_call_id: Option<String>,
|
||
/// 工具输出最大字符数（继承自 AgentConfig，工具可据此截断输出）
|
||
pub max_output_chars: usize,
|
||
}
|
||
|
||
impl ToolContext {
|
||
/// 创建标准上下文
|
||
pub fn new(app_state: Arc<AppState>) -> Self {
|
||
ToolContext {
|
||
app_state,
|
||
session_id: String::new(),
|
||
silent: false,
|
||
read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())),
|
||
sse_tx: None,
|
||
enable_thinking: false,
|
||
additional_allowed_dirs: Vec::new(),
|
||
tool_call_id: None,
|
||
max_output_chars: 4000,
|
||
}
|
||
}
|
||
|
||
/// 创建带共享文件缓存的上下文（用于 AgentRuntime 保持同一 cache 实例）
|
||
pub fn with_file_cache(
|
||
app_state: Arc<AppState>,
|
||
read_file_state: Arc<std::sync::Mutex<FileStateCache>>,
|
||
) -> Self {
|
||
ToolContext {
|
||
app_state,
|
||
session_id: String::new(),
|
||
silent: false,
|
||
read_file_state,
|
||
sse_tx: None,
|
||
enable_thinking: false,
|
||
additional_allowed_dirs: Vec::new(),
|
||
tool_call_id: None,
|
||
max_output_chars: 4000,
|
||
}
|
||
}
|
||
|
||
/// 设置 SSE 通道
|
||
pub fn with_sse_tx(
|
||
mut self,
|
||
tx: tokio::sync::mpsc::UnboundedSender<crate::agent::runtime::AgentStreamEvent>,
|
||
) -> Self {
|
||
self.sse_tx = Some(tx);
|
||
self
|
||
}
|
||
|
||
/// 设置会话 ID
|
||
pub fn with_session_id(mut self, id: String) -> Self {
|
||
self.session_id = id;
|
||
self
|
||
}
|
||
|
||
/// 设置思考模式
|
||
pub fn with_thinking(mut self, enable: bool) -> Self {
|
||
self.enable_thinking = enable;
|
||
self
|
||
}
|
||
|
||
/// 设置附加允许目录
|
||
pub fn with_additional_dirs(mut self, dirs: Vec<String>) -> Self {
|
||
self.additional_allowed_dirs = dirs;
|
||
self
|
||
}
|
||
|
||
/// 设置工具调用 ID（用于工具的流式输出关联到自身）
|
||
pub fn with_tool_call_id(mut self, id: String) -> Self {
|
||
self.tool_call_id = Some(id);
|
||
self
|
||
}
|
||
|
||
/// 设置最大输出字符数
|
||
pub fn with_max_output_chars(mut self, max_chars: usize) -> Self {
|
||
self.max_output_chars = max_chars;
|
||
self
|
||
}
|
||
|
||
/// 创建静默上下文（子代理使用）
|
||
pub fn silent(app_state: Arc<AppState>) -> Self {
|
||
ToolContext {
|
||
app_state,
|
||
session_id: String::new(),
|
||
silent: true,
|
||
read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())),
|
||
sse_tx: None,
|
||
enable_thinking: false,
|
||
additional_allowed_dirs: Vec::new(),
|
||
tool_call_id: None,
|
||
max_output_chars: 4000,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 工具执行结果
|
||
#[derive(Debug, Clone)]
|
||
pub struct ToolOutput {
|
||
/// 给大模型阅读的截断文本
|
||
pub content: String,
|
||
/// 是否为错误
|
||
pub is_error: bool,
|
||
/// 结构化元数据（给前端 Timeline 直接渲染）
|
||
pub metadata: serde_json::Value,
|
||
/// 跳过持久化到磁盘（用于 read_file 等已从磁盘读取内容的工具，避免级联持久化）
|
||
pub skip_persist: bool,
|
||
}
|
||
|
||
impl ToolOutput {
|
||
/// 创建成功结果
|
||
pub fn success(content: impl Into<String>, metadata: serde_json::Value) -> Self {
|
||
ToolOutput {
|
||
content: content.into(),
|
||
is_error: false,
|
||
metadata,
|
||
skip_persist: false,
|
||
}
|
||
}
|
||
|
||
/// 创建成功结果，跳过持久化
|
||
pub fn success_skip_persist(content: impl Into<String>, metadata: serde_json::Value) -> Self {
|
||
ToolOutput {
|
||
content: content.into(),
|
||
is_error: false,
|
||
metadata,
|
||
skip_persist: true,
|
||
}
|
||
}
|
||
|
||
/// 创建错误结果
|
||
pub fn error(msg: impl Into<String>) -> Self {
|
||
ToolOutput {
|
||
content: msg.into(),
|
||
is_error: true,
|
||
metadata: json!({}),
|
||
skip_persist: false,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 工具被中断时的行为策略
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub enum InterruptBehavior {
|
||
/// 取消执行并返回错误（默认，适用于只读工具）
|
||
Cancel,
|
||
/// 阻塞中断信号直到执行完成（适用于有副作用的写入工具）
|
||
Block,
|
||
}
|
||
|
||
/// 权限规则来源
|
||
#[derive(Debug, Clone, PartialEq)]
|
||
pub enum PermissionRuleSource {
|
||
/// 环境变量加载（AGENT_PERMISSIONS_*）
|
||
Env,
|
||
/// 会话内动态添加（API / Always Allow）
|
||
Session,
|
||
}
|
||
|
||
impl std::fmt::Display for PermissionRuleSource {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
PermissionRuleSource::Env => write!(f, "env"),
|
||
PermissionRuleSource::Session => write!(f, "session"),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 权限规则 — 工具自定义的权限限制
|
||
#[derive(Debug, Clone)]
|
||
pub enum PermissionRule {
|
||
/// 不可覆盖的拒绝
|
||
Deny {
|
||
tool_name: String,
|
||
reason: String,
|
||
source: PermissionRuleSource,
|
||
},
|
||
/// 允许
|
||
Allow {
|
||
tool_name: String,
|
||
source: PermissionRuleSource,
|
||
},
|
||
/// 需要用户确认
|
||
Ask {
|
||
tool_name: String,
|
||
message: String,
|
||
source: PermissionRuleSource,
|
||
},
|
||
}
|
||
|
||
/// 智能体工具 trait
|
||
#[async_trait]
|
||
pub trait AgentTool: Send + Sync {
|
||
/// 工具名称（与 LLM function calling 的 name 保持一致）
|
||
fn name(&self) -> &str;
|
||
/// 工具描述（告知 LLM 何时应该调用该工具）
|
||
fn description(&self) -> &str;
|
||
/// JSON Schema 格式的参数定义
|
||
fn parameters(&self) -> serde_json::Value;
|
||
/// 执行工具逻辑
|
||
async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput;
|
||
|
||
// ── P1 优化：新增默认方法 ──
|
||
|
||
/// 中断行为策略。默认 Cancel — 可以安全中断。
|
||
fn interrupt_behavior(&self) -> InterruptBehavior {
|
||
InterruptBehavior::Cancel
|
||
}
|
||
|
||
/// 该工具是否支持并发安全执行。
|
||
/// 默认 false（保守策略），只读工具应覆写为 true。
|
||
fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
|
||
false
|
||
}
|
||
|
||
/// 工具自定义权限检查。默认无额外限制。
|
||
fn check_permissions(&self, _args: &serde_json::Value) -> Vec<PermissionRule> {
|
||
Vec::new()
|
||
}
|
||
|
||
/// 该工具错误时是否应中止兄弟并行执行。
|
||
/// 默认 false（只读工具不触发）。下载/解析类工具可覆写为 true。
|
||
fn causes_sibling_abort(&self) -> bool {
|
||
false
|
||
}
|
||
|
||
/// 带进度流式执行。默认委托给 execute()。
|
||
/// 长时间操作的工具可覆写以发送进度更新。
|
||
async fn execute_with_progress(
|
||
&self,
|
||
args: serde_json::Value,
|
||
ctx: &ToolContext,
|
||
_progress_tx: Option<&tokio::sync::mpsc::UnboundedSender<String>>,
|
||
) -> ToolOutput {
|
||
self.execute(args, ctx).await
|
||
}
|
||
|
||
// ── P3 优化：延迟加载与分类器支持 ──
|
||
|
||
/// 工具简要分类描述（~15 词），供 LLM 判断是否需要加载延迟工具。
|
||
/// 默认取 description 的第一句，截断到 100 字符。
|
||
fn classifier_summary(&self) -> String {
|
||
self.description()
|
||
.split('.')
|
||
.next()
|
||
.unwrap_or("")
|
||
.chars()
|
||
.take(100)
|
||
.collect()
|
||
}
|
||
|
||
/// 该工具是否为只读（无副作用）。auto-mode 分类器和权限系统可用此信息。
|
||
/// 默认 false（保守——必须显式选择加入）。
|
||
fn is_readonly(&self) -> bool {
|
||
false
|
||
}
|
||
|
||
/// 该工具是否应延迟加载（不在初始 tool definitions 中，通过 tool_search 可发现）。
|
||
/// 默认 false（常驻——始终在提示词中）。
|
||
fn defer_loading(&self) -> bool {
|
||
false
|
||
}
|
||
|
||
/// 工具所属分组。用于按功能域组织工具目录和模式过滤。
|
||
/// 默认返回 "general"，天文科研工具应覆写为 "system" 或 "research"。
|
||
fn group(&self) -> &str {
|
||
"general"
|
||
}
|
||
}
|
||
|
||
/// 工具注册表，管理所有可用工具。
|
||
/// 内部使用 HashMap 实现 O(1) 按名查找，同时保留插入顺序供 definitions() 使用。
|
||
///
|
||
/// `definition_filter` 用于子代理最小权限：设置后 `definitions()` 仅返回白名单工具，
|
||
/// 但 `get()` 仍可查找所有工具（以便对未授权调用返回友好错误消息）。
|
||
pub struct ToolRegistry {
|
||
tools: std::collections::HashMap<String, Box<dyn AgentTool>>,
|
||
ordered_names: Vec<String>,
|
||
/// 可选的工具白名单（子代理最小权限）
|
||
definition_filter: Option<std::collections::HashSet<String>>,
|
||
/// 缓存的工具 schema（session 生命周期内有效，工具注册变更时失效）
|
||
schema_cache: Option<Vec<ToolDefinition>>,
|
||
/// schema 缓存版本号（递增使缓存失效）
|
||
schema_generation: u64,
|
||
}
|
||
|
||
// ── 工具注册辅助函数（消除重复代码） ──
|
||
|
||
/// 注册所有基础研究工具（文件、文献、RAG、笔记等工具）。
|
||
fn add_base_tools(registry: &mut ToolRegistry, skill_registry: Arc<RwLock<SkillRegistry>>) {
|
||
let tools: Vec<Box<dyn AgentTool>> = vec![
|
||
Box::new(ReadFileTool),
|
||
Box::new(GrepFilesTool),
|
||
Box::new(GlobFilesTool),
|
||
Box::new(RunBashTool),
|
||
Box::new(FileWriteTool),
|
||
Box::new(FileEditTool),
|
||
// 系统级
|
||
Box::new(SearchPapersTool),
|
||
Box::new(ProcessPaperTool),
|
||
// 研究级
|
||
Box::new(SearchLocalLibraryTool),
|
||
Box::new(GetPaperMetadataTool),
|
||
Box::new(GetPaperOutlineTool),
|
||
Box::new(GetPaperContentTool),
|
||
Box::new(GetCitationNetworkTool),
|
||
Box::new(RagSearchTool),
|
||
Box::new(QueryTargetTool),
|
||
Box::new(SaveNoteTool),
|
||
Box::new(TodoWriteTool),
|
||
Box::new(CompressTool),
|
||
Box::new(AskUserTool),
|
||
Box::new(LoadSkillTool::new(skill_registry)),
|
||
Box::new(SubAgentTool::new()),
|
||
Box::new(SearchHistoryTool),
|
||
];
|
||
for tool in tools {
|
||
registry.ordered_names.push(tool.name().to_string());
|
||
registry.tools.insert(tool.name().to_string(), tool);
|
||
}
|
||
}
|
||
|
||
/// 注册后台任务工具（bg_task_run, bg_task_check）。
|
||
fn add_background_tools(
|
||
registry: &mut ToolRegistry,
|
||
queue: Arc<crate::agent::background::BgNotificationQueue>,
|
||
) {
|
||
let run_tool = Box::new(BgTaskRunTool::new(queue.clone()));
|
||
let check_tool = Box::new(BgTaskCheckTool::new(queue));
|
||
registry.ordered_names.push(run_tool.name().to_string());
|
||
registry.tools.insert(run_tool.name().to_string(), run_tool);
|
||
registry.ordered_names.push(check_tool.name().to_string());
|
||
registry
|
||
.tools
|
||
.insert(check_tool.name().to_string(), check_tool);
|
||
}
|
||
|
||
/// 注册团队协作工具（spawn_teammate, send_teammate_message, team_broadcast, check_team_inbox）。
|
||
fn add_team_tools(
|
||
registry: &mut ToolRegistry,
|
||
team_manager: Arc<tokio::sync::Mutex<Option<crate::agent::team::manager::TeamManager>>>,
|
||
) {
|
||
let team_tools: Vec<Box<dyn AgentTool>> = vec![
|
||
Box::new(SpawnTeammateTool::new(team_manager.clone())),
|
||
Box::new(SendTeammateMessageTool::new(team_manager.clone())),
|
||
Box::new(TeamBroadcastTool::new(team_manager.clone())),
|
||
Box::new(CheckTeamInboxTool::new(team_manager)),
|
||
];
|
||
for tool in team_tools {
|
||
registry.ordered_names.push(tool.name().to_string());
|
||
registry.tools.insert(tool.name().to_string(), tool);
|
||
}
|
||
}
|
||
|
||
impl ToolRegistry {
|
||
/// 创建空工具注册表（调用者通过 add_tool 手动添加工具）。
|
||
/// 用于受限场景（如记忆提取子代理只需要只读 + save_memory）。
|
||
pub fn empty() -> Self {
|
||
ToolRegistry {
|
||
tools: std::collections::HashMap::new(),
|
||
ordered_names: Vec::new(),
|
||
definition_filter: None,
|
||
schema_cache: None,
|
||
schema_generation: 0,
|
||
}
|
||
}
|
||
|
||
/// 创建默认工具注册表（包含全部科研工具，不含后台工具）
|
||
pub fn new(skill_registry: Arc<RwLock<SkillRegistry>>) -> Self {
|
||
Self::new_with_queue(None, skill_registry)
|
||
}
|
||
|
||
/// 创建工具注册表，可选注入后台通知队列以启用 bg_task_run/bg_task_check
|
||
pub fn new_with_queue(
|
||
queue: Option<Arc<crate::agent::background::BgNotificationQueue>>,
|
||
skill_registry: Arc<RwLock<SkillRegistry>>,
|
||
) -> Self {
|
||
let mut registry = ToolRegistry {
|
||
tools: std::collections::HashMap::new(),
|
||
ordered_names: Vec::new(),
|
||
definition_filter: None,
|
||
schema_cache: None,
|
||
schema_generation: 0,
|
||
};
|
||
add_base_tools(&mut registry, skill_registry);
|
||
if let Some(q) = queue {
|
||
add_background_tools(&mut registry, q);
|
||
}
|
||
registry.precompute_definitions();
|
||
registry
|
||
}
|
||
|
||
/// 创建包含团队工具的注册表
|
||
pub fn new_with_team(
|
||
queue: Option<Arc<crate::agent::background::BgNotificationQueue>>,
|
||
team_manager: Arc<tokio::sync::Mutex<Option<crate::agent::team::manager::TeamManager>>>,
|
||
skill_registry: Arc<RwLock<SkillRegistry>>,
|
||
) -> Self {
|
||
// 使用 new_with_queue 获取基础 + 后台工具，再添加团队工具
|
||
let mut registry = Self::new_with_queue(queue, skill_registry);
|
||
add_team_tools(&mut registry, team_manager);
|
||
registry.precompute_definitions();
|
||
registry
|
||
}
|
||
|
||
/// 动态添加工具（用于需要共享状态的工具，如 MemoryManager）
|
||
pub fn add_tool(&mut self, tool: Box<dyn AgentTool>) {
|
||
let name = tool.name().to_string();
|
||
self.ordered_names.push(name.clone());
|
||
self.tools.insert(name, tool);
|
||
// 工具变更 → 使 schema 缓存失效
|
||
self.schema_cache = None;
|
||
self.schema_generation += 1;
|
||
}
|
||
|
||
/// 替换已存在的工具（保持名称在 ordered_names 中的位置不变）。
|
||
/// 如果工具不存在，行为等同于 add_tool。
|
||
pub fn replace_tool(&mut self, tool: Box<dyn AgentTool>) {
|
||
let name = tool.name().to_string();
|
||
if !self.tools.contains_key(&name) {
|
||
self.ordered_names.push(name.clone());
|
||
}
|
||
self.tools.insert(name, tool);
|
||
// 工具变更 → 使 schema 缓存失效
|
||
self.schema_cache = None;
|
||
self.schema_generation += 1;
|
||
}
|
||
|
||
/// 根据名称查找工具 (O(1))
|
||
pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
|
||
self.tools.get(name).map(|t| t.as_ref())
|
||
}
|
||
|
||
/// 设置工具定义白名单过滤器（子代理最小权限）。
|
||
/// 设置后 `definitions()` 仅返回白名单中的工具。
|
||
pub fn set_definition_filter(&mut self, allowed: Vec<String>) {
|
||
self.definition_filter = Some(allowed.into_iter().collect());
|
||
// 过滤器变更 → 使 schema 缓存失效
|
||
self.schema_cache = None;
|
||
self.schema_generation += 1;
|
||
}
|
||
|
||
/// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM）。
|
||
/// 按名称字母序排序以保证跨调用的稳定性，提升 prompt cache 命中率。
|
||
/// 若设置了 definition_filter，仅返回白名单中的工具。
|
||
///
|
||
/// 使用内部 schema 缓存：工具注册表不变时复用上次计算结果。
|
||
pub fn definitions(&self) -> Vec<ToolDefinition> {
|
||
// 缓存命中：直接返回
|
||
if let Some(ref cached) = self.schema_cache {
|
||
return cached.clone();
|
||
}
|
||
|
||
// 缓存未命中：重新计算
|
||
let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
|
||
self.tools
|
||
.iter()
|
||
.filter(|(name, _)| filter.contains(*name))
|
||
.map(|(_, tool)| tool)
|
||
.collect()
|
||
} else {
|
||
self.tools.values().collect()
|
||
};
|
||
let mut defs: Vec<_> = values
|
||
.iter()
|
||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||
.collect();
|
||
defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
|
||
|
||
// 缓存命中逻辑在上方 precompute_definitions() 中（&mut self）预填充；
|
||
// definitions() 持 &self 不可写入，仅读取缓存或当场计算返回。
|
||
// 调用方应在初始化时调用 precompute_definitions() 确保缓存已预热。
|
||
defs
|
||
}
|
||
|
||
/// 获取当前 schema 缓存版本号（用于外部检测工具是否变更）。
|
||
pub fn schema_generation(&self) -> u64 {
|
||
self.schema_generation
|
||
}
|
||
|
||
/// 预计算并缓存 tool definitions（在 AgentRuntime 初始化时调用，
|
||
/// 此时拥有 &mut self，可以安全写入缓存）。
|
||
pub fn precompute_definitions(&mut self) {
|
||
if self.schema_cache.is_some() {
|
||
return;
|
||
}
|
||
let values: Vec<&Box<dyn AgentTool>> = if let Some(filter) = &self.definition_filter {
|
||
self.tools
|
||
.iter()
|
||
.filter(|(name, _)| filter.contains(*name))
|
||
.map(|(_, tool)| tool)
|
||
.collect()
|
||
} else {
|
||
self.tools.values().collect()
|
||
};
|
||
let mut defs: Vec<_> = values
|
||
.iter()
|
||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||
.collect();
|
||
defs.sort_by(|a, b| a.function.name.cmp(&b.function.name));
|
||
self.schema_cache = Some(defs);
|
||
}
|
||
|
||
/// 设置工具白名单（子代理最小权限）。
|
||
/// 设置后 `definitions()` 仅暴露白名单中的工具给 LLM，
|
||
/// 但 `get()` 仍可访问所有工具以便对未授权调用返回友好错误消息。
|
||
pub fn with_filter(mut self, allowed: &[&str]) -> Self {
|
||
self.definition_filter = Some(allowed.iter().map(|s| s.to_string()).collect());
|
||
self
|
||
}
|
||
|
||
/// 返回只常驻工具定义（defer_loading() == false）。
|
||
/// 这是每次 LLM 调用的主工具列表。
|
||
pub fn resident_definitions(&self) -> Vec<ToolDefinition> {
|
||
self.tools
|
||
.values()
|
||
.filter(|t| !t.defer_loading())
|
||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||
.collect()
|
||
}
|
||
|
||
/// 返回延迟加载工具定义（defer_loading() == true）。
|
||
pub fn deferred_definitions(&self) -> Vec<ToolDefinition> {
|
||
self.tools
|
||
.values()
|
||
.filter(|t| t.defer_loading())
|
||
.map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters()))
|
||
.collect()
|
||
}
|
||
|
||
/// 构建紧凑的"可用工具目录"字符串，供系统提示词使用。
|
||
/// 格式: "tool_name: classifier_summary"，按名称排序，延迟工具加 [deferred] 前缀。
|
||
pub fn tool_catalog(&self) -> String {
|
||
let mut entries: Vec<String> = self
|
||
.tools
|
||
.values()
|
||
.map(|t| {
|
||
let prefix = if t.defer_loading() { "[deferred]" } else { "" };
|
||
format!("{}{}: {}", prefix, t.name(), t.classifier_summary())
|
||
})
|
||
.collect();
|
||
entries.sort();
|
||
entries.join("\n")
|
||
}
|
||
|
||
/// 返回当前所有工具名称列表
|
||
pub fn tool_names(&self) -> Vec<String> {
|
||
self.ordered_names.clone()
|
||
}
|
||
|
||
/// 按分组名返回工具名称列表
|
||
pub fn tools_by_group(&self, group: &str) -> Vec<String> {
|
||
self.tools
|
||
.values()
|
||
.filter(|t| t.group() == group)
|
||
.map(|t| t.name().to_string())
|
||
.collect()
|
||
}
|
||
|
||
/// 构建按 group 分组的紧凑工具目录字符串。
|
||
/// 格式:
|
||
/// [system]
|
||
/// search_papers: ...
|
||
/// [research]
|
||
/// get_paper_metadata: ...
|
||
pub fn tool_catalog_grouped(&self) -> String {
|
||
let mut groups: std::collections::BTreeMap<&str, Vec<String>> =
|
||
std::collections::BTreeMap::new();
|
||
for tool in self.tools.values() {
|
||
let g = tool.group();
|
||
let prefix = if tool.defer_loading() {
|
||
"[deferred]"
|
||
} else {
|
||
""
|
||
};
|
||
groups.entry(g).or_default().push(format!(
|
||
"{}{}: {}",
|
||
prefix,
|
||
tool.name(),
|
||
tool.classifier_summary()
|
||
));
|
||
}
|
||
let mut out = String::new();
|
||
for (group, mut entries) in groups {
|
||
entries.sort();
|
||
out.push_str(&format!("[{}]\n", group));
|
||
for e in entries {
|
||
out.push_str(&format!("{}\n", e));
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
/// 截断文本到指定最大字符数
|
||
pub fn truncate_content(s: &str, max_chars: usize) -> String {
|
||
let char_count = s.chars().count();
|
||
if char_count <= max_chars {
|
||
s.to_string()
|
||
} else {
|
||
let truncated: String = s.chars().take(max_chars).collect();
|
||
format!(
|
||
"{}\n\n[... 内容已截断，共 {} 字符 ...]",
|
||
truncated, char_count
|
||
)
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::path::PathBuf;
|
||
|
||
#[test]
|
||
fn test_truncate_content_short() {
|
||
let text = "Hello, world!";
|
||
assert_eq!(truncate_content(text, 100), text);
|
||
}
|
||
|
||
#[test]
|
||
fn test_truncate_content_long() {
|
||
let text = "a".repeat(5000);
|
||
let result = truncate_content(&text, 100);
|
||
assert!(result.contains("内容已截断"));
|
||
assert!(result.contains("5000"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_tool_output_success() {
|
||
let output = ToolOutput::success("ok", json!({"key": "value"}));
|
||
assert!(!output.is_error);
|
||
assert_eq!(output.content, "ok");
|
||
}
|
||
|
||
#[test]
|
||
fn test_tool_output_error() {
|
||
let output = ToolOutput::error("something went wrong");
|
||
assert!(output.is_error);
|
||
}
|
||
|
||
#[test]
|
||
fn test_tool_registry_definitions() {
|
||
let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
|
||
"./skills",
|
||
)))));
|
||
let defs = registry.definitions();
|
||
assert_eq!(defs.len(), 22);
|
||
assert!(defs.iter().any(|d| d.function.name == "read_file"));
|
||
assert!(defs.iter().any(|d| d.function.name == "grep_files"));
|
||
assert!(defs.iter().any(|d| d.function.name == "glob_files"));
|
||
assert!(defs.iter().any(|d| d.function.name == "run_bash"));
|
||
assert!(defs.iter().any(|d| d.function.name == "file_write"));
|
||
assert!(defs.iter().any(|d| d.function.name == "file_edit"));
|
||
assert!(defs.iter().any(|d| d.function.name == "search_papers"));
|
||
assert!(defs.iter().any(|d| d.function.name == "process_paper"));
|
||
assert!(defs
|
||
.iter()
|
||
.any(|d| d.function.name == "search_local_library"));
|
||
assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata"));
|
||
assert!(defs.iter().any(|d| d.function.name == "get_paper_outline"));
|
||
assert!(defs.iter().any(|d| d.function.name == "get_paper_content"));
|
||
assert!(defs
|
||
.iter()
|
||
.any(|d| d.function.name == "get_citation_network"));
|
||
assert!(defs.iter().any(|d| d.function.name == "rag_search"));
|
||
assert!(defs.iter().any(|d| d.function.name == "query_target"));
|
||
assert!(defs.iter().any(|d| d.function.name == "save_note"));
|
||
assert!(defs.iter().any(|d| d.function.name == "todo_write"));
|
||
assert!(defs.iter().any(|d| d.function.name == "compress_context"));
|
||
assert!(defs.iter().any(|d| d.function.name == "load_skill"));
|
||
assert!(defs.iter().any(|d| d.function.name == "subagent"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_tool_registry_get() {
|
||
let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from(
|
||
"./skills",
|
||
)))));
|
||
assert!(registry.get("search_papers").is_some());
|
||
assert!(registry.get("nonexistent").is_none());
|
||
}
|
||
}
|