monothink 0.4.5

/*
* This file is part of monothink.
*
* monothink is free software: you can redistribute it and/or modify it under the terms
* of the GNU General Public License as published by the Free Software Foundation,
* either version 3 of the License, or (at your option) any later version.
*
* monothink is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
* See the GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License along with Wavething.
* If not, see <https://www.gnu.org/licenses/>.

*/

use std::{
    fs,
    path::{Path, PathBuf},
};

use rig::tool::{Tool, ToolContext};
use serde::{Deserialize, Serialize};
use serde_json::json;

#[derive(Deserialize, Serialize)]
pub struct ReadFileArgs {
    path: PathBuf,
    #[serde(default)]
    start: Option<usize>,
    #[serde(default)]
    end: Option<usize>,
}

#[derive(Deserialize, Serialize)]
pub struct ReadFile {
    base_dir: PathBuf,
}

impl ReadFile {
    pub fn new(base_dir: &Path) -> Self {
        Self {
            base_dir: base_dir.to_path_buf(),
        }
    }
}

// @todo check access rights (i.e. calculated path must be below base_dir)
impl Tool for ReadFile {
    const NAME: &'static str = "read_file";

    type Error = super::ToolError;
    type Args = ReadFileArgs;
    type Output = super::FileContent;

    async fn call(
        &self,
        _context: &mut ToolContext,
        args: Self::Args,
    ) -> Result<Self::Output, Self::Error> {
        let path = super::allowed_canonical_path(&args.path, &self.base_dir)?;

        if path.is_dir() {
            return Err(super::ToolError::NotAFile {
                path: args.path.to_string_lossy().to_string(),
                suggestion: "Use list_dir to explore this directory instead.".to_string(),
            });
        }

        let content: Vec<String> = fs::read_to_string(path)
            .map_err(|err| super::io_error_to_tool_error(err, &args.path))?
            .lines()
            .map(|line| line.to_string())
            .collect();

        let mut start = match args.start {
            Some(v) => v.saturating_sub(1),
            None => 0,
        };
        let mut end = match args.end {
            Some(v) => {
                if v > 0 {
                    v // we do not substract to convert from one-based to zero-based because the line shall be included which it isn't for slicing
                } else {
                    1 // LLMs like to provide 0 or -1 for end of file, we do not support -1 but can support 0
                }
            }
            None => content.len(),
        };
        if end > content.len() {
            end = content.len();
        }
        if start > end {
            start = end;
        };

        let extract = &content[start..end].join("\n");

        // @todo clip large files
        Ok(super::FileContent {
            range: super::Range {
                path: args.path,
                start: start + 1,
                end: end + 1,
            },
            content: extract.to_owned(),
        })
    }

    fn description(&self) -> String {
        "Read the content of a file, optionally specifying a line range. \
            Lines are 1-indexed and inclusive. \
            If no range is specified, the entire file is returned."
            .to_string()
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Path to the file relative to workspace root directory"
                },
                "start": {
                    "type": "integer",
                    "description": "First line number to read (1-indexed, inclusive). \
                    If not provided, reads from the beginning of the file."
                },
                "end": {
                    "type": "integer",
                    "description": "Last line number to read (1-indexed, inclusive). \
                    If not provided, reads to the end of the file. \
                      If end exceeds the file length, reads until the end."
                }
            },
            "required": ["path"]
        })
    }

    fn map_error(&self, error: Self::Error) -> rig::tool::ToolExecutionError {
        rig::tool::ToolExecutionError::from_error(error)
    }
}

#[tokio::test]
async fn test_read_file() {
    // fake test to just see the output
    // Should output the GPLv3 license header of main.rs
    let base_dir = std::env::current_dir().expect("Cannot get current directory");
    let ld = ReadFile::new(&base_dir);
    let result = ld
        .call(
            &mut ToolContext::default(),
            ReadFileArgs {
                path: PathBuf::from("src/main.rs"),
                start: Some(1),
                end: Some(14),
            },
        )
        .await
        .unwrap();
    println!("Result: {:?}", result);
}