11use std:: str:: FromStr ;
22
33use anyhow:: { anyhow, Result } ;
4- use kcl_lib:: native_engine:: EngineConnection ;
5- use kcl_lib:: EngineManager ;
6- use kcmc:: each_cmd as mcmd;
7- use kcmc:: websocket:: OkWebSocketResponseData ;
8- use kittycad:: types:: { ApiCallStatus , AsyncApiCallOutput , TextToCad , TextToCadCreateBody , TextToCadIteration } ;
4+ use kcl_lib:: { native_engine:: EngineConnection , EngineManager } ;
5+ use kcmc:: { each_cmd as mcmd, websocket:: OkWebSocketResponseData } ;
6+ use kittycad:: types:: { ApiCallStatus , AsyncApiCallOutput , TextToCad , TextToCadCreateBody , TextToCadMultiFileIteration } ;
97use kittycad_modeling_cmds:: { self as kcmc, shared:: FileExportFormat , websocket:: ModelingSessionData , ModelingCmd } ;
108
119use crate :: { config:: Config , config_file:: get_env_var, kcl_error_fmt, types:: FormatOutput } ;
@@ -288,12 +286,13 @@ impl Context<'_> {
288286 pub async fn get_edit_for_prompt (
289287 & self ,
290288 hostname : & str ,
291- body : & kittycad:: types:: TextToCadIterationBody ,
292- ) -> Result < TextToCadIteration > {
289+ body : & kittycad:: types:: TextToCadMultiFileIterationBody ,
290+ files : Vec < kittycad:: types:: multipart:: Attachment > ,
291+ ) -> Result < TextToCadMultiFileIteration > {
293292 let client = self . api_client ( hostname) ?;
294293
295294 // Create the text-to-cad request.
296- let mut gen_model = client. ml ( ) . create_text_to_cad_iteration ( body) . await ?;
295+ let mut gen_model = client. ml ( ) . create_text_to_cad_multi_file_iteration ( files , body) . await ?;
297296
298297 // Poll until the model is ready.
299298 let mut status = gen_model. status . clone ( ) ;
@@ -308,7 +307,7 @@ impl Context<'_> {
308307 // Poll for the status.
309308 let result = client. api_calls ( ) . get_async_operation ( gen_model. id ) . await ?;
310309
311- if let AsyncApiCallOutput :: TextToCadIteration {
310+ if let AsyncApiCallOutput :: TextToCadMultiFileIteration {
312311 completed_at,
313312 created_at,
314313 error,
@@ -320,13 +319,12 @@ impl Context<'_> {
320319 status,
321320 updated_at,
322321 user_id,
323- code,
324322 model,
325- original_source_code,
326323 source_ranges,
324+ outputs,
327325 } = result
328326 {
329- gen_model = TextToCadIteration {
327+ gen_model = TextToCadMultiFileIteration {
330328 completed_at,
331329 created_at,
332330 error,
@@ -338,10 +336,9 @@ impl Context<'_> {
338336 status,
339337 updated_at,
340338 user_id,
341- code,
342339 model,
343- original_source_code,
344340 source_ranges,
341+ outputs,
345342 } ;
346343 } else {
347344 anyhow:: bail!( "Unexpected response type: {:?}" , result) ;
@@ -475,6 +472,75 @@ impl Context<'_> {
475472 let code = std:: str:: from_utf8 ( & b) ?;
476473 Ok ( ( code. to_string ( ) , path) )
477474 }
475+
476+ /// Collect all the kcl files in the given directory or parent directory to the given path.
477+ pub async fn collect_kcl_files (
478+ & mut self ,
479+ path : & std:: path:: Path ,
480+ ) -> Result < ( Vec < kittycad:: types:: multipart:: Attachment > , std:: path:: PathBuf ) > {
481+ let mut files = Vec :: new ( ) ;
482+
483+ let ( code, filepath) = self . get_code_and_file_path ( path) . await ?;
484+ files. push ( kittycad:: types:: multipart:: Attachment {
485+ name : filepath. to_string_lossy ( ) . to_string ( ) ,
486+ filepath : Some ( filepath. clone ( ) ) ,
487+ content_type : Some ( "text/plain" . parse ( ) ?) ,
488+ data : code. as_bytes ( ) . to_vec ( ) ,
489+ } ) ;
490+
491+ // Walk the directory and collect all the kcl files.
492+ let parent = filepath
493+ . parent ( )
494+ . ok_or_else ( || anyhow ! ( "Could not get parent directory to: `{}`" , filepath. display( ) ) ) ?;
495+ let walked_kcl = kcl_lib:: walk_dir ( & parent. to_path_buf ( ) ) . await ?;
496+
497+ // Get all the attachements async.
498+ let futures = walked_kcl
499+ . into_iter ( )
500+ . filter ( |file| * file != filepath)
501+ . map ( |file| {
502+ tokio:: spawn ( async move {
503+ let contents = tokio:: fs:: read ( & file)
504+ . await
505+ . map_err ( |err| anyhow:: anyhow!( "Failed to read file `{}`: {:?}" , file. display( ) , err) ) ?;
506+
507+ Ok :: < kittycad:: types:: multipart:: Attachment , anyhow:: Error > (
508+ kittycad:: types:: multipart:: Attachment {
509+ name : file. to_string_lossy ( ) . to_string ( ) ,
510+ filepath : Some ( file) ,
511+ content_type : Some ( "text/plain" . parse ( ) ?) ,
512+ data : contents,
513+ } ,
514+ )
515+ } )
516+ } )
517+ . collect :: < Vec < _ > > ( ) ;
518+
519+ // Join all futures and await their completion
520+ let results = futures:: future:: join_all ( futures) . await ;
521+
522+ // Check if any of the futures failed.
523+ let mut errors = Vec :: new ( ) ;
524+ for result in results {
525+ match result {
526+ Ok ( Ok ( attachment) ) => {
527+ files. push ( attachment) ;
528+ }
529+ Ok ( Err ( err) ) => {
530+ errors. push ( err) ;
531+ }
532+ Err ( err) => {
533+ errors. push ( anyhow:: anyhow!( "Failed to join future: {:?}" , err) ) ;
534+ }
535+ }
536+ }
537+
538+ if !errors. is_empty ( ) {
539+ anyhow:: bail!( "Failed to walk some kcl files: {:?}" , errors) ;
540+ }
541+
542+ Ok ( ( files, filepath) )
543+ }
478544}
479545
480546#[ cfg( test) ]
0 commit comments