Add kcl edit to cli

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zoo"
-version = "0.2.91"
+version = "0.2.92"
 edition = "2021"
 build = "build.rs"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

src/cmd_ml/cmd_kcl.rs

@@ -0,0 +1,162 @@
+use anyhow::Result;
+use clap::Parser;
+
+/// Edit a KCL file with machine learning.
+#[derive(Parser, Debug, Clone)]
+#[clap(verbatim_doc_comment)]
+pub struct CmdKcl {
+    #[clap(subcommand)]
+    subcmd: SubCommand,
+}
+
+#[derive(Parser, Debug, Clone)]
+enum SubCommand {
+    Edit(CmdKclEdit),
+}
+
+#[async_trait::async_trait(?Send)]
+impl crate::cmd::Command for CmdKcl {
+    async fn run(&self, ctx: &mut crate::context::Context) -> Result<()> {
+        match &self.subcmd {
+            SubCommand::Edit(cmd) => cmd.run(ctx).await,
+        }
+    }
+}
+
+/// Edit a `kcl` file with a prompt.
+///
+///     $ zoo ml kcl edit --prompt "Make it blue"
+///
+/// This command outputs the edited `kcl` file to stdout.
+#[derive(Parser, Debug, Clone)]
+#[clap(verbatim_doc_comment)]
+pub struct CmdKclEdit {
+    /// The path to the input file.
+    /// If you pass `-` as the path, the file will be read from stdin.
+    #[clap(name = "input", required = true)]
+    pub input: std::path::PathBuf,
+
+    /// Your prompt.
+    #[clap(name = "prompt", required = true)]
+    pub prompt: Vec<String>,
+
+    /// The source ranges to edit. This is optional.
+    /// If you don't pass this, the entire file will be edited.
+    #[clap(name = "source_range", long, short = 'r')]
+    pub source_range: Option<String>,
+}
+
+#[async_trait::async_trait(?Send)]
+impl crate::cmd::Command for CmdKclEdit {
+    async fn run(&self, ctx: &mut crate::context::Context) -> Result<()> {
+        // Get the contents of the input file.
+        let input = ctx.read_file(self.input.to_str().unwrap_or(""))?;
+        // Parse the input as a string.
+        let input = std::str::from_utf8(&input)?;
+
+        let prompt = self.prompt.join(" ");
+
+        if prompt.is_empty() {
+            anyhow::bail!("prompt cannot be empty");
+        }
+
+        let source_ranges = if let Some(source_range) = &self.source_range {
+            vec![kittycad::types::SourceRangePrompt {
+                range: convert_to_source_range(source_range)?,
+                prompt: prompt.clone(),
+            }]
+        } else {
+            Default::default()
+        };
+
+        let body = kittycad::types::TextToCadIterationBody {
+            original_source_code: input.to_string(),
+            prompt: if source_ranges.is_empty() { Some(prompt) } else { None },
+            source_ranges,
+        };
+
+        let model = ctx.get_edit_for_prompt("", &body).await?;
+
+        // Print the output of the conversion.
+        writeln!(ctx.io.out, "{}", model.code)?;
+
+        Ok(())
+    }
+}
+
+/// Convert from a string like "4:2-4:5" to a source range.
+/// Where 4 is the line number and 2 and 5 are the column numbers.
+fn convert_to_source_range(source_range: &str) -> Result<kittycad::types::SourceRange> {
+    let parts: Vec<&str> = source_range.split('-').collect();
+    if parts.len() != 2 {
+        anyhow::bail!("source range must be in the format 'line:column-line:column'");
+    }
+
+    let inner_parts_start = parts[0].split(':').collect::<Vec<&str>>();
+    if inner_parts_start.len() != 2 {
+        anyhow::bail!("source range must be in the format 'line:column'");
+    }
+
+    let inner_parts_end = parts[1].split(':').collect::<Vec<&str>>();
+    if inner_parts_end.len() != 2 {
+        anyhow::bail!("source range must be in the format 'line:column'");
+    }
+
+    let start = kittycad::types::SourcePosition {
+        line: inner_parts_start[0].parse::<u32>()?,
+        column: inner_parts_start[1].parse::<u32>()?,
+    };
+    let end = kittycad::types::SourcePosition {
+        line: inner_parts_end[0].parse::<u32>()?,
+        column: inner_parts_end[1].parse::<u32>()?,
+    };
+
+    Ok(kittycad::types::SourceRange { start, end })
+}
+
+#[cfg(test)]
+mod test {
+    use pretty_assertions::assert_eq;
+
+    use super::*;
+
+    #[test]
+    fn test_convert_to_source_range() {
+        let source_range = "4:2-4:5";
+        let result = convert_to_source_range(source_range).unwrap();
+        assert_eq!(
+            result,
+            kittycad::types::SourceRange {
+                start: kittycad::types::SourcePosition { line: 4, column: 2 },
+                end: kittycad::types::SourcePosition { line: 4, column: 5 }
+            }
+        );
+    }
+
+    #[test]
+    fn test_convert_to_source_range_invalid() {
+        let source_range = "4:2-4";
+        let result = convert_to_source_range(source_range);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_convert_to_source_range_invalid_inner() {
+        let source_range = "4:2-4:5:6";
+        let result = convert_to_source_range(source_range);
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_convert_to_source_range_bigger() {
+        let source_range = "14:12-15:25";
+        let result = convert_to_source_range(source_range).unwrap();
+        assert_eq!(
+            result,
+            kittycad::types::SourceRange {
+                start: kittycad::types::SourcePosition { line: 14, column: 12 },
+                end: kittycad::types::SourcePosition { line: 15, column: 25 }
+            }
+        );
+    }
+}

src/cmd_ml/cmd_text_to_cad.rs

@@ -34,6 +34,59 @@ impl crate::cmd::Command for CmdTextToCad {
     }
 }
 
+#[doc = "The valid types of output file formats."]
+#[derive(serde :: Serialize, serde :: Deserialize, PartialEq, Hash, Debug, Clone, clap::ValueEnum)]
+pub enum FileExportFormat {
+    /// KCL file format. <https://kittycad.com/docs/kcl/>
+    #[serde(rename = "kcl")]
+    Kcl,
+    #[doc = "Autodesk Filmbox (FBX) format. <https://en.wikipedia.org/wiki/FBX>"]
+    #[serde(rename = "fbx")]
+    Fbx,
+    #[doc = "Binary glTF 2.0.\n\nThis is a single binary with .glb extension.\n\nThis is better \
+             if you want a compressed format as opposed to the human readable glTF that lacks \
+             compression."]
+    #[serde(rename = "glb")]
+    Glb,
+    #[doc = "glTF 2.0. Embedded glTF 2.0 (pretty printed).\n\nSingle JSON file with .gltf \
+             extension binary data encoded as base64 data URIs.\n\nThe JSON contents are pretty \
+             printed.\n\nIt is human readable, single file, and you can view the diff easily in a \
+             git commit."]
+    #[serde(rename = "gltf")]
+    Gltf,
+    #[doc = "The OBJ file format. <https://en.wikipedia.org/wiki/Wavefront_.obj_file> It may or \
+             may not have an an attached material (mtl // mtllib) within the file, but we \
+             interact with it as if it does not."]
+    #[serde(rename = "obj")]
+    Obj,
+    #[doc = "The PLY file format. <https://en.wikipedia.org/wiki/PLY_(file_format)>"]
+    #[serde(rename = "ply")]
+    Ply,
+    #[doc = "The STEP file format. <https://en.wikipedia.org/wiki/ISO_10303-21>"]
+    #[serde(rename = "step")]
+    Step,
+    #[doc = "The STL file format. <https://en.wikipedia.org/wiki/STL_(file_format)>"]
+    #[serde(rename = "stl")]
+    Stl,
+}
+
+impl TryFrom<FileExportFormat> for kittycad::types::FileExportFormat {
+    type Error = anyhow::Error;
+
+    fn try_from(value: FileExportFormat) -> Result<Self> {
+        match value {
+            FileExportFormat::Kcl => anyhow::bail!("KCL file format is not supported"),
+            FileExportFormat::Fbx => Ok(kittycad::types::FileExportFormat::Fbx),
+            FileExportFormat::Glb => Ok(kittycad::types::FileExportFormat::Glb),
+            FileExportFormat::Gltf => Ok(kittycad::types::FileExportFormat::Gltf),
+            FileExportFormat::Obj => Ok(kittycad::types::FileExportFormat::Obj),
+            FileExportFormat::Ply => Ok(kittycad::types::FileExportFormat::Ply),
+            FileExportFormat::Step => Ok(kittycad::types::FileExportFormat::Step),
+            FileExportFormat::Stl => Ok(kittycad::types::FileExportFormat::Stl),
+        }
+    }
+}
+
 /// Run a Text-to-CAD prompt and export it as any other supported CAD file format.
 ///
 ///     $ zoo ml text-to-cad export --output-format=obj A 2x4 lego brick
@@ -51,7 +104,7 @@ pub struct CmdTextToCadExport {
 
     /// A valid output file format.
     #[clap(short = 't', long = "output-format", value_enum)]
-    output_format: kittycad::types::FileExportFormat,
+    output_format: FileExportFormat,
 
     /// Command output format.
     #[clap(long, short, value_enum)]
@@ -82,24 +135,49 @@ impl crate::cmd::Command for CmdTextToCadExport {
         }
 
         let mut model = ctx
-            .get_model_for_prompt("", &prompt, self.output_format.clone())
+            .get_model_for_prompt(
+                "",
+                &prompt,
+                self.output_format == FileExportFormat::Kcl,
+                if self.output_format == FileExportFormat::Kcl {
+                    kittycad::types::FileExportFormat::Gltf
+                } else {
+                    self.output_format.clone().try_into()?
+                },
+            )
             .await?;
 
-        if let Some(outputs) = model.outputs {
-            // Write the contents of the files to the output directory.
-            for (filename, data) in outputs.iter() {
-                let path = output_dir.clone().join(filename);
-                std::fs::write(&path, data)?;
-                writeln!(
-                    ctx.io.out,
-                    "wrote file `{}` to {}",
-                    filename,
-                    path.to_str().unwrap_or("")
-                )?;
+        if self.output_format != FileExportFormat::Kcl {
+            if let Some(outputs) = model.outputs {
+                // Write the contents of the files to the output directory.
+                for (filename, data) in outputs.iter() {
+                    let path = output_dir.clone().join(filename);
+                    std::fs::write(&path, data)?;
+                    writeln!(
+                        ctx.io.out,
+                        "wrote file `{}` to {}",
+                        filename,
+                        path.to_str().unwrap_or("")
+                    )?;
+                }
+            } else {
+                anyhow::bail!(
+                "no output was generated! (this is probably a bug in the API) you should report it to support@zoo.dev"
+            );
             }
+        } else if let Some(code) = &model.code {
+            let filename = prompt.replace(" ", "_").to_lowercase() + ".kcl";
+            let path = output_dir.clone().join(&filename);
+            std::fs::write(&path, code)?;
+            writeln!(
+                ctx.io.out,
+                "wrote file `{}` to {}",
+                filename,
+                path.to_str().unwrap_or("")
+            )?;
         } else {
             anyhow::bail!(
-                "no output was generated! (this is probably a bug in the API) you should report it to support@zoo.dev"
+                "no code was generated! (this is probably a bug in the API) you should report it to support@zoo.dev"
             );
         }
 
@@ -164,7 +242,7 @@ impl crate::cmd::Command for CmdTextToCadSnapshot {
         }
 
         let model = ctx
-            .get_model_for_prompt("", &prompt, kittycad::types::FileExportFormat::Gltf)
+            .get_model_for_prompt("", &prompt, false, kittycad::types::FileExportFormat::Gltf)
             .await?;
 
         // Get the gltf bytes.
@@ -230,7 +308,7 @@ impl crate::cmd::Command for CmdTextToCadView {
         }
 
         let model = ctx
-            .get_model_for_prompt("", &prompt, kittycad::types::FileExportFormat::Gltf)
+            .get_model_for_prompt("", &prompt, false, kittycad::types::FileExportFormat::Gltf)
             .await?;
 
         // Get the gltf bytes.

src/cmd_ml/mod.rs

@@ -1,6 +1,8 @@
 use anyhow::Result;
 use clap::Parser;
 
+/// Kcl commands.
+mod cmd_kcl;
 /// Text-to-CAD commands.
 mod cmd_text_to_cad;
 
@@ -16,13 +18,15 @@ pub struct CmdMl {
 enum SubCommand {
     #[clap(name = "text-to-cad")]
     TextToCad(crate::cmd_ml::cmd_text_to_cad::CmdTextToCad),
+    Kcl(crate::cmd_ml::cmd_kcl::CmdKcl),
 }
 
 #[async_trait::async_trait(?Send)]
 impl crate::cmd::Command for CmdMl {
     async fn run(&self, ctx: &mut crate::context::Context) -> Result<()> {
         match &self.subcmd {
             SubCommand::TextToCad(cmd) => cmd.run(ctx).await,
+            SubCommand::Kcl(cmd) => cmd.run(ctx).await,
         }
     }
 }