Compare commits
3 Commits
bafa7a606c
...
a479b92cdf
| Author | SHA1 | Date | |
|---|---|---|---|
| a479b92cdf | |||
| af07eaf820 | |||
| c81b1e42c7 |
@ -37,3 +37,4 @@ chrono = "0.4"
|
|||||||
hostname = "0.3"
|
hostname = "0.3"
|
||||||
sqlx = { version = "0.8", features = ["sqlite", "macros", "chrono", "runtime-tokio"] }
|
sqlx = { version = "0.8", features = ["sqlite", "macros", "chrono", "runtime-tokio"] }
|
||||||
jieba-rs = "0.9"
|
jieba-rs = "0.9"
|
||||||
|
which = "7"
|
||||||
|
|||||||
19
resources/templates/AGENTS.md
Normal file
19
resources/templates/AGENTS.md
Normal file
@ -0,0 +1,19 @@
|
|||||||
|
# Agent Instructions
|
||||||
|
|
||||||
|
You are PicoBot, a personal AI assistant.
|
||||||
|
|
||||||
|
## Personality
|
||||||
|
- Helpful and friendly
|
||||||
|
- Concise and to the point
|
||||||
|
- Proactive when useful, respects user boundaries
|
||||||
|
|
||||||
|
## Values
|
||||||
|
- Accuracy over speed
|
||||||
|
- User privacy and safety
|
||||||
|
- Transparency in actions
|
||||||
|
|
||||||
|
## Communication Style
|
||||||
|
- Be clear and direct
|
||||||
|
- Use Chinese or English based on the user's language
|
||||||
|
- Explain reasoning when helpful
|
||||||
|
- Ask clarifying questions when needed
|
||||||
31
resources/templates/USER.md
Normal file
31
resources/templates/USER.md
Normal file
@ -0,0 +1,31 @@
|
|||||||
|
# 用户配置
|
||||||
|
|
||||||
|
PicoBot 会根据此文件了解你的偏好。
|
||||||
|
|
||||||
|
## 基本信息
|
||||||
|
|
||||||
|
- **称呼**: 用户
|
||||||
|
- **时区**: Asia/Shanghai (UTC+8)
|
||||||
|
- **语言**: 中文
|
||||||
|
|
||||||
|
## 偏好设置
|
||||||
|
|
||||||
|
### 回复风格
|
||||||
|
- [ ] 简洁扼要
|
||||||
|
- [ ] 详细解释
|
||||||
|
- [ ] 根据问题自适应
|
||||||
|
|
||||||
|
### 沟通风格
|
||||||
|
- [ ] 随意
|
||||||
|
- [ ] 专业
|
||||||
|
- [ ] 技术导向
|
||||||
|
|
||||||
|
## 工作环境
|
||||||
|
|
||||||
|
- **主要角色**: 开发者
|
||||||
|
- **当前项目**:
|
||||||
|
- **常用工具**:
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
*编辑此文件来定制 PicoBot 的行为偏好。*
|
||||||
@ -605,7 +605,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn test_threshold() {
|
fn test_threshold() {
|
||||||
let compressor = ContextCompressor::new(mock_provider(), 128_000, test_memory_manager());
|
let compressor = ContextCompressor::new(mock_provider(), 128_000, test_memory_manager());
|
||||||
assert_eq!(compressor.threshold(), 64_000);
|
assert_eq!(compressor.threshold(), 89_600);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
|
|||||||
@ -3,8 +3,11 @@
|
|||||||
//! This module provides a modular framework for building system prompts
|
//! This module provides a modular framework for building system prompts
|
||||||
//! using the SystemPromptBuilder pattern.
|
//! using the SystemPromptBuilder pattern.
|
||||||
//!
|
//!
|
||||||
//! Configuration:
|
//! Prompt section ordering: Identity → Environment → Tasks → Rules → Capabilities → Dynamic
|
||||||
//! - USER.md is loaded from ~/.picobot/USER.md (user's personal configuration)
|
//!
|
||||||
|
//! Configuration files loaded from ~/.picobot/:
|
||||||
|
//! - AGENTS.md — agent identity and behavior
|
||||||
|
//! - USER.md — user preferences and profile
|
||||||
|
|
||||||
use crate::tools::ToolRegistry;
|
use crate::tools::ToolRegistry;
|
||||||
use std::fmt::Write;
|
use std::fmt::Write;
|
||||||
@ -42,16 +45,17 @@ impl SystemPromptBuilder {
|
|||||||
pub fn with_defaults() -> Self {
|
pub fn with_defaults() -> Self {
|
||||||
Self {
|
Self {
|
||||||
sections: vec![
|
sections: vec![
|
||||||
Box::new(ToolHonestySection),
|
Box::new(AgentProfileSection),
|
||||||
Box::new(YourTaskSection),
|
|
||||||
Box::new(SafetySection),
|
|
||||||
Box::new(WorkspaceSection),
|
|
||||||
Box::new(UserProfileSection),
|
Box::new(UserProfileSection),
|
||||||
|
Box::new(RuntimeSection),
|
||||||
|
Box::new(DateTimeSection),
|
||||||
|
Box::new(WorkspaceSection),
|
||||||
|
Box::new(YourTaskSection),
|
||||||
|
Box::new(ToolHonestySection),
|
||||||
|
Box::new(SafetySection),
|
||||||
|
Box::new(CrossChannelSection),
|
||||||
Box::new(MemorySection),
|
Box::new(MemorySection),
|
||||||
Box::new(HistorySection),
|
Box::new(HistorySection),
|
||||||
Box::new(DateTimeSection),
|
|
||||||
Box::new(RuntimeSection),
|
|
||||||
Box::new(CrossChannelSection),
|
|
||||||
],
|
],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -202,6 +206,29 @@ impl PromptSection for UserProfileSection {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Agent profile from ~/.picobot/AGENTS.md.
|
||||||
|
pub struct AgentProfileSection;
|
||||||
|
|
||||||
|
impl PromptSection for AgentProfileSection {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"agent_profile"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build(&self, _ctx: &PromptContext<'_>) -> String {
|
||||||
|
let mut output = String::from("## Agent 配置\n\n");
|
||||||
|
|
||||||
|
if let Some(user_config_dir) = get_user_config_dir()
|
||||||
|
&& let Some(content) =
|
||||||
|
load_file_from_dir(&user_config_dir, "AGENTS.md", BOOTSTRAP_MAX_CHARS)
|
||||||
|
{
|
||||||
|
output.push_str(&content);
|
||||||
|
return output;
|
||||||
|
}
|
||||||
|
|
||||||
|
String::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Current date and time.
|
/// Current date and time.
|
||||||
pub struct DateTimeSection;
|
pub struct DateTimeSection;
|
||||||
|
|
||||||
|
|||||||
@ -36,6 +36,9 @@ impl GatewayState {
|
|||||||
|
|
||||||
tracing::info!("Using workspace directory: {}", workspace_path.display());
|
tracing::info!("Using workspace directory: {}", workspace_path.display());
|
||||||
|
|
||||||
|
// Release default AGENTS.md and USER.md to ~/.picobot/ if not exist
|
||||||
|
ensure_default_config_files();
|
||||||
|
|
||||||
// Get provider config for SessionManager
|
// Get provider config for SessionManager
|
||||||
let mut provider_config = config.get_provider_config("default")?;
|
let mut provider_config = config.get_provider_config("default")?;
|
||||||
// Override workspace_dir with the ensured path
|
// Override workspace_dir with the ensured path
|
||||||
@ -347,3 +350,32 @@ pub async fn run(host: Option<String>, port: Option<u16>) -> Result<(), Box<dyn
|
|||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Release default AGENTS.md and USER.md templates to ~/.picobot/ if not already present.
|
||||||
|
fn ensure_default_config_files() {
|
||||||
|
let picobot_dir = dirs::home_dir().unwrap_or_default().join(".picobot");
|
||||||
|
if let Err(e) = std::fs::create_dir_all(&picobot_dir) {
|
||||||
|
tracing::warn!(dir = %picobot_dir.display(), error = %e, "Failed to create ~/.picobot directory");
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
let agents_path = picobot_dir.join("AGENTS.md");
|
||||||
|
if !agents_path.exists() {
|
||||||
|
let content = include_str!("../../resources/templates/AGENTS.md");
|
||||||
|
if let Err(e) = std::fs::write(&agents_path, content) {
|
||||||
|
tracing::warn!(path = %agents_path.display(), error = %e, "Failed to write AGENTS.md template");
|
||||||
|
} else {
|
||||||
|
tracing::info!(path = %agents_path.display(), "Released default AGENTS.md template");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let user_path = picobot_dir.join("USER.md");
|
||||||
|
if !user_path.exists() {
|
||||||
|
let content = include_str!("../../resources/templates/USER.md");
|
||||||
|
if let Err(e) = std::fs::write(&user_path, content) {
|
||||||
|
tracing::warn!(path = %user_path.display(), error = %e, "Failed to write USER.md template");
|
||||||
|
} else {
|
||||||
|
tracing::info!(path = %user_path.display(), "Released default USER.md template");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@ -487,7 +487,7 @@ impl Session {
|
|||||||
if skills_prompt.trim().is_empty() {
|
if skills_prompt.trim().is_empty() {
|
||||||
base_prompt
|
base_prompt
|
||||||
} else {
|
} else {
|
||||||
format!("{}\n\n## Skills\n\n{}\n\nUse the `get_skill` tool to load a skill's full content when needed.", base_prompt, skills_prompt)
|
format!("{}\n\n{}", base_prompt, skills_prompt)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -810,8 +810,9 @@ impl SessionManager {
|
|||||||
bus: Arc<MessageBus>,
|
bus: Arc<MessageBus>,
|
||||||
memory_manager: Arc<crate::memory::MemoryManager>,
|
memory_manager: Arc<crate::memory::MemoryManager>,
|
||||||
) -> Result<Self, AgentError> {
|
) -> Result<Self, AgentError> {
|
||||||
let skills_loader = SkillsLoader::new();
|
let mut skills_loader = SkillsLoader::new();
|
||||||
skills_loader.load_skills();
|
skills_loader.load_skills();
|
||||||
|
skills_loader.set_workspace_skills_dir(provider_config.workspace_dir.clone());
|
||||||
let skills_loader = Arc::new(skills_loader);
|
let skills_loader = Arc::new(skills_loader);
|
||||||
|
|
||||||
let tools = Arc::new(create_default_tools(skills_loader.clone(), memory_manager.clone()));
|
let tools = Arc::new(create_default_tools(skills_loader.clone(), memory_manager.clone()));
|
||||||
|
|||||||
@ -24,6 +24,7 @@ struct SkillsState {
|
|||||||
loaded_skills: Vec<Skill>,
|
loaded_skills: Vec<Skill>,
|
||||||
last_picobot_mtime: Option<SystemTime>,
|
last_picobot_mtime: Option<SystemTime>,
|
||||||
last_agent_mtime: Option<SystemTime>,
|
last_agent_mtime: Option<SystemTime>,
|
||||||
|
last_workspace_mtime: Option<SystemTime>,
|
||||||
last_load_time: SystemTime,
|
last_load_time: SystemTime,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -33,6 +34,7 @@ impl Default for SkillsState {
|
|||||||
loaded_skills: Vec::new(),
|
loaded_skills: Vec::new(),
|
||||||
last_picobot_mtime: None,
|
last_picobot_mtime: None,
|
||||||
last_agent_mtime: None,
|
last_agent_mtime: None,
|
||||||
|
last_workspace_mtime: None,
|
||||||
last_load_time: SystemTime::now(),
|
last_load_time: SystemTime::now(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -43,6 +45,7 @@ impl Default for SkillsState {
|
|||||||
pub struct SkillsLoader {
|
pub struct SkillsLoader {
|
||||||
picobot_skills_dir: PathBuf,
|
picobot_skills_dir: PathBuf,
|
||||||
agent_skills_dir: PathBuf,
|
agent_skills_dir: PathBuf,
|
||||||
|
workspace_skills_dir: Option<PathBuf>,
|
||||||
state: Arc<Mutex<SkillsState>>,
|
state: Arc<Mutex<SkillsState>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -53,6 +56,7 @@ impl SkillsLoader {
|
|||||||
Self {
|
Self {
|
||||||
picobot_skills_dir: home.join(".picobot/skills"),
|
picobot_skills_dir: home.join(".picobot/skills"),
|
||||||
agent_skills_dir: home.join(".agent/skills"),
|
agent_skills_dir: home.join(".agent/skills"),
|
||||||
|
workspace_skills_dir: None,
|
||||||
state: Arc::new(Mutex::new(SkillsState::default())),
|
state: Arc::new(Mutex::new(SkillsState::default())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@ -62,10 +66,16 @@ impl SkillsLoader {
|
|||||||
Self {
|
Self {
|
||||||
picobot_skills_dir: picobot_dir,
|
picobot_skills_dir: picobot_dir,
|
||||||
agent_skills_dir: agent_dir,
|
agent_skills_dir: agent_dir,
|
||||||
|
workspace_skills_dir: None,
|
||||||
state: Arc::new(Mutex::new(SkillsState::default())),
|
state: Arc::new(Mutex::new(SkillsState::default())),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Set the workspace skills directory (./skills under workspace root)
|
||||||
|
pub fn set_workspace_skills_dir(&mut self, workspace_path: PathBuf) {
|
||||||
|
self.workspace_skills_dir = Some(workspace_path.join("skills"));
|
||||||
|
}
|
||||||
|
|
||||||
/// Load all skills from both directories and record modification times
|
/// Load all skills from both directories and record modification times
|
||||||
pub fn load_skills(&self) {
|
pub fn load_skills(&self) {
|
||||||
let mut state = self.state.lock().unwrap();
|
let mut state = self.state.lock().unwrap();
|
||||||
@ -104,6 +114,20 @@ impl SkillsLoader {
|
|||||||
state.last_agent_mtime = Self::get_dir_mtime(&self.agent_skills_dir);
|
state.last_agent_mtime = Self::get_dir_mtime(&self.agent_skills_dir);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Load from workspace ./skills (if set)
|
||||||
|
if let Some(ref ws_dir) = self.workspace_skills_dir {
|
||||||
|
if ws_dir.exists() {
|
||||||
|
let loaded = self.load_skills_from_dir(ws_dir);
|
||||||
|
tracing::debug!(
|
||||||
|
dir = %ws_dir.display(),
|
||||||
|
count = loaded.len(),
|
||||||
|
"Loaded skills from workspace directory"
|
||||||
|
);
|
||||||
|
state.loaded_skills.extend(loaded);
|
||||||
|
state.last_workspace_mtime = Self::get_dir_mtime(ws_dir);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
state.last_load_time = SystemTime::now();
|
state.last_load_time = SystemTime::now();
|
||||||
|
|
||||||
if state.loaded_skills.is_empty() {
|
if state.loaded_skills.is_empty() {
|
||||||
@ -130,7 +154,18 @@ impl SkillsLoader {
|
|||||||
false
|
false
|
||||||
};
|
};
|
||||||
|
|
||||||
picobot_changed || agent_changed
|
let workspace_changed = if let Some(ref ws_dir) = self.workspace_skills_dir {
|
||||||
|
if ws_dir.exists() {
|
||||||
|
let current_mtime = Self::get_dir_mtime(ws_dir);
|
||||||
|
current_mtime != state.last_workspace_mtime
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
false
|
||||||
|
};
|
||||||
|
|
||||||
|
picobot_changed || agent_changed || workspace_changed
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Reload skills if changes are detected
|
/// Reload skills if changes are detected
|
||||||
@ -247,46 +282,53 @@ impl SkillsLoader {
|
|||||||
parts.join("\n\n---\n\n")
|
parts.join("\n\n---\n\n")
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build full skills prompt combining always skills and summary (checks for changes first)
|
/// Build full skills prompt: directory conventions, always-skill summary, always-skill content
|
||||||
pub fn build_skills_prompt(&self) -> String {
|
pub fn build_skills_prompt(&self) -> String {
|
||||||
self.reload_if_changed();
|
self.reload_if_changed();
|
||||||
let state = self.state.lock().unwrap();
|
let state = self.state.lock().unwrap();
|
||||||
|
|
||||||
let mut prompt = String::new();
|
if state.loaded_skills.is_empty() {
|
||||||
|
return String::new();
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut prompt = String::from("## Skills\n\n");
|
||||||
|
|
||||||
|
// Directory conventions
|
||||||
|
prompt.push_str("### 目录说明\n\n");
|
||||||
|
prompt.push_str("- `~/.agent/skills/` — 外部共享 skill 目录(第三方、系统级 skill)\n");
|
||||||
|
prompt.push_str("- `~/.picobot/skills/` — 安装 skill 的默认目录\n");
|
||||||
|
prompt.push_str("- `./skills/` — 工作目录下的 skill,picobot 自行创建的 skill 存放于此\n\n");
|
||||||
|
prompt.push_str("安装或创建 skill 时请按上述目录规范存放。\n\n");
|
||||||
|
|
||||||
|
// Always skills summary
|
||||||
let always_skills: Vec<_> = state.loaded_skills.iter().filter(|s| s.always).collect();
|
let always_skills: Vec<_> = state.loaded_skills.iter().filter(|s| s.always).collect();
|
||||||
if !always_skills.is_empty() {
|
if !always_skills.is_empty() {
|
||||||
|
prompt.push_str("### 常用技能\n\n");
|
||||||
|
for skill in &always_skills {
|
||||||
|
let path_str = skill.path.as_ref()
|
||||||
|
.map(|p| p.to_string_lossy().to_string())
|
||||||
|
.unwrap_or_else(|| "—".to_string());
|
||||||
|
prompt.push_str(&format!(
|
||||||
|
"- **{}**: {} [路径: `{}`]\n",
|
||||||
|
skill.name, skill.description, path_str
|
||||||
|
));
|
||||||
|
}
|
||||||
|
prompt.push('\n');
|
||||||
|
}
|
||||||
|
|
||||||
|
// Usage instructions
|
||||||
|
prompt.push_str("### 使用方法\n\n");
|
||||||
|
prompt.push_str("- 使用 `get_skill` 工具 action=\"list\" 列出所有可用 skill 及其名称、简介、路径\n");
|
||||||
|
prompt.push_str("- 使用 `get_skill` 工具 action=\"get\" 并提供 `skill_name` 获取指定 skill 完整内容\n");
|
||||||
|
|
||||||
|
// Always skills full content
|
||||||
|
if !always_skills.is_empty() {
|
||||||
|
prompt.push_str("\n---\n\n");
|
||||||
let mut parts = Vec::new();
|
let mut parts = Vec::new();
|
||||||
for skill in always_skills {
|
for skill in &always_skills {
|
||||||
parts.push(format!("## Skill: {}\n\n{}", skill.name, skill.content));
|
parts.push(format!("## Skill: {}\n\n{}", skill.name, skill.content));
|
||||||
}
|
}
|
||||||
prompt.push_str(&parts.join("\n\n---\n\n"));
|
prompt.push_str(&parts.join("\n\n---\n\n"));
|
||||||
prompt.push_str("\n\n");
|
|
||||||
}
|
|
||||||
|
|
||||||
let has_other_skills = state.loaded_skills.iter().any(|s| !s.always);
|
|
||||||
if has_other_skills {
|
|
||||||
prompt.push_str("## Available Skills\n\n");
|
|
||||||
prompt.push_str("Skills teach the agent how to use specific capabilities. Use the `get_skill` tool to load a skill's full content when needed.\n\n");
|
|
||||||
|
|
||||||
let mut lines = vec!["<skills>".to_string()];
|
|
||||||
for skill in &state.loaded_skills {
|
|
||||||
if skill.always {
|
|
||||||
continue;
|
|
||||||
}
|
|
||||||
lines.push(" <skill>".to_string());
|
|
||||||
lines.push(format!(" <name>{}</name>", escape_xml(&skill.name)));
|
|
||||||
lines.push(format!(
|
|
||||||
" <description>{}</description>",
|
|
||||||
escape_xml(&skill.description)
|
|
||||||
));
|
|
||||||
if let Some(path) = &skill.path {
|
|
||||||
lines.push(format!(" <path>{}</path>", escape_xml(&path.to_string_lossy())));
|
|
||||||
}
|
|
||||||
lines.push(" </skill>".to_string());
|
|
||||||
}
|
|
||||||
lines.push("</skills>".to_string());
|
|
||||||
prompt.push_str(&lines.join("\n"));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
prompt
|
prompt
|
||||||
|
|||||||
460
src/tools/content_search.rs
Normal file
460
src/tools/content_search.rs
Normal file
@ -0,0 +1,460 @@
|
|||||||
|
use std::path::Path;
|
||||||
|
use std::process::Stdio;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::json;
|
||||||
|
use tokio::process::Command;
|
||||||
|
use tokio::time::timeout;
|
||||||
|
|
||||||
|
use crate::tools::traits::{Tool, ToolResult};
|
||||||
|
|
||||||
|
const MAX_RESULTS: usize = 100;
|
||||||
|
const MAX_OUTPUT_CHARS: usize = 50_000;
|
||||||
|
const TIMEOUT_SECS: u64 = 60;
|
||||||
|
|
||||||
|
pub struct ContentSearchTool;
|
||||||
|
|
||||||
|
impl ContentSearchTool {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_dir(&self, dir: Option<&str>) -> String {
|
||||||
|
match dir {
|
||||||
|
Some(d) if !d.is_empty() => d.to_string(),
|
||||||
|
_ => ".".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truncate_output(&self, lines: &[String]) -> String {
|
||||||
|
let mut output = String::new();
|
||||||
|
for line in lines {
|
||||||
|
if output.len() + line.len() + 1 > MAX_OUTPUT_CHARS {
|
||||||
|
output.push_str(&format!(
|
||||||
|
"\n... ({} chars truncated, {} matches omitted) ...",
|
||||||
|
output.len(),
|
||||||
|
lines.len()
|
||||||
|
));
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if !output.is_empty() {
|
||||||
|
output.push('\n');
|
||||||
|
}
|
||||||
|
output.push_str(line);
|
||||||
|
}
|
||||||
|
output
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for ContentSearchTool {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for ContentSearchTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"content_search"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Search file contents by regex or text pattern. Uses ripgrep (rg) if available, falls back to grep, then pure Rust."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"pattern": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Regex or text pattern to search for in file contents"
|
||||||
|
},
|
||||||
|
"dir": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Directory to search in (default: current working directory)"
|
||||||
|
},
|
||||||
|
"file_pattern": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Optional glob to restrict which files to search (e.g. '*.rs', '*.{rs,toml}')"
|
||||||
|
},
|
||||||
|
"case_sensitive": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Whether to match case-sensitively (default: false)"
|
||||||
|
},
|
||||||
|
"context_lines": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Number of context lines to show before and after each match (default: 0)"
|
||||||
|
},
|
||||||
|
"max_results": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Maximum number of matching lines to return (default: 100)"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["pattern"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_only(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||||
|
let pattern = match args.get("pattern").and_then(|v| v.as_str()) {
|
||||||
|
Some(p) if !p.is_empty() => p,
|
||||||
|
_ => {
|
||||||
|
return Ok(ToolResult {
|
||||||
|
success: false,
|
||||||
|
output: String::new(),
|
||||||
|
error: Some("Missing required parameter: pattern".to_string()),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let dir = self.resolve_dir(args.get("dir").and_then(|v| v.as_str()));
|
||||||
|
let file_pattern = args.get("file_pattern").and_then(|v| v.as_str());
|
||||||
|
let case_sensitive = args.get("case_sensitive").and_then(|v| v.as_bool()).unwrap_or(false);
|
||||||
|
let context_lines = args.get("context_lines").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
|
||||||
|
let max_results = args.get("max_results").and_then(|v| v.as_u64()).unwrap_or(MAX_RESULTS as u64) as usize;
|
||||||
|
|
||||||
|
let result = self.run_search(pattern, &dir, file_pattern, case_sensitive, context_lines, max_results).await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(lines) => {
|
||||||
|
let count = lines.len();
|
||||||
|
let mut output = self.truncate_output(&lines);
|
||||||
|
output.push_str(&format!("\n\n---\n共 {} 条匹配", count));
|
||||||
|
Ok(ToolResult { success: true, output, error: None })
|
||||||
|
}
|
||||||
|
Err(e) => Ok(ToolResult {
|
||||||
|
success: false,
|
||||||
|
output: String::new(),
|
||||||
|
error: Some(e.to_string()),
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ContentSearchTool {
|
||||||
|
async fn run_search(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
file_pattern: Option<&str>,
|
||||||
|
case_sensitive: bool,
|
||||||
|
context_lines: usize,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
if which::which("rg").is_ok() {
|
||||||
|
match self.search_with_rg(pattern, dir, file_pattern, case_sensitive, context_lines, max_results).await {
|
||||||
|
Ok(lines) => return Ok(lines),
|
||||||
|
Err(e) => tracing::warn!("rg failed: {}, falling back", e),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if which::which("grep").is_ok() {
|
||||||
|
match self.search_with_grep(pattern, dir, file_pattern, case_sensitive, context_lines, max_results).await {
|
||||||
|
Ok(lines) if !lines.is_empty() => return Ok(lines),
|
||||||
|
Ok(_) => {},
|
||||||
|
Err(e) => tracing::warn!("grep failed: {}, falling back", e),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::warn!("No rg/grep available, using built-in content search (much slower). Install ripgrep for better performance.");
|
||||||
|
self.search_with_rust(pattern, dir, file_pattern, case_sensitive, context_lines, max_results).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_rg(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
file_pattern: Option<&str>,
|
||||||
|
case_sensitive: bool,
|
||||||
|
context_lines: usize,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let mut cmd = Command::new("rg");
|
||||||
|
cmd.arg("-n")
|
||||||
|
.arg("--no-heading")
|
||||||
|
.arg("--color").arg("never")
|
||||||
|
.arg("--max-count").arg(max_results.to_string())
|
||||||
|
.arg(pattern)
|
||||||
|
.arg(dir)
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
if !case_sensitive {
|
||||||
|
cmd.arg("-i");
|
||||||
|
}
|
||||||
|
if context_lines > 0 {
|
||||||
|
cmd.arg("-C").arg(context_lines.to_string());
|
||||||
|
}
|
||||||
|
if let Some(fp) = file_pattern {
|
||||||
|
cmd.arg("--glob").arg(fp);
|
||||||
|
}
|
||||||
|
|
||||||
|
let output = timeout(
|
||||||
|
std::time::Duration::from_secs(TIMEOUT_SECS),
|
||||||
|
cmd.output(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow::anyhow!("rg timed out after {}s", TIMEOUT_SECS))??;
|
||||||
|
|
||||||
|
if !output.status.success() && output.status.code() != Some(1) {
|
||||||
|
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||||
|
return Err(anyhow::anyhow!("rg error: {}", stderr.trim()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let text = String::from_utf8_lossy(&output.stdout);
|
||||||
|
let lines: Vec<String> = text.lines()
|
||||||
|
.take(max_results)
|
||||||
|
.map(|l| l.to_string())
|
||||||
|
.collect();
|
||||||
|
Ok(lines)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_grep(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
file_pattern: Option<&str>,
|
||||||
|
case_sensitive: bool,
|
||||||
|
context_lines: usize,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let mut cmd = Command::new("grep");
|
||||||
|
cmd.arg("-rn")
|
||||||
|
.arg("-E")
|
||||||
|
.arg("--color=never")
|
||||||
|
.arg("--binary-files=without-match")
|
||||||
|
.arg(pattern)
|
||||||
|
.arg(dir)
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
if !case_sensitive {
|
||||||
|
cmd.arg("-i");
|
||||||
|
}
|
||||||
|
if context_lines > 0 {
|
||||||
|
cmd.arg("-C").arg(context_lines.to_string());
|
||||||
|
}
|
||||||
|
if let Some(fp) = file_pattern {
|
||||||
|
cmd.arg("--include").arg(fp);
|
||||||
|
}
|
||||||
|
|
||||||
|
let output = timeout(
|
||||||
|
std::time::Duration::from_secs(TIMEOUT_SECS),
|
||||||
|
cmd.output(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow::anyhow!("grep timed out after {}s", TIMEOUT_SECS))??;
|
||||||
|
|
||||||
|
let text = String::from_utf8_lossy(&output.stdout);
|
||||||
|
let lines: Vec<String> = text.lines()
|
||||||
|
.take(max_results)
|
||||||
|
.map(|l| l.to_string())
|
||||||
|
.collect();
|
||||||
|
Ok(lines)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_rust(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
file_pattern: Option<&str>,
|
||||||
|
case_sensitive: bool,
|
||||||
|
_context_lines: usize,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let re = if case_sensitive {
|
||||||
|
regex::Regex::new(pattern)
|
||||||
|
} else {
|
||||||
|
regex::RegexBuilder::new(pattern)
|
||||||
|
.case_insensitive(true)
|
||||||
|
.build()
|
||||||
|
}
|
||||||
|
.map_err(|e| anyhow::anyhow!("Invalid regex pattern '{}': {}", pattern, e))?;
|
||||||
|
|
||||||
|
let file_re = file_pattern.map(|fp| {
|
||||||
|
let re_str = glob_to_regex(fp);
|
||||||
|
if case_sensitive {
|
||||||
|
regex::Regex::new(&re_str)
|
||||||
|
} else {
|
||||||
|
regex::RegexBuilder::new(&re_str).case_insensitive(true).build()
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let file_re = match file_re {
|
||||||
|
Some(Ok(r)) => Some(r),
|
||||||
|
Some(Err(e)) => return Err(anyhow::anyhow!("Invalid file pattern: {}", e)),
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let mut results = Vec::new();
|
||||||
|
grep_dir(Path::new(dir), Path::new(dir), &re, file_re.as_ref(), &mut results, max_results)?;
|
||||||
|
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn glob_to_regex(glob: &str) -> String {
|
||||||
|
let mut regex = String::from("^");
|
||||||
|
let chars: Vec<char> = glob.chars().collect();
|
||||||
|
let mut i = 0;
|
||||||
|
while i < chars.len() {
|
||||||
|
match chars[i] {
|
||||||
|
'*' => {
|
||||||
|
if i + 1 < chars.len() && chars[i + 1] == '*' {
|
||||||
|
regex.push_str(".*");
|
||||||
|
i += 1;
|
||||||
|
} else {
|
||||||
|
regex.push_str("[^/]*");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
'?' => regex.push_str("[^/]"),
|
||||||
|
'.' | '+' | '(' | ')' | '[' | ']' | '{' | '}' | '^' | '$' | '|' | '\\' => {
|
||||||
|
regex.push('\\');
|
||||||
|
regex.push(chars[i]);
|
||||||
|
}
|
||||||
|
c => regex.push(c),
|
||||||
|
}
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
regex.push('$');
|
||||||
|
regex
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grep_dir(
|
||||||
|
base: &Path,
|
||||||
|
current: &Path,
|
||||||
|
re: ®ex::Regex,
|
||||||
|
file_re: Option<®ex::Regex>,
|
||||||
|
results: &mut Vec<String>,
|
||||||
|
max: usize,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
if results.len() >= max {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let entries = match std::fs::read_dir(current) {
|
||||||
|
Ok(e) => e,
|
||||||
|
Err(_) => return Ok(()),
|
||||||
|
};
|
||||||
|
|
||||||
|
for entry in entries.flatten() {
|
||||||
|
let path = entry.path();
|
||||||
|
let rel = match path.strip_prefix(base) {
|
||||||
|
Ok(r) => r,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
if path.is_dir() {
|
||||||
|
if let Some(name) = rel.file_name().and_then(|n| n.to_str()) {
|
||||||
|
if name.starts_with('.') && name.len() > 1 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
grep_dir(base, &path, re, file_re, results, max)?;
|
||||||
|
} else if path.is_file() {
|
||||||
|
if let Some(file_re) = file_re {
|
||||||
|
if let Some(name) = rel.file_name().and_then(|n| n.to_str()) {
|
||||||
|
if !file_re.is_match(name) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Ok(content) = std::fs::read_to_string(&path) {
|
||||||
|
for (line_num, line) in content.lines().enumerate() {
|
||||||
|
if re.is_match(line) {
|
||||||
|
results.push(format!(
|
||||||
|
"{}:{}:{}",
|
||||||
|
rel.to_string_lossy(),
|
||||||
|
line_num + 1,
|
||||||
|
line
|
||||||
|
));
|
||||||
|
if results.len() >= max {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::fs;
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_content_search_rust_fallback() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
fs::write(dir.path().join("main.rs"), "fn main() {\n let x = 42;\n println!(\"hello\");\n}").unwrap();
|
||||||
|
fs::write(dir.path().join("lib.rs"), "pub fn foo() -> u32 {\n let y = 42;\n y\n}").unwrap();
|
||||||
|
fs::write(dir.path().join("README.md"), "# Project\nHello world").unwrap();
|
||||||
|
|
||||||
|
let tool = ContentSearchTool::new();
|
||||||
|
let result = tool
|
||||||
|
.execute(json!({
|
||||||
|
"pattern": "let.*=.*42",
|
||||||
|
"dir": dir.path().to_str().unwrap()
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(result.success);
|
||||||
|
assert!(result.output.contains("main.rs"));
|
||||||
|
assert!(result.output.contains("lib.rs"));
|
||||||
|
assert!(!result.output.contains("README.md"));
|
||||||
|
assert!(result.output.contains("共 2 条匹配"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_content_search_file_filter() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
fs::write(dir.path().join("main.rs"), "fn main() {}").unwrap();
|
||||||
|
fs::write(dir.path().join("config.toml"), "name = \"test\"").unwrap();
|
||||||
|
|
||||||
|
let tool = ContentSearchTool::new();
|
||||||
|
let result = tool
|
||||||
|
.execute(json!({
|
||||||
|
"pattern": "test",
|
||||||
|
"dir": dir.path().to_str().unwrap(),
|
||||||
|
"file_pattern": "*.toml"
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(result.success);
|
||||||
|
assert!(result.output.contains("config.toml"));
|
||||||
|
assert!(!result.output.contains("main.rs"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_content_search_max_results() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
let mut content = String::new();
|
||||||
|
for i in 0..10 {
|
||||||
|
content.push_str(&format!("match line {}\n", i));
|
||||||
|
}
|
||||||
|
fs::write(dir.path().join("data.txt"), &content).unwrap();
|
||||||
|
|
||||||
|
let tool = ContentSearchTool::new();
|
||||||
|
let result = tool
|
||||||
|
.execute(json!({
|
||||||
|
"pattern": "match line",
|
||||||
|
"dir": dir.path().to_str().unwrap(),
|
||||||
|
"max_results": 3
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(result.success);
|
||||||
|
assert!(result.output.contains("共 3 条匹配"));
|
||||||
|
}
|
||||||
|
}
|
||||||
375
src/tools/file_search.rs
Normal file
375
src/tools/file_search.rs
Normal file
@ -0,0 +1,375 @@
|
|||||||
|
use std::path::Path;
|
||||||
|
use std::process::Stdio;
|
||||||
|
|
||||||
|
use async_trait::async_trait;
|
||||||
|
use serde_json::json;
|
||||||
|
use tokio::process::Command;
|
||||||
|
use tokio::time::timeout;
|
||||||
|
|
||||||
|
use crate::tools::traits::{Tool, ToolResult};
|
||||||
|
|
||||||
|
const MAX_RESULTS: usize = 200;
|
||||||
|
const MAX_OUTPUT_CHARS: usize = 50_000;
|
||||||
|
const TIMEOUT_SECS: u64 = 60;
|
||||||
|
|
||||||
|
pub struct FileSearchTool;
|
||||||
|
|
||||||
|
impl FileSearchTool {
|
||||||
|
pub fn new() -> Self {
|
||||||
|
Self
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_dir(&self, dir: Option<&str>) -> String {
|
||||||
|
match dir {
|
||||||
|
Some(d) if !d.is_empty() => d.to_string(),
|
||||||
|
_ => ".".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn truncate_output(&self, lines: &[String]) -> String {
|
||||||
|
let mut output = String::new();
|
||||||
|
for line in lines {
|
||||||
|
if output.len() + line.len() + 1 > MAX_OUTPUT_CHARS {
|
||||||
|
output.push_str(&format!("\n... ({} chars truncated) ...", output.len()));
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if !output.is_empty() {
|
||||||
|
output.push('\n');
|
||||||
|
}
|
||||||
|
output.push_str(line);
|
||||||
|
}
|
||||||
|
output
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for FileSearchTool {
|
||||||
|
fn default() -> Self {
|
||||||
|
Self::new()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl Tool for FileSearchTool {
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"file_search"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Search for files by glob pattern (e.g. '*.rs', 'test_*.rs'). Uses fd if available, falls back to find, then pure Rust."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parameters_schema(&self) -> serde_json::Value {
|
||||||
|
json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"pattern": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "File glob pattern to search for (e.g. *.rs, test_*.rs, src/**/*.py)"
|
||||||
|
},
|
||||||
|
"dir": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Directory to search in (default: current working directory)"
|
||||||
|
},
|
||||||
|
"case_sensitive": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Whether to match case-sensitively (default: true)"
|
||||||
|
},
|
||||||
|
"max_results": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Maximum number of results to return (default: 200)"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["pattern"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_only(&self) -> bool {
|
||||||
|
true
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||||
|
let pattern = match args.get("pattern").and_then(|v| v.as_str()) {
|
||||||
|
Some(p) if !p.is_empty() => p,
|
||||||
|
_ => {
|
||||||
|
return Ok(ToolResult {
|
||||||
|
success: false,
|
||||||
|
output: String::new(),
|
||||||
|
error: Some("Missing required parameter: pattern".to_string()),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let dir = self.resolve_dir(args.get("dir").and_then(|v| v.as_str()));
|
||||||
|
let case_sensitive = args.get("case_sensitive").and_then(|v| v.as_bool()).unwrap_or(true);
|
||||||
|
let max_results = args.get("max_results").and_then(|v| v.as_u64()).unwrap_or(MAX_RESULTS as u64) as usize;
|
||||||
|
|
||||||
|
let result = self.run_search(pattern, &dir, case_sensitive, max_results).await;
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(lines) => {
|
||||||
|
let count = lines.len();
|
||||||
|
let mut output = self.truncate_output(&lines);
|
||||||
|
output.push_str(&format!("\n\n---\n共 {} 个文件", count));
|
||||||
|
Ok(ToolResult { success: true, output, error: None })
|
||||||
|
}
|
||||||
|
Err(e) => Ok(ToolResult {
|
||||||
|
success: false,
|
||||||
|
output: String::new(),
|
||||||
|
error: Some(e.to_string()),
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FileSearchTool {
|
||||||
|
async fn run_search(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
case_sensitive: bool,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
if which::which("fd").is_ok() {
|
||||||
|
match self.search_with_fd(pattern, dir, case_sensitive, max_results).await {
|
||||||
|
Ok(lines) if !lines.is_empty() => return Ok(lines),
|
||||||
|
Ok(_) => {},
|
||||||
|
Err(e) => tracing::warn!("fd failed: {}, falling back", e),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if which::which("find").is_ok() {
|
||||||
|
match self.search_with_find(pattern, dir, max_results).await {
|
||||||
|
Ok(lines) if !lines.is_empty() => return Ok(lines),
|
||||||
|
Ok(_) => {},
|
||||||
|
Err(e) => tracing::warn!("find failed: {}, falling back", e),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tracing::warn!("No fd/find available, using built-in file search (slower)");
|
||||||
|
self.search_with_rust(pattern, dir, case_sensitive, max_results).await
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_fd(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
case_sensitive: bool,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let mut cmd = Command::new("fd");
|
||||||
|
cmd.arg("--search-path").arg(dir)
|
||||||
|
.arg("--glob").arg(pattern)
|
||||||
|
.arg("--color").arg("never")
|
||||||
|
.arg("--strip-cwd-prefix")
|
||||||
|
.arg("--max-results").arg(max_results.to_string())
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
if !case_sensitive {
|
||||||
|
cmd.arg("--ignore-case");
|
||||||
|
}
|
||||||
|
|
||||||
|
let output = timeout(
|
||||||
|
std::time::Duration::from_secs(TIMEOUT_SECS),
|
||||||
|
cmd.output(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow::anyhow!("fd timed out after {}s", TIMEOUT_SECS))??;
|
||||||
|
|
||||||
|
if !output.status.success() {
|
||||||
|
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||||
|
return Err(anyhow::anyhow!("fd error: {}", stderr.trim()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let text = String::from_utf8_lossy(&output.stdout);
|
||||||
|
let lines: Vec<String> = text.lines()
|
||||||
|
.filter(|l| !l.is_empty())
|
||||||
|
.map(|l| l.to_string())
|
||||||
|
.collect();
|
||||||
|
Ok(lines)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_find(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let limit_str = max_results.to_string();
|
||||||
|
let mut cmd = Command::new("sh");
|
||||||
|
cmd.arg("-c")
|
||||||
|
.arg(format!(
|
||||||
|
"find '{}' -name '{}' -not -path '*/.*' 2>/dev/null | head -n {}",
|
||||||
|
dir, pattern, limit_str
|
||||||
|
))
|
||||||
|
.stdout(Stdio::piped())
|
||||||
|
.stderr(Stdio::piped());
|
||||||
|
|
||||||
|
let output = timeout(
|
||||||
|
std::time::Duration::from_secs(TIMEOUT_SECS),
|
||||||
|
cmd.output(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| anyhow::anyhow!("find timed out after {}s", TIMEOUT_SECS))??;
|
||||||
|
|
||||||
|
let text = String::from_utf8_lossy(&output.stdout);
|
||||||
|
let lines: Vec<String> = text.lines()
|
||||||
|
.filter(|l| !l.is_empty())
|
||||||
|
.map(|l| {
|
||||||
|
let p = Path::new(l);
|
||||||
|
p.to_string_lossy().to_string()
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
Ok(lines)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn search_with_rust(
|
||||||
|
&self,
|
||||||
|
pattern: &str,
|
||||||
|
dir: &str,
|
||||||
|
case_sensitive: bool,
|
||||||
|
max_results: usize,
|
||||||
|
) -> anyhow::Result<Vec<String>> {
|
||||||
|
let regex_str = glob_to_regex(pattern);
|
||||||
|
let re = if case_sensitive {
|
||||||
|
regex::Regex::new(®ex_str)
|
||||||
|
} else {
|
||||||
|
regex::RegexBuilder::new(®ex_str)
|
||||||
|
.case_insensitive(true)
|
||||||
|
.build()
|
||||||
|
}
|
||||||
|
.map_err(|e| anyhow::anyhow!("Invalid glob pattern '{}': {}", pattern, e))?;
|
||||||
|
|
||||||
|
let mut results = Vec::new();
|
||||||
|
walk_dir(Path::new(dir), Path::new(dir), &re, &mut results, max_results)?;
|
||||||
|
Ok(results)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn glob_to_regex(glob: &str) -> String {
|
||||||
|
let mut regex = String::from("^");
|
||||||
|
let chars: Vec<char> = glob.chars().collect();
|
||||||
|
let mut i = 0;
|
||||||
|
while i < chars.len() {
|
||||||
|
match chars[i] {
|
||||||
|
'*' => {
|
||||||
|
if i + 1 < chars.len() && chars[i + 1] == '*' {
|
||||||
|
regex.push_str(".*");
|
||||||
|
i += 1;
|
||||||
|
} else {
|
||||||
|
regex.push_str("[^/]*");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
'?' => regex.push_str("[^/]"),
|
||||||
|
'.' | '+' | '(' | ')' | '[' | ']' | '{' | '}' | '^' | '$' | '|' | '\\' => {
|
||||||
|
regex.push('\\');
|
||||||
|
regex.push(chars[i]);
|
||||||
|
}
|
||||||
|
c => regex.push(c),
|
||||||
|
}
|
||||||
|
i += 1;
|
||||||
|
}
|
||||||
|
regex.push('$');
|
||||||
|
regex
|
||||||
|
}
|
||||||
|
|
||||||
|
fn walk_dir(
|
||||||
|
base: &Path,
|
||||||
|
current: &Path,
|
||||||
|
re: ®ex::Regex,
|
||||||
|
results: &mut Vec<String>,
|
||||||
|
max: usize,
|
||||||
|
) -> anyhow::Result<()> {
|
||||||
|
if results.len() >= max {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
|
||||||
|
let entries = match std::fs::read_dir(current) {
|
||||||
|
Ok(e) => e,
|
||||||
|
Err(_) => return Ok(()),
|
||||||
|
};
|
||||||
|
|
||||||
|
for entry in entries.flatten() {
|
||||||
|
let path = entry.path();
|
||||||
|
let rel = match path.strip_prefix(base) {
|
||||||
|
Ok(r) => r,
|
||||||
|
Err(_) => continue,
|
||||||
|
};
|
||||||
|
|
||||||
|
if path.is_dir() {
|
||||||
|
if let Some(name) = rel.file_name().and_then(|n| n.to_str()) {
|
||||||
|
if name.starts_with('.') && name.len() > 1 {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
walk_dir(base, &path, re, results, max)?;
|
||||||
|
} else if path.is_file() {
|
||||||
|
if let Some(name) = rel.file_name().and_then(|n| n.to_str()) {
|
||||||
|
if re.is_match(name) {
|
||||||
|
results.push(rel.to_string_lossy().to_string());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if results.len() >= max {
|
||||||
|
return Ok(());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use std::fs;
|
||||||
|
use tempfile::TempDir;
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_file_search_rust_fallback() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
fs::write(dir.path().join("main.rs"), "fn main() {}").unwrap();
|
||||||
|
fs::write(dir.path().join("lib.rs"), "pub fn foo() {}").unwrap();
|
||||||
|
fs::write(dir.path().join("test.rs"), "#[test] fn t() {}").unwrap();
|
||||||
|
fs::write(dir.path().join("README.md"), "# Readme").unwrap();
|
||||||
|
fs::create_dir(dir.path().join("src")).unwrap();
|
||||||
|
fs::write(dir.path().join("src/nested.rs"), "fn nested() {}").unwrap();
|
||||||
|
|
||||||
|
let tool = FileSearchTool::new();
|
||||||
|
let result = tool
|
||||||
|
.execute(json!({
|
||||||
|
"pattern": "*.rs",
|
||||||
|
"dir": dir.path().to_str().unwrap()
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(result.success);
|
||||||
|
assert!(result.output.contains("main.rs"));
|
||||||
|
assert!(result.output.contains("lib.rs"));
|
||||||
|
assert!(result.output.contains("test.rs"));
|
||||||
|
assert!(result.output.contains("nested.rs"));
|
||||||
|
assert!(!result.output.contains("README.md"));
|
||||||
|
assert!(result.output.contains("共 4 个文件"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_file_search_max_results() {
|
||||||
|
let dir = TempDir::new().unwrap();
|
||||||
|
for i in 0..5 {
|
||||||
|
fs::write(dir.path().join(format!("file_{}.rs", i)), "").unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
let tool = FileSearchTool::new();
|
||||||
|
let result = tool
|
||||||
|
.execute(json!({
|
||||||
|
"pattern": "*.rs",
|
||||||
|
"dir": dir.path().to_str().unwrap(),
|
||||||
|
"max_results": 3
|
||||||
|
}))
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
assert!(result.success);
|
||||||
|
assert!(result.output.contains("共 3 个文件"));
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -44,12 +44,17 @@ impl Tool for GetSkillTool {
|
|||||||
json!({
|
json!({
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
|
"action": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["get", "list"],
|
||||||
|
"description": "操作类型: get 获取指定 skill 完整内容, list 列出所有可用 skill"
|
||||||
|
},
|
||||||
"skill_name": {
|
"skill_name": {
|
||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "Name of the skill to retrieve"
|
"description": "Name of the skill to retrieve,仅在 action 为 get 时必填"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"required": ["skill_name"]
|
"required": []
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -58,6 +63,17 @@ impl Tool for GetSkillTool {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||||
|
let action = args.get("action").and_then(|v| v.as_str()).unwrap_or("get");
|
||||||
|
|
||||||
|
match action {
|
||||||
|
"list" => self.list_skills_full(),
|
||||||
|
_ => self.get_skill_by_name(&args),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GetSkillTool {
|
||||||
|
fn get_skill_by_name(&self, args: &serde_json::Value) -> anyhow::Result<ToolResult> {
|
||||||
let skill_name = match args.get("skill_name").and_then(|v| v.as_str()) {
|
let skill_name = match args.get("skill_name").and_then(|v| v.as_str()) {
|
||||||
Some(name) => name,
|
Some(name) => name,
|
||||||
None => {
|
None => {
|
||||||
@ -100,6 +116,33 @@ impl Tool for GetSkillTool {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn list_skills_full(&self) -> anyhow::Result<ToolResult> {
|
||||||
|
let skills = self.skills_loader.get_loaded_skills();
|
||||||
|
if skills.is_empty() {
|
||||||
|
return Ok(ToolResult {
|
||||||
|
success: true,
|
||||||
|
output: "当前没有安装任何 skill".to_string(),
|
||||||
|
error: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
let mut output = format!("可用 skill (共 {} 个):\n", skills.len());
|
||||||
|
for s in &skills {
|
||||||
|
let always_mark = if s.always { " [常驻]" } else { "" };
|
||||||
|
let path_str = s.path.as_ref()
|
||||||
|
.map(|p| p.to_string_lossy().to_string())
|
||||||
|
.unwrap_or_else(|| "—".to_string());
|
||||||
|
output.push_str(&format!(
|
||||||
|
"- {}{}\n 简介: {}\n 路径: {}\n",
|
||||||
|
s.name, always_mark, s.description, path_str
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(ToolResult {
|
||||||
|
success: true,
|
||||||
|
output,
|
||||||
|
error: None,
|
||||||
|
})
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
|
|||||||
@ -1,9 +1,11 @@
|
|||||||
pub mod bash;
|
pub mod bash;
|
||||||
pub mod calculator;
|
pub mod calculator;
|
||||||
pub mod chat_manager;
|
pub mod chat_manager;
|
||||||
|
pub mod content_search;
|
||||||
pub mod cron;
|
pub mod cron;
|
||||||
pub mod file_edit;
|
pub mod file_edit;
|
||||||
pub mod file_read;
|
pub mod file_read;
|
||||||
|
pub mod file_search;
|
||||||
pub mod file_write;
|
pub mod file_write;
|
||||||
pub mod get_skill;
|
pub mod get_skill;
|
||||||
pub mod http_request;
|
pub mod http_request;
|
||||||
@ -17,8 +19,10 @@ pub mod web_fetch;
|
|||||||
pub use bash::BashTool;
|
pub use bash::BashTool;
|
||||||
pub use calculator::CalculatorTool;
|
pub use calculator::CalculatorTool;
|
||||||
pub use chat_manager::ChatManagerTool;
|
pub use chat_manager::ChatManagerTool;
|
||||||
|
pub use content_search::ContentSearchTool;
|
||||||
pub use file_edit::FileEditTool;
|
pub use file_edit::FileEditTool;
|
||||||
pub use file_read::FileReadTool;
|
pub use file_read::FileReadTool;
|
||||||
|
pub use file_search::FileSearchTool;
|
||||||
pub use file_write::FileWriteTool;
|
pub use file_write::FileWriteTool;
|
||||||
pub use get_skill::GetSkillTool;
|
pub use get_skill::GetSkillTool;
|
||||||
pub use http_request::HttpRequestTool;
|
pub use http_request::HttpRequestTool;
|
||||||
@ -45,6 +49,8 @@ pub fn create_default_tools(
|
|||||||
registry.register(FileReadTool::new());
|
registry.register(FileReadTool::new());
|
||||||
registry.register(FileWriteTool::new());
|
registry.register(FileWriteTool::new());
|
||||||
registry.register(FileEditTool::new());
|
registry.register(FileEditTool::new());
|
||||||
|
registry.register(FileSearchTool::new());
|
||||||
|
registry.register(ContentSearchTool::new());
|
||||||
registry.register(BashTool::new());
|
registry.register(BashTool::new());
|
||||||
registry.register(HttpRequestTool::new(
|
registry.register(HttpRequestTool::new(
|
||||||
vec!["*".to_string()],
|
vec!["*".to_string()],
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user