Skip to content

Commit

Permalink
Add kcl edit to cli (#1083)
Browse files Browse the repository at this point in the history
* add rtests

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* tests

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* text-to-kcl edit in cli

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

* updates

Signed-off-by: Jess Frazelle <github@jessfraz.com>

---------

Signed-off-by: Jess Frazelle <github@jessfraz.com>
  • Loading branch information
jessfraz authored Dec 17, 2024
1 parent 99bb205 commit 86e8491
Show file tree
Hide file tree
Showing 11 changed files with 152,272 additions and 8,246 deletions.
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
@@ -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
Expand Down
162 changes: 162 additions & 0 deletions src/cmd_ml/cmd_kcl.rs
Original file line number Diff line number Diff line change
@@ -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 }
}
);
}
}
110 changes: 94 additions & 16 deletions src/cmd_ml/cmd_text_to_cad.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)]
Expand Down Expand Up @@ -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"
);
}

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
4 changes: 4 additions & 0 deletions src/cmd_ml/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use anyhow::Result;
use clap::Parser;

/// Kcl commands.
mod cmd_kcl;
/// Text-to-CAD commands.
mod cmd_text_to_cad;

Expand All @@ -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,
}
}
}
Loading

0 comments on commit 86e8491

Please sign in to comment.