diff --git a/src/agent.rs b/src/agent.rs index 1475a4f..097281d 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -16,6 +16,8 @@ pub enum ToolError { Io(#[from] std::io::Error), #[error("Command execution failed: {0}")] CommandFailed(String), + #[error("Invalid line range: {0}")] + InvalidLineRange(String), } // Create file tool @@ -105,6 +107,10 @@ impl Tool for DeleteFileTool { #[derive(Deserialize)] pub struct ReadFileArgs { path: String, + #[serde(default)] + start_line: Option, + #[serde(default)] + end_line: Option, } #[derive(Serialize, Deserialize)] @@ -119,13 +125,21 @@ impl Tool for ReadFileTool { async fn definition(&self, _prompt: String) -> ToolDefinition { ToolDefinition { name: Self::NAME.to_string(), - description: "Read the content of a file".to_string(), + description: "Read the content of a file, optionally with line range".to_string(), parameters: json!({ "type": "object", "properties": { "path": { "type": "string", "description": "Path of the file to read" + }, + "start_line": { + "type": "number", + "description": "Optional: Starting line number (1-based, inclusive)" + }, + "end_line": { + "type": "number", + "description": "Optional: Ending line number (1-based, inclusive)" } }, "required": ["path"] @@ -135,8 +149,81 @@ impl Tool for ReadFileTool { async fn call(&self, args: Self::Args) -> Result { let content = fs::read_to_string(&args.path)?; - println!("file content:\n{}", content); - Ok(format!("file content:\n{}", content)) + + let result = match (args.start_line, args.end_line) { + (Some(start), Some(end)) if start > end || start == 0 => { + return Err(ToolError::InvalidLineRange( + format!("Invalid range: start={}, end={}", start, end) + )); + } + (Some(start), Some(end)) => { + content.lines() + .skip(start.saturating_sub(1)) + .take(end - start + 1) + .collect::>() + .join("\n") + } + (Some(start), None) => { + content.lines() + .skip(start.saturating_sub(1)) + .collect::>() + .join("\n") + } + (None, Some(end)) => { + content.lines() + .take(end) + .collect::>() + .join("\n") + } + (None, None) => content, + }; + + println!("file content:\n{}", result); + Ok(format!("file content:\n{}", result)) + } +} + +// Edit file tool +#[derive(Deserialize)] +pub struct EditFileArgs { + path: String, + content: String, +} + +#[derive(Serialize, Deserialize)] +pub struct EditFileTool; + +impl Tool for EditFileTool { + const NAME: &'static str = "edit_file"; + type Error = ToolError; + type Args = EditFileArgs; + type Output = String; + + async fn definition(&self, _prompt: String) -> ToolDefinition { + ToolDefinition { + name: Self::NAME.to_string(), + description: "Edit or overwrite an existing file with new content".to_string(), + parameters: json!({ + "type": "object", + "properties": { + "path": { + "type": "string", + "description": "Path of the file to edit" + }, + "content": { + "type": "string", + "description": "New content to write to the file" + } + }, + "required": ["path", "content"] + }), + } + } + + async fn call(&self, args: Self::Args) -> Result { + fs::write(&args.path, &args.content)?; + println!("file edited: {}", args.path); + Ok(format!("file edited: {}", args.path)) } } diff --git a/src/main.rs b/src/main.rs index d275748..e8bbb33 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,6 +26,7 @@ async fn main() { .tool(agent::CreateFileTool) .tool(agent::DeleteFileTool) .tool(agent::ReadFileTool) + .tool(agent::EditFileTool) .tool(agent::ExecuteCommandTool) .build();