switch to the multi-file iteration endpoint (#1219)

* updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

fixes

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

updates

Signed-off-by: Jess Frazelle <[email protected]>

* updates

Signed-off-by: Jess Frazelle <[email protected]>

* updates

Signed-off-by: Jess Frazelle <[email protected]>

---------

Signed-off-by: Jess Frazelle <[email protected]>

Author: jessfraz

Cargo.lock

@@ -17,14 +17,15 @@ colored_json = "4.1"
 data-encoding = "2.6.0"
 dialoguer = "0.11.0"
 dirs = "6"
+futures = "0.3"
 git_rev = "0.1.0"
 heck = "0.5.0"
 http = "1"
 itertools = "0.12.1"
-kcl-lib = { version = "0.2.62", features = ["disable-println"] }
+kcl-lib = { version = "0.2.64", features = ["disable-println"] }
 #kcl-lib = { path = "../modeling-app/src/wasm-lib/kcl" }
-kcl-test-server = "0.1.62"
-kittycad = { version = "0.3.33", features = ["clap", "tabled", "requests", "retry"] }
+kcl-test-server = "0.1.64"
+kittycad = { version = "0.3.34", features = ["clap", "tabled", "requests", "retry"] }
 kittycad-modeling-cmds = { version = "0.2.107", features = ["websocket", "convert_client_crate", "tabled"] }
 log = "0.4.27"
 miette = { version = "7.5.0", features = ["fancy"] }

Cargo.toml

@@ -22,6 +22,12 @@
             rustToolchain = super.rust-bin.stable.latest.default.override {
               extensions = ["rustfmt" "clippy" "rust-src"];
             };
+
+            # stand-alone nightly formatter so we get the fancy unstable flags
+            nightlyRustfmt = super.rust-bin.selectLatestNightlyWith (toolchain:
+              toolchain.default.override {
+                extensions = ["rustfmt"]; # just the formatter
+              });
           })
         ];
         pkgs = import nixpkgs {
@@ -33,6 +39,8 @@
             nativeBuildInputs =
               (with pkgs; [
                 rustToolchain
+                nightlyRustfmt
+                cargo-sort
                 toml-cli
                 openssl
                 postgresql
@@ -41,7 +49,7 @@
               ++ pkgs.lib.optionals pkgs.stdenv.isDarwin (with pkgs; [
                 ]);
 
-            RUSTFMT = "${pkgs.rust-bin.nightly.latest.rustfmt}/bin/rustfmt";
+            RUSTFMT = "${pkgs.nightlyRustfmt}/bin/rustfmt";
           };
       }
     );

flake.nix

@@ -1,7 +1,8 @@
-use crate::config::{ConfigOption, CONFIG_OPTIONS};
 use anyhow::{bail, Result};
 use clap::Parser;
 
+use crate::config::{ConfigOption, CONFIG_OPTIONS};
+
 // TODO: make this doc a function that parses from the config the options so it's not hardcoded
 /// Manage configuration for `zoo`.
 ///
@@ -170,9 +171,7 @@ impl crate::cmd::Command for CmdConfigList {
 mod test {
     use pretty_assertions::assert_eq;
 
-    use crate::cmd::Command;
-    use crate::config;
-    use crate::config::Config;
+    use crate::{cmd::Command, config, config::Config};
 
     pub struct TestItem {
         name: String,

src/cmd_config.rs

@@ -1,14 +1,13 @@
 use std::{net::SocketAddr, str::FromStr};
 
-use crate::kcl_error_fmt;
 use anyhow::Result;
 use clap::Parser;
 use kcmc::format::OutputFormat3d as OutputFormat;
 use kittycad::types as kt;
 use kittycad_modeling_cmds as kcmc;
 use url::Url;
 
-use crate::iostreams::IoStreams;
+use crate::{iostreams::IoStreams, kcl_error_fmt};
 
 /// Perform actions on `kcl` files.
 #[derive(Parser, Debug, Clone)]
@@ -134,6 +133,7 @@ impl crate::cmd::Command for CmdKclExport {
         for file in files {
             let path = self.output_dir.join(file.name);
             std::fs::write(&path, file.contents)?;
+
             writeln!(ctx.io.out, "Wrote file: {}", path.display())?;
         }
 

src/cmd_kcl.rs

@@ -23,15 +23,18 @@ impl crate::cmd::Command for CmdKcl {
     }
 }
 
-/// Edit a `kcl` file with a prompt.
+/// Edit `kcl` file(s) with a prompt.
 ///
 ///     $ zoo ml kcl edit --prompt "Make it blue"
 ///
-/// This command outputs the edited `kcl` file to stdout.
+/// This command outputs the edited `kcl` files back to the same location.
+/// We do not output to stdout, because for projects with multiple files,
+/// it would be difficult to know which file the output corresponds to.
 #[derive(Parser, Debug, Clone)]
 #[clap(verbatim_doc_comment)]
 pub struct CmdKclEdit {
-    /// The path to the input file.
+    /// The path to the input file or directory containing a main.kcl file.
+    /// We will read in the contents of all the project's `kcl` files.
     /// If you pass `-` as the path, the file will be read from stdin.
     #[clap(name = "input", required = true)]
     pub input: std::path::PathBuf,
@@ -50,9 +53,7 @@ pub struct CmdKclEdit {
 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 (files, filepath) = ctx.collect_kcl_files(&self.input).await?;
 
         let prompt = self.prompt.join(" ");
 
@@ -61,27 +62,34 @@ impl crate::cmd::Command for CmdKclEdit {
         }
 
         let source_ranges = if let Some(source_range) = &self.source_range {
-            vec![kittycad::types::SourceRangePrompt {
+            Some(vec![kittycad::types::SourceRangePrompt {
                 range: convert_to_source_range(source_range)?,
                 prompt: prompt.clone(),
-                file: None,
-            }]
+                file: Some(filepath.to_string_lossy().to_string()),
+            }])
         } else {
             Default::default()
         };
 
-        let body = kittycad::types::TextToCadIterationBody {
-            original_source_code: input.to_string(),
-            prompt: if source_ranges.is_empty() { Some(prompt) } else { None },
+        let body = kittycad::types::TextToCadMultiFileIterationBody {
+            prompt: if source_ranges.is_none() { Some(prompt) } else { None },
             source_ranges,
             project_name: None,
             kcl_version: Some(kcl_lib::version().to_owned()),
         };
 
-        let model = ctx.get_edit_for_prompt("", &body).await?;
+        let model = ctx.get_edit_for_prompt("", &body, files).await?;
 
-        // Print the output of the conversion.
-        writeln!(ctx.io.out, "{}", model.code)?;
+        let Some(outputs) = model.outputs else {
+            anyhow::bail!("model did not return any outputs");
+        };
+
+        // Write the output to each file locally.
+        for (file, output) in outputs {
+            // We could do these in parallel...
+            tokio::fs::write(&file, output).await?;
+            writeln!(ctx.io.out, "Wrote to {}", file)?;
+        }
 
         Ok(())
     }

src/cmd_ml/cmd_kcl.rs

@@ -1,11 +1,10 @@
 use anyhow::Result;
 use clap::Parser;
 use kcl_lib::EngineManager;
-use kcmc::each_cmd as mcmd;
-use kcmc::format::InputFormat3d;
-use kcmc::ok_response::OkModelingCmdResponse;
-use kcmc::websocket::OkWebSocketResponseData;
-use kcmc::{ImageFormat, ModelingCmd};
+use kcmc::{
+    each_cmd as mcmd, format::InputFormat3d, ok_response::OkModelingCmdResponse, websocket::OkWebSocketResponseData,
+    ImageFormat, ModelingCmd,
+};
 use kittycad_modeling_cmds::{self as kcmc, ImportFile};
 
 /// Perform Text-to-CAD commands.

src/cmd_ml/cmd_text_to_cad.rs

@@ -1,11 +1,9 @@
 use std::str::FromStr;
 
 use anyhow::{anyhow, Result};
-use kcl_lib::native_engine::EngineConnection;
-use kcl_lib::EngineManager;
-use kcmc::each_cmd as mcmd;
-use kcmc::websocket::OkWebSocketResponseData;
-use kittycad::types::{ApiCallStatus, AsyncApiCallOutput, TextToCad, TextToCadCreateBody, TextToCadIteration};
+use kcl_lib::{native_engine::EngineConnection, EngineManager};
+use kcmc::{each_cmd as mcmd, websocket::OkWebSocketResponseData};
+use kittycad::types::{ApiCallStatus, AsyncApiCallOutput, TextToCad, TextToCadCreateBody, TextToCadMultiFileIteration};
 use kittycad_modeling_cmds::{self as kcmc, shared::FileExportFormat, websocket::ModelingSessionData, ModelingCmd};
 
 use crate::{config::Config, config_file::get_env_var, kcl_error_fmt, types::FormatOutput};
@@ -288,12 +286,13 @@ impl Context<'_> {
     pub async fn get_edit_for_prompt(
         &self,
         hostname: &str,
-        body: &kittycad::types::TextToCadIterationBody,
-    ) -> Result<TextToCadIteration> {
+        body: &kittycad::types::TextToCadMultiFileIterationBody,
+        files: Vec<kittycad::types::multipart::Attachment>,
+    ) -> Result<TextToCadMultiFileIteration> {
         let client = self.api_client(hostname)?;
 
         // Create the text-to-cad request.
-        let mut gen_model = client.ml().create_text_to_cad_iteration(body).await?;
+        let mut gen_model = client.ml().create_text_to_cad_multi_file_iteration(files, body).await?;
 
         // Poll until the model is ready.
         let mut status = gen_model.status.clone();
@@ -308,7 +307,7 @@ impl Context<'_> {
             // Poll for the status.
             let result = client.api_calls().get_async_operation(gen_model.id).await?;
 
-            if let AsyncApiCallOutput::TextToCadIteration {
+            if let AsyncApiCallOutput::TextToCadMultiFileIteration {
                 completed_at,
                 created_at,
                 error,
@@ -320,13 +319,12 @@ impl Context<'_> {
                 status,
                 updated_at,
                 user_id,
-                code,
                 model,
-                original_source_code,
                 source_ranges,
+                outputs,
             } = result
             {
-                gen_model = TextToCadIteration {
+                gen_model = TextToCadMultiFileIteration {
                     completed_at,
                     created_at,
                     error,
@@ -338,10 +336,9 @@ impl Context<'_> {
                     status,
                     updated_at,
                     user_id,
-                    code,
                     model,
-                    original_source_code,
                     source_ranges,
+                    outputs,
                 };
             } else {
                 anyhow::bail!("Unexpected response type: {:?}", result);
@@ -475,6 +472,75 @@ impl Context<'_> {
         let code = std::str::from_utf8(&b)?;
         Ok((code.to_string(), path))
     }
+
+    /// Collect all the kcl files in the given directory or parent directory to the given path.
+    pub async fn collect_kcl_files(
+        &mut self,
+        path: &std::path::Path,
+    ) -> Result<(Vec<kittycad::types::multipart::Attachment>, std::path::PathBuf)> {
+        let mut files = Vec::new();
+
+        let (code, filepath) = self.get_code_and_file_path(path).await?;
+        files.push(kittycad::types::multipart::Attachment {
+            name: filepath.to_string_lossy().to_string(),
+            filepath: Some(filepath.clone()),
+            content_type: Some("text/plain".parse()?),
+            data: code.as_bytes().to_vec(),
+        });
+
+        // Walk the directory and collect all the kcl files.
+        let parent = filepath
+            .parent()
+            .ok_or_else(|| anyhow!("Could not get parent directory to: `{}`", filepath.display()))?;
+        let walked_kcl = kcl_lib::walk_dir(&parent.to_path_buf()).await?;
+
+        // Get all the attachements async.
+        let futures = walked_kcl
+            .into_iter()
+            .filter(|file| *file != filepath)
+            .map(|file| {
+                tokio::spawn(async move {
+                    let contents = tokio::fs::read(&file)
+                        .await
+                        .map_err(|err| anyhow::anyhow!("Failed to read file `{}`: {:?}", file.display(), err))?;
+
+                    Ok::<kittycad::types::multipart::Attachment, anyhow::Error>(
+                        kittycad::types::multipart::Attachment {
+                            name: file.to_string_lossy().to_string(),
+                            filepath: Some(file),
+                            content_type: Some("text/plain".parse()?),
+                            data: contents,
+                        },
+                    )
+                })
+            })
+            .collect::<Vec<_>>();
+
+        // Join all futures and await their completion
+        let results = futures::future::join_all(futures).await;
+
+        // Check if any of the futures failed.
+        let mut errors = Vec::new();
+        for result in results {
+            match result {
+                Ok(Ok(attachment)) => {
+                    files.push(attachment);
+                }
+                Ok(Err(err)) => {
+                    errors.push(err);
+                }
+                Err(err) => {
+                    errors.push(anyhow::anyhow!("Failed to join future: {:?}", err));
+                }
+            }
+        }
+
+        if !errors.is_empty() {
+            anyhow::bail!("Failed to walk some kcl files: {:?}", errors);
+        }
+
+        Ok((files, filepath))
+    }
 }
 
 #[cfg(test)]

src/context.rs

@@ -1,5 +1,4 @@
-use std::io::IsTerminal;
-use std::{collections::HashMap, env, process::Command};
+use std::{collections::HashMap, env, io::IsTerminal, process::Command};
 
 use anyhow::{anyhow, Result};
 use terminal_size::{terminal_size, Height, Width};

src/iostreams.rs

@@ -767,13 +767,13 @@ async fn test_main(ctx: &mut MainContext) {
                 "ml".to_string(),
                 "kcl".to_string(),
                 "edit".to_string(),
-                "tests/gear.kcl".to_string(),
+                "tests/assembly-edit".to_string(),
                 "make".to_string(),
-                "the".to_string(),
-                "teeth".to_string(),
+                "it".to_string(),
                 "blue".to_string(),
             ],
-            want_out: r#"appearance("#.to_string(),
+            want_out: r#"Wrote to tests/assembly-edit/main.kcl"#.to_string(), // Make sure it keeps
+            // the path.
             want_err: "".to_string(),
             want_code: 0,
             ..Default::default()