fix: update fs_write tool to try to accomodate hallucinations, other UX changes for fs_read (#730)

Co-authored-by: Chay Nabors <[email protected]>

Author: brandonskiser

Cargo.lock

@@ -128,6 +128,7 @@ serde = { version = "1.0.216", features = ["derive", "rc"] }
 serde_json = "1.0.115"
 sha2 = "0.10.6"
 shlex = "1.3.0"
+similar = "2.6.0"
 strum = { version = "0.26.3", features = ["derive"] }
 sysinfo = "0.32.0"
 thiserror = "2.0.3"

Cargo.toml

@@ -68,6 +68,7 @@ semver.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 shlex.workspace = true
+similar.workspace = true
 spinners = "4.1.0"
 sysinfo.workspace = true
 tempfile.workspace = true

crates/q_cli/Cargo.toml

@@ -26,7 +26,6 @@ use super::{
     OutputKind,
     format_path,
     sanitize_path_tool_arg,
-    stylize_output_if_able,
 };
 
 #[derive(Debug, Clone, Deserialize)]
@@ -52,6 +51,7 @@ impl FsRead {
     pub async fn invoke(&self, ctx: &Context, updates: &mut impl Write) -> Result<InvokeOutput> {
         let path = sanitize_path_tool_arg(ctx, &self.path);
         let is_file = ctx.fs().symlink_metadata(&path).await?.is_file();
+        let relative_path = format_path(ctx.env().current_dir()?, &path);
         debug!(?path, is_file, "Reading");
 
         if is_file {
@@ -84,8 +84,22 @@ impl FsRead {
                     .collect::<Vec<_>>()
                     .join("\n");
 
-                let formatted = stylize_output_if_able(ctx, &path, file_contents.as_str(), Some(start + 1), None);
-                queue!(updates, style::Print(formatted), style::ResetColor, style::Print("\n"))?;
+                queue!(
+                    updates,
+                    style::Print("Reading: "),
+                    style::SetForegroundColor(Color::Green),
+                    style::Print(relative_path),
+                    style::ResetColor,
+                    style::Print(", lines "),
+                    style::SetForegroundColor(Color::Green),
+                    style::Print(start + 1),
+                    style::ResetColor,
+                    style::Print("-"),
+                    style::SetForegroundColor(Color::Green),
+                    style::Print(end + 1),
+                    style::ResetColor,
+                    style::Print("\n"),
+                )?;
 
                 // TODO: We copy and paste this below so the control flow needs work in this tool
                 let byte_count = file_contents.len();
@@ -108,13 +122,11 @@ impl FsRead {
                 );
             }
 
-            let file_text = stylize_output_if_able(ctx, path, file.as_str(), None, None);
             queue!(
                 updates,
+                style::Print("Reading: "),
                 style::SetForegroundColor(Color::Green),
-                style::Print("Reading:\n"),
-                style::ResetColor,
-                style::Print(file_text),
+                style::Print(relative_path),
                 style::ResetColor,
                 style::Print("\n"),
             )?;
@@ -135,8 +147,9 @@ impl FsRead {
                 let relative_path = format_path(&cwd, &path);
                 queue!(
                     updates,
+                    style::Print("Reading: "),
                     style::SetForegroundColor(Color::Green),
-                    style::Print(format!("Reading directory {}", &relative_path)),
+                    style::Print(&relative_path),
                     style::ResetColor,
                     style::Print("\n"),
                 )?;

crates/q_cli/src/cli/chat/tools/fs_read.rs

@@ -13,6 +13,7 @@ use eyre::{
 };
 use fig_os_shim::Context;
 use serde::Deserialize;
+use tracing::warn;
 
 use super::{
     InvokeOutput,
@@ -24,8 +25,14 @@ use super::{
 #[derive(Debug, Clone, Deserialize)]
 #[serde(tag = "command")]
 pub enum FsWrite {
+    /// The tool spec should only require `file_text`, but the model sometimes doesn't want to
+    /// provide it. Thus, including `new_str` as a fallback check, if it's available.
     #[serde(rename = "create")]
-    Create { path: String, file_text: String },
+    Create {
+        path: String,
+        file_text: Option<String>,
+        new_str: Option<String>,
+    },
     #[serde(rename = "str_replace")]
     StrReplace {
         path: String,
@@ -45,18 +52,15 @@ impl FsWrite {
         let fs = ctx.fs();
         let cwd = ctx.env().current_dir()?;
         match self {
-            FsWrite::Create { path, file_text } => {
+            FsWrite::Create { path, .. } => {
+                let file_text = self.canonical_create_command_text();
                 let path = sanitize_path_tool_arg(ctx, path);
-                let relative_path = format_path(cwd, fs.chroot_path(&path));
-                let invoke_description = if fs.exists(&path) {
-                    "Replacing the current file contents at"
-                } else {
-                    "Creating a new file at"
-                };
+                let invoke_description = if fs.exists(&path) { "Replacing: " } else { "Creating: " };
                 queue!(
                     updates,
+                    style::Print(invoke_description),
                     style::SetForegroundColor(Color::Green),
-                    style::Print(format!("{} {}", invoke_description, relative_path)),
+                    style::Print(format_path(cwd, &path)),
                     style::ResetColor,
                     style::Print("\n"),
                 )?;
@@ -69,8 +73,9 @@ impl FsWrite {
                 let matches = file.match_indices(old_str).collect::<Vec<_>>();
                 queue!(
                     updates,
+                    style::Print("Updating: "),
                     style::SetForegroundColor(Color::Green),
-                    style::Print(format!("Updating {}", format_path(cwd, &path))),
+                    style::Print(format_path(cwd, &path)),
                     style::ResetColor,
                     style::Print("\n"),
                 )?;
@@ -89,19 +94,16 @@ impl FsWrite {
                 insert_line,
                 new_str,
             } => {
+                let path = sanitize_path_tool_arg(ctx, path);
+                let mut file = fs.read_to_string(&path).await?;
                 queue!(
                     updates,
+                    style::Print("Updating: "),
                     style::SetForegroundColor(Color::Green),
-                    style::Print(format!(
-                        "Inserting at line {} in {}",
-                        insert_line,
-                        format_path(cwd, path)
-                    )),
+                    style::Print(format_path(cwd, &path)),
                     style::ResetColor,
                     style::Print("\n"),
                 )?;
-                let path = fs.chroot_path_str(path);
-                let mut file = fs.read_to_string(&path).await?;
 
                 // Get the index of the start of the line to insert at.
                 let num_lines = file.lines().enumerate().map(|(i, _)| i + 1).last().unwrap_or(1);
@@ -121,7 +123,8 @@ impl FsWrite {
     pub fn queue_description(&self, ctx: &Context, updates: &mut impl Write) -> Result<()> {
         let cwd = ctx.env().current_dir()?;
         match self {
-            FsWrite::Create { path, file_text } => {
+            FsWrite::Create { path, .. } => {
+                let file_text = self.canonical_create_command_text();
                 let relative_path = format_path(cwd, path);
                 queue!(
                     updates,
@@ -134,7 +137,7 @@ impl FsWrite {
                 if ctx.fs().exists(path) {
                     let prev = ctx.fs().read_to_string_sync(path)?;
                     let prev = stylize_output_if_able(ctx, &relative_path, prev.as_str(), None, Some("-"));
-                    let new = stylize_output_if_able(ctx, &relative_path, file_text, None, Some("+"));
+                    let new = stylize_output_if_able(ctx, &relative_path, &file_text, None, Some("+"));
                     queue!(
                         updates,
                         style::Print("Replacing:\n"),
@@ -147,7 +150,7 @@ impl FsWrite {
                         style::Print("\n\n")
                     )?;
                 } else {
-                    let file = stylize_output_if_able(ctx, &relative_path, file_text, None, None);
+                    let file = stylize_output_if_able(ctx, &relative_path, &file_text, None, None);
                     queue!(
                         updates,
                         style::Print("Contents:\n"),
@@ -223,6 +226,25 @@ impl FsWrite {
 
         Ok(())
     }
+
+    /// Returns the text to use for the [FsWrite::Create] command. This is required since we can't
+    /// rely on the model always providing `file_text`.
+    fn canonical_create_command_text(&self) -> String {
+        match self {
+            FsWrite::Create { file_text, new_str, .. } => match (file_text, new_str) {
+                (Some(file_text), _) => file_text.clone(),
+                (None, Some(new_str)) => {
+                    warn!("required field `file_text` is missing, using the provided `new_str` instead");
+                    new_str.clone()
+                },
+                _ => {
+                    warn!("no content provided for the create command");
+                    String::new()
+                },
+            },
+            _ => String::new(),
+        }
+    }
 }
 
 /// Limits the passed str to `max_len`.
@@ -358,6 +380,20 @@ mod tests {
             .unwrap();
 
         assert_eq!(ctx.fs().read_to_string("/my-file").await.unwrap(), file_text);
+
+        let file_text = "This is a new string";
+        let v = serde_json::json!({
+            "path": "/my-file",
+            "command": "create",
+            "new_str": file_text
+        });
+        serde_json::from_value::<FsWrite>(v)
+            .unwrap()
+            .invoke(&ctx, &mut stdout)
+            .await
+            .unwrap();
+
+        assert_eq!(ctx.fs().read_to_string("/my-file").await.unwrap(), file_text);
     }
 
     #[tokio::test]

crates/q_cli/src/cli/chat/tools/fs_write.rs

@@ -224,7 +224,9 @@ fn sanitize_path_tool_arg(ctx: &Context, path: impl AsRef<Path>) -> PathBuf {
         Some(p) => res.push(p),
         None => return res,
     }
-    res.push(path);
+    for p in path {
+        res.push(p);
+    }
     // For testing scenarios, we need to make sure paths are appropriately handled in chroot test
     // file systems since they are passed directly from the model.
     ctx.fs().chroot_path(res)