140 lines
4.6 KiB
Rust
140 lines
4.6 KiB
Rust
use std::sync::Arc;
|
|
|
|
use async_trait::async_trait;
|
|
use llm_engine::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
|
|
use schemars::JsonSchema;
|
|
use serde::Deserialize;
|
|
use workdir::{GrepOutputMode, GrepRequest, WorkdirPath, WorkdirSessionHandle};
|
|
|
|
use crate::ToolsError;
|
|
|
|
const DEFAULT_HEAD_LIMIT: usize = 250;
|
|
|
|
#[derive(Debug, Clone, Copy, Default, Deserialize, JsonSchema)]
|
|
#[serde(rename_all = "snake_case")]
|
|
enum OutputMode {
|
|
Content,
|
|
#[default]
|
|
FilesWithMatches,
|
|
Count,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize, JsonSchema)]
|
|
struct GrepParams {
|
|
pattern: String,
|
|
/// Logical Workdir-relative path to search. Defaults to the Workdir root.
|
|
#[serde(default)]
|
|
path: Option<String>,
|
|
#[serde(default)]
|
|
glob: Option<String>,
|
|
#[serde(default, rename = "type")]
|
|
file_type: Option<String>,
|
|
#[serde(default)]
|
|
case_insensitive: bool,
|
|
#[serde(default, rename = "-B")]
|
|
before: Option<usize>,
|
|
#[serde(default, rename = "-A")]
|
|
after: Option<usize>,
|
|
#[serde(default, rename = "-C")]
|
|
context: Option<usize>,
|
|
#[serde(default)]
|
|
multiline: bool,
|
|
#[serde(default)]
|
|
output_mode: Option<OutputMode>,
|
|
#[serde(default)]
|
|
head_limit: Option<usize>,
|
|
#[serde(default)]
|
|
offset: Option<usize>,
|
|
}
|
|
|
|
struct GrepTool {
|
|
session: WorkdirSessionHandle,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Tool for GrepTool {
|
|
async fn execute(
|
|
&self,
|
|
input_json: &str,
|
|
_ctx: llm_engine::tool::ToolExecutionContext,
|
|
) -> Result<ToolOutput, ToolError> {
|
|
let params: GrepParams = serde_json::from_str(input_json)
|
|
.map_err(|error| ToolError::InvalidArgument(format!("invalid Grep input: {error}")))?;
|
|
let path = match params.path {
|
|
Some(path) => WorkdirPath::new(&path).map_err(ToolsError::from)?,
|
|
None => WorkdirPath::root(),
|
|
};
|
|
let mode = match params.output_mode.unwrap_or_default() {
|
|
OutputMode::FilesWithMatches => GrepOutputMode::FilesWithMatches,
|
|
OutputMode::Content => GrepOutputMode::Content,
|
|
OutputMode::Count => GrepOutputMode::Count,
|
|
};
|
|
let (before_context, after_context) = params
|
|
.context
|
|
.map(|context| (context, context))
|
|
.unwrap_or((params.before.unwrap_or(0), params.after.unwrap_or(0)));
|
|
let head_limit = params.head_limit.unwrap_or(DEFAULT_HEAD_LIMIT);
|
|
let result = self
|
|
.session
|
|
.grep(GrepRequest {
|
|
pattern: params.pattern,
|
|
path,
|
|
glob: params.glob,
|
|
file_type: params.file_type,
|
|
case_insensitive: params.case_insensitive,
|
|
before_context,
|
|
after_context,
|
|
multiline: params.multiline,
|
|
output_mode: mode,
|
|
limit: head_limit,
|
|
offset: params.offset.unwrap_or(0),
|
|
})
|
|
.await
|
|
.map_err(ToolsError::from)?;
|
|
|
|
let summary = if result.match_count == 0 {
|
|
match mode {
|
|
GrepOutputMode::Content => "No matches".to_owned(),
|
|
_ => "No files matched".to_owned(),
|
|
}
|
|
} else {
|
|
match mode {
|
|
GrepOutputMode::FilesWithMatches => {
|
|
format!("Found matches in {} file(s)", result.matched_files)
|
|
}
|
|
GrepOutputMode::Count => format!(
|
|
"Found matches in {} file(s), {} total line(s)",
|
|
result.matched_files, result.match_count
|
|
),
|
|
GrepOutputMode::Content => format!(
|
|
"{} matching line(s) in {} file(s)",
|
|
result.match_count, result.matched_files
|
|
),
|
|
}
|
|
};
|
|
let summary = if result.truncated {
|
|
format!("{summary} (truncated at {head_limit})")
|
|
} else {
|
|
summary
|
|
};
|
|
Ok(ToolOutput {
|
|
summary,
|
|
content: (!result.output.is_empty()).then_some(result.output),
|
|
attachments: Vec::new(),
|
|
})
|
|
}
|
|
}
|
|
|
|
pub fn grep_tool(session: WorkdirSessionHandle) -> ToolDefinition {
|
|
Arc::new(move || {
|
|
let schema = schemars::schema_for!(GrepParams);
|
|
let meta = ToolMeta::new("Grep")
|
|
.description("Search Workdir file contents with a regex. Glob/Grep traversal executes inside the WorkdirSession provider. Results are bounded and Workdir-relative.")
|
|
.input_schema(serde_json::to_value(schema).expect("Grep schema serialization"));
|
|
let tool: Arc<dyn Tool> = Arc::new(GrepTool {
|
|
session: session.clone(),
|
|
});
|
|
(meta, tool)
|
|
})
|
|
}
|