main.rs

  1mod example;
  2mod headless;
  3mod source_location;
  4mod syntax_retrieval_stats;
  5mod util;
  6
  7use crate::example::{ExampleFormat, NamedExample};
  8use crate::syntax_retrieval_stats::retrieval_stats;
  9use ::serde::Serialize;
 10use ::util::paths::PathStyle;
 11use ::util::rel_path::RelPath;
 12use anyhow::{Context as _, Result, anyhow};
 13use clap::{Args, Parser, Subcommand};
 14use cloud_llm_client::predict_edits_v3::{self, Excerpt};
 15use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
 16use edit_prediction_context::{
 17    EditPredictionContextOptions, EditPredictionExcerpt, EditPredictionExcerptOptions,
 18    EditPredictionScoreOptions, Line,
 19};
 20use futures::StreamExt as _;
 21use futures::channel::mpsc;
 22use gpui::{Application, AsyncApp, Entity, prelude::*};
 23use language::{Bias, Buffer, BufferSnapshot, OffsetRangeExt, Point};
 24use language_model::LanguageModelRegistry;
 25use project::{Project, ProjectPath, Worktree};
 26use reqwest_client::ReqwestClient;
 27use serde_json::json;
 28use std::io;
 29use std::time::Duration;
 30use std::{collections::HashSet, path::PathBuf, process::exit, str::FromStr, sync::Arc};
 31use zeta2::{ContextMode, LlmContextOptions, SearchToolQuery};
 32
 33use crate::headless::ZetaCliAppState;
 34use crate::source_location::SourceLocation;
 35use crate::util::{open_buffer, open_buffer_with_language_server};
 36
 37#[derive(Parser, Debug)]
 38#[command(name = "zeta")]
 39struct ZetaCliArgs {
 40    #[command(subcommand)]
 41    command: Command,
 42}
 43
 44#[derive(Subcommand, Debug)]
 45enum Command {
 46    Zeta1 {
 47        #[command(subcommand)]
 48        command: Zeta1Command,
 49    },
 50    Zeta2 {
 51        #[command(subcommand)]
 52        command: Zeta2Command,
 53    },
 54    ConvertExample {
 55        path: PathBuf,
 56        #[arg(long, value_enum, default_value_t = ExampleFormat::Md)]
 57        output_format: ExampleFormat,
 58    },
 59}
 60
 61#[derive(Subcommand, Debug)]
 62enum Zeta1Command {
 63    Context {
 64        #[clap(flatten)]
 65        context_args: ContextArgs,
 66    },
 67}
 68
 69#[derive(Subcommand, Debug)]
 70enum Zeta2Command {
 71    Syntax {
 72        #[clap(flatten)]
 73        args: Zeta2Args,
 74        #[clap(flatten)]
 75        syntax_args: Zeta2SyntaxArgs,
 76        #[command(subcommand)]
 77        command: Zeta2SyntaxCommand,
 78    },
 79    Llm {
 80        #[clap(flatten)]
 81        args: Zeta2Args,
 82        #[command(subcommand)]
 83        command: Zeta2LlmCommand,
 84    },
 85    Predict {
 86        example_path: PathBuf,
 87    },
 88}
 89
 90#[derive(Subcommand, Debug)]
 91enum Zeta2SyntaxCommand {
 92    Context {
 93        #[clap(flatten)]
 94        context_args: ContextArgs,
 95    },
 96    Stats {
 97        #[arg(long)]
 98        worktree: PathBuf,
 99        #[arg(long)]
100        extension: Option<String>,
101        #[arg(long)]
102        limit: Option<usize>,
103        #[arg(long)]
104        skip: Option<usize>,
105    },
106}
107
108#[derive(Subcommand, Debug)]
109enum Zeta2LlmCommand {
110    Context {
111        #[clap(flatten)]
112        context_args: ContextArgs,
113    },
114}
115
116#[derive(Debug, Args)]
117#[group(requires = "worktree")]
118struct ContextArgs {
119    #[arg(long)]
120    worktree: PathBuf,
121    #[arg(long)]
122    cursor: SourceLocation,
123    #[arg(long)]
124    use_language_server: bool,
125    #[arg(long)]
126    edit_history: Option<FileOrStdin>,
127}
128
129#[derive(Debug, Args)]
130struct Zeta2Args {
131    #[arg(long, default_value_t = 8192)]
132    max_prompt_bytes: usize,
133    #[arg(long, default_value_t = 2048)]
134    max_excerpt_bytes: usize,
135    #[arg(long, default_value_t = 1024)]
136    min_excerpt_bytes: usize,
137    #[arg(long, default_value_t = 0.66)]
138    target_before_cursor_over_total_bytes: f32,
139    #[arg(long, default_value_t = 1024)]
140    max_diagnostic_bytes: usize,
141    #[arg(long, value_enum, default_value_t = PromptFormat::default())]
142    prompt_format: PromptFormat,
143    #[arg(long, value_enum, default_value_t = Default::default())]
144    output_format: OutputFormat,
145    #[arg(long, default_value_t = 42)]
146    file_indexing_parallelism: usize,
147}
148
149#[derive(Debug, Args)]
150struct Zeta2SyntaxArgs {
151    #[arg(long, default_value_t = false)]
152    disable_imports_gathering: bool,
153    #[arg(long, default_value_t = u8::MAX)]
154    max_retrieved_definitions: u8,
155}
156
157fn syntax_args_to_options(
158    zeta2_args: &Zeta2Args,
159    syntax_args: &Zeta2SyntaxArgs,
160    omit_excerpt_overlaps: bool,
161) -> zeta2::ZetaOptions {
162    zeta2::ZetaOptions {
163        context: ContextMode::Syntax(EditPredictionContextOptions {
164            max_retrieved_declarations: syntax_args.max_retrieved_definitions,
165            use_imports: !syntax_args.disable_imports_gathering,
166            excerpt: EditPredictionExcerptOptions {
167                max_bytes: zeta2_args.max_excerpt_bytes,
168                min_bytes: zeta2_args.min_excerpt_bytes,
169                target_before_cursor_over_total_bytes: zeta2_args
170                    .target_before_cursor_over_total_bytes,
171            },
172            score: EditPredictionScoreOptions {
173                omit_excerpt_overlaps,
174            },
175        }),
176        max_diagnostic_bytes: zeta2_args.max_diagnostic_bytes,
177        max_prompt_bytes: zeta2_args.max_prompt_bytes,
178        prompt_format: zeta2_args.prompt_format.clone().into(),
179        file_indexing_parallelism: zeta2_args.file_indexing_parallelism,
180        buffer_change_grouping_interval: Duration::ZERO,
181    }
182}
183
184#[derive(clap::ValueEnum, Default, Debug, Clone)]
185enum PromptFormat {
186    MarkedExcerpt,
187    LabeledSections,
188    OnlySnippets,
189    #[default]
190    NumberedLines,
191}
192
193impl Into<predict_edits_v3::PromptFormat> for PromptFormat {
194    fn into(self) -> predict_edits_v3::PromptFormat {
195        match self {
196            Self::MarkedExcerpt => predict_edits_v3::PromptFormat::MarkedExcerpt,
197            Self::LabeledSections => predict_edits_v3::PromptFormat::LabeledSections,
198            Self::OnlySnippets => predict_edits_v3::PromptFormat::OnlySnippets,
199            Self::NumberedLines => predict_edits_v3::PromptFormat::NumLinesUniDiff,
200        }
201    }
202}
203
204#[derive(clap::ValueEnum, Default, Debug, Clone)]
205enum OutputFormat {
206    #[default]
207    Prompt,
208    Request,
209    Full,
210}
211
212#[derive(Debug, Clone)]
213enum FileOrStdin {
214    File(PathBuf),
215    Stdin,
216}
217
218impl FileOrStdin {
219    async fn read_to_string(&self) -> Result<String, std::io::Error> {
220        match self {
221            FileOrStdin::File(path) => smol::fs::read_to_string(path).await,
222            FileOrStdin::Stdin => smol::unblock(|| std::io::read_to_string(std::io::stdin())).await,
223        }
224    }
225}
226
227impl FromStr for FileOrStdin {
228    type Err = <PathBuf as FromStr>::Err;
229
230    fn from_str(s: &str) -> Result<Self, Self::Err> {
231        match s {
232            "-" => Ok(Self::Stdin),
233            _ => Ok(Self::File(PathBuf::from_str(s)?)),
234        }
235    }
236}
237
238struct LoadedContext {
239    full_path_str: String,
240    snapshot: BufferSnapshot,
241    clipped_cursor: Point,
242    worktree: Entity<Worktree>,
243    project: Entity<Project>,
244    buffer: Entity<Buffer>,
245}
246
247async fn load_context(
248    args: &ContextArgs,
249    app_state: &Arc<ZetaCliAppState>,
250    cx: &mut AsyncApp,
251) -> Result<LoadedContext> {
252    let ContextArgs {
253        worktree: worktree_path,
254        cursor,
255        use_language_server,
256        ..
257    } = args;
258
259    let worktree_path = worktree_path.canonicalize()?;
260
261    let project = cx.update(|cx| {
262        Project::local(
263            app_state.client.clone(),
264            app_state.node_runtime.clone(),
265            app_state.user_store.clone(),
266            app_state.languages.clone(),
267            app_state.fs.clone(),
268            None,
269            cx,
270        )
271    })?;
272
273    let worktree = project
274        .update(cx, |project, cx| {
275            project.create_worktree(&worktree_path, true, cx)
276        })?
277        .await?;
278
279    let mut ready_languages = HashSet::default();
280    let (_lsp_open_handle, buffer) = if *use_language_server {
281        let (lsp_open_handle, _, buffer) = open_buffer_with_language_server(
282            project.clone(),
283            worktree.clone(),
284            cursor.path.clone(),
285            &mut ready_languages,
286            cx,
287        )
288        .await?;
289        (Some(lsp_open_handle), buffer)
290    } else {
291        let buffer =
292            open_buffer(project.clone(), worktree.clone(), cursor.path.clone(), cx).await?;
293        (None, buffer)
294    };
295
296    let full_path_str = worktree
297        .read_with(cx, |worktree, _| worktree.root_name().join(&cursor.path))?
298        .display(PathStyle::local())
299        .to_string();
300
301    let snapshot = cx.update(|cx| buffer.read(cx).snapshot())?;
302    let clipped_cursor = snapshot.clip_point(cursor.point, Bias::Left);
303    if clipped_cursor != cursor.point {
304        let max_row = snapshot.max_point().row;
305        if cursor.point.row < max_row {
306            return Err(anyhow!(
307                "Cursor position {:?} is out of bounds (line length is {})",
308                cursor.point,
309                snapshot.line_len(cursor.point.row)
310            ));
311        } else {
312            return Err(anyhow!(
313                "Cursor position {:?} is out of bounds (max row is {})",
314                cursor.point,
315                max_row
316            ));
317        }
318    }
319
320    Ok(LoadedContext {
321        full_path_str,
322        snapshot,
323        clipped_cursor,
324        worktree,
325        project,
326        buffer,
327    })
328}
329
330async fn zeta2_predict(
331    example: NamedExample,
332    app_state: &Arc<ZetaCliAppState>,
333    cx: &mut AsyncApp,
334) -> Result<()> {
335    let worktree_path = example.setup_worktree().await?;
336
337    cx.update(|cx| {
338        LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
339            registry
340                .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
341                .unwrap()
342                .authenticate(cx)
343        })
344    })?
345    .await?;
346
347    app_state
348        .client
349        .sign_in_with_optional_connect(true, cx)
350        .await?;
351
352    let project = cx.update(|cx| {
353        Project::local(
354            app_state.client.clone(),
355            app_state.node_runtime.clone(),
356            app_state.user_store.clone(),
357            app_state.languages.clone(),
358            app_state.fs.clone(),
359            None,
360            cx,
361        )
362    })?;
363
364    let worktree = project
365        .update(cx, |project, cx| {
366            project.create_worktree(&worktree_path, true, cx)
367        })?
368        .await?;
369    worktree
370        .read_with(cx, |worktree, _cx| {
371            worktree.as_local().unwrap().scan_complete()
372        })?
373        .await;
374
375    let _edited_buffers = example.apply_edit_history(&project, cx).await?;
376
377    let cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
378
379    let cursor_buffer = project
380        .update(cx, |project, cx| {
381            project.open_buffer(
382                ProjectPath {
383                    worktree_id: worktree.read(cx).id(),
384                    path: cursor_path,
385                },
386                cx,
387            )
388        })?
389        .await?;
390
391    let cursor_offset_within_excerpt = example
392        .example
393        .cursor_position
394        .find(CURSOR_MARKER)
395        .ok_or_else(|| anyhow!("missing cursor marker"))?;
396    let mut cursor_excerpt = example.example.cursor_position.clone();
397    cursor_excerpt.replace_range(
398        cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
399        "",
400    );
401    let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
402        let text = buffer.text();
403
404        let mut matches = text.match_indices(&cursor_excerpt);
405        let Some((excerpt_offset, _)) = matches.next() else {
406            anyhow::bail!(
407                "Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
408            );
409        };
410        assert!(matches.next().is_none());
411
412        Ok(excerpt_offset)
413    })??;
414    let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
415
416    let cursor_anchor =
417        cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
418
419    let zeta = cx.update(|cx| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
420
421    zeta.update(cx, |zeta, cx| {
422        zeta.register_buffer(&cursor_buffer, &project, cx);
423    })?;
424
425    let (prediction_task, mut debug_rx) = zeta.update(cx, |zeta, cx| {
426        let receiver = zeta.debug_info();
427        let prediction_task = zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx);
428        (prediction_task, receiver)
429    })?;
430
431    let mut response = None;
432
433    let mut excerpts_text = String::new();
434    while let Some(event) = debug_rx.next().await {
435        match event {
436            zeta2::ZetaDebugInfo::EditPredicted(request) => {
437                response = Some(request.response_rx.await?);
438                for included_file in request.request.included_files {
439                    let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
440                    write_codeblock(
441                        &included_file.path,
442                        included_file.excerpts.iter(),
443                        if included_file.path == request.request.excerpt_path {
444                            &insertions
445                        } else {
446                            &[]
447                        },
448                        included_file.max_row,
449                        false,
450                        &mut excerpts_text,
451                    );
452                }
453                break;
454            }
455            _ => {}
456        }
457    }
458
459    prediction_task.await.context("No prediction")?;
460
461    println!("## Excerpts\n");
462    println!("{excerpts_text}");
463
464    println!("## Prediction\n");
465    let response = response
466        .unwrap()
467        .map(|r| r.debug_info.unwrap().model_response)
468        .unwrap_or_else(|s| s);
469    println!("{response}");
470
471    anyhow::Ok(())
472}
473
474async fn zeta2_syntax_context(
475    zeta2_args: Zeta2Args,
476    syntax_args: Zeta2SyntaxArgs,
477    args: ContextArgs,
478    app_state: &Arc<ZetaCliAppState>,
479    cx: &mut AsyncApp,
480) -> Result<String> {
481    let LoadedContext {
482        worktree,
483        project,
484        buffer,
485        clipped_cursor,
486        ..
487    } = load_context(&args, app_state, cx).await?;
488
489    // wait for worktree scan before starting zeta2 so that wait_for_initial_indexing waits for
490    // the whole worktree.
491    worktree
492        .read_with(cx, |worktree, _cx| {
493            worktree.as_local().unwrap().scan_complete()
494        })?
495        .await;
496    let output = cx
497        .update(|cx| {
498            let zeta = cx.new(|cx| {
499                zeta2::Zeta::new(app_state.client.clone(), app_state.user_store.clone(), cx)
500            });
501            let indexing_done_task = zeta.update(cx, |zeta, cx| {
502                zeta.set_options(syntax_args_to_options(&zeta2_args, &syntax_args, true));
503                zeta.register_buffer(&buffer, &project, cx);
504                zeta.wait_for_initial_indexing(&project, cx)
505            });
506            cx.spawn(async move |cx| {
507                indexing_done_task.await?;
508                let request = zeta
509                    .update(cx, |zeta, cx| {
510                        let cursor = buffer.read(cx).snapshot().anchor_before(clipped_cursor);
511                        zeta.cloud_request_for_zeta_cli(&project, &buffer, cursor, cx)
512                    })?
513                    .await?;
514
515                let (prompt_string, section_labels) = cloud_zeta2_prompt::build_prompt(&request)?;
516
517                match zeta2_args.output_format {
518                    OutputFormat::Prompt => anyhow::Ok(prompt_string),
519                    OutputFormat::Request => anyhow::Ok(serde_json::to_string_pretty(&request)?),
520                    OutputFormat::Full => anyhow::Ok(serde_json::to_string_pretty(&json!({
521                        "request": request,
522                        "prompt": prompt_string,
523                        "section_labels": section_labels,
524                    }))?),
525                }
526            })
527        })?
528        .await?;
529
530    Ok(output)
531}
532
533async fn zeta2_llm_context(
534    zeta2_args: Zeta2Args,
535    context_args: ContextArgs,
536    app_state: &Arc<ZetaCliAppState>,
537    cx: &mut AsyncApp,
538) -> Result<String> {
539    let LoadedContext {
540        buffer,
541        clipped_cursor,
542        snapshot: cursor_snapshot,
543        project,
544        ..
545    } = load_context(&context_args, app_state, cx).await?;
546
547    let cursor_position = cursor_snapshot.anchor_after(clipped_cursor);
548
549    cx.update(|cx| {
550        LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
551            registry
552                .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
553                .unwrap()
554                .authenticate(cx)
555        })
556    })?
557    .await?;
558
559    let edit_history_unified_diff = match context_args.edit_history {
560        Some(events) => events.read_to_string().await?,
561        None => String::new(),
562    };
563
564    let (debug_tx, mut debug_rx) = mpsc::unbounded();
565
566    let excerpt_options = EditPredictionExcerptOptions {
567        max_bytes: zeta2_args.max_excerpt_bytes,
568        min_bytes: zeta2_args.min_excerpt_bytes,
569        target_before_cursor_over_total_bytes: zeta2_args.target_before_cursor_over_total_bytes,
570    };
571
572    let related_excerpts = cx
573        .update(|cx| {
574            zeta2::related_excerpts::find_related_excerpts(
575                buffer,
576                cursor_position,
577                &project,
578                edit_history_unified_diff,
579                &LlmContextOptions {
580                    excerpt: excerpt_options.clone(),
581                },
582                Some(debug_tx),
583                cx,
584            )
585        })?
586        .await?;
587
588    let cursor_excerpt = EditPredictionExcerpt::select_from_buffer(
589        clipped_cursor,
590        &cursor_snapshot,
591        &excerpt_options,
592        None,
593    )
594    .context("line didn't fit")?;
595
596    #[derive(Serialize)]
597    struct Output {
598        excerpts: Vec<OutputExcerpt>,
599        formatted_excerpts: String,
600        meta: OutputMeta,
601    }
602
603    #[derive(Default, Serialize)]
604    struct OutputMeta {
605        search_prompt: String,
606        search_queries: Vec<SearchToolQuery>,
607    }
608
609    #[derive(Serialize)]
610    struct OutputExcerpt {
611        path: PathBuf,
612        #[serde(flatten)]
613        excerpt: Excerpt,
614    }
615
616    let mut meta = OutputMeta::default();
617
618    while let Some(debug_info) = debug_rx.next().await {
619        match debug_info {
620            zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
621                meta.search_prompt = info.search_prompt;
622            }
623            zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
624                meta.search_queries = info.queries
625            }
626            _ => {}
627        }
628    }
629
630    cx.update(|cx| {
631        let mut excerpts = Vec::new();
632        let mut formatted_excerpts = String::new();
633
634        let cursor_insertions = [(
635            predict_edits_v3::Point {
636                line: Line(clipped_cursor.row),
637                column: clipped_cursor.column,
638            },
639            CURSOR_MARKER,
640        )];
641
642        let mut cursor_excerpt_added = false;
643
644        for (buffer, ranges) in related_excerpts {
645            let excerpt_snapshot = buffer.read(cx).snapshot();
646
647            let mut line_ranges = ranges
648                .into_iter()
649                .map(|range| {
650                    let point_range = range.to_point(&excerpt_snapshot);
651                    Line(point_range.start.row)..Line(point_range.end.row)
652                })
653                .collect::<Vec<_>>();
654
655            let Some(file) = excerpt_snapshot.file() else {
656                continue;
657            };
658            let path = file.full_path(cx);
659
660            let is_cursor_file = path == cursor_snapshot.file().unwrap().full_path(cx);
661            if is_cursor_file {
662                let insertion_ix = line_ranges
663                    .binary_search_by(|probe| {
664                        probe
665                            .start
666                            .cmp(&cursor_excerpt.line_range.start)
667                            .then(cursor_excerpt.line_range.end.cmp(&probe.end))
668                    })
669                    .unwrap_or_else(|ix| ix);
670                line_ranges.insert(insertion_ix, cursor_excerpt.line_range.clone());
671                cursor_excerpt_added = true;
672            }
673
674            let merged_excerpts =
675                zeta2::merge_excerpts::merge_excerpts(&excerpt_snapshot, line_ranges)
676                    .into_iter()
677                    .map(|excerpt| OutputExcerpt {
678                        path: path.clone(),
679                        excerpt,
680                    });
681
682            let excerpt_start_ix = excerpts.len();
683            excerpts.extend(merged_excerpts);
684
685            write_codeblock(
686                &path,
687                excerpts[excerpt_start_ix..].iter().map(|e| &e.excerpt),
688                if is_cursor_file {
689                    &cursor_insertions
690                } else {
691                    &[]
692                },
693                Line(excerpt_snapshot.max_point().row),
694                true,
695                &mut formatted_excerpts,
696            );
697        }
698
699        if !cursor_excerpt_added {
700            write_codeblock(
701                &cursor_snapshot.file().unwrap().full_path(cx),
702                &[Excerpt {
703                    start_line: cursor_excerpt.line_range.start,
704                    text: cursor_excerpt.text(&cursor_snapshot).body.into(),
705                }],
706                &cursor_insertions,
707                Line(cursor_snapshot.max_point().row),
708                true,
709                &mut formatted_excerpts,
710            );
711        }
712
713        let output = Output {
714            excerpts,
715            formatted_excerpts,
716            meta,
717        };
718
719        Ok(serde_json::to_string_pretty(&output)?)
720    })
721    .unwrap()
722}
723
724async fn zeta1_context(
725    args: ContextArgs,
726    app_state: &Arc<ZetaCliAppState>,
727    cx: &mut AsyncApp,
728) -> Result<zeta::GatherContextOutput> {
729    let LoadedContext {
730        full_path_str,
731        snapshot,
732        clipped_cursor,
733        ..
734    } = load_context(&args, app_state, cx).await?;
735
736    let events = match args.edit_history {
737        Some(events) => events.read_to_string().await?,
738        None => String::new(),
739    };
740
741    let prompt_for_events = move || (events, 0);
742    cx.update(|cx| {
743        zeta::gather_context(
744            full_path_str,
745            &snapshot,
746            clipped_cursor,
747            prompt_for_events,
748            cx,
749        )
750    })?
751    .await
752}
753
754fn main() {
755    zlog::init();
756    zlog::init_output_stderr();
757    let args = ZetaCliArgs::parse();
758    let http_client = Arc::new(ReqwestClient::new());
759    let app = Application::headless().with_http_client(http_client);
760
761    app.run(move |cx| {
762        let app_state = Arc::new(headless::init(cx));
763        cx.spawn(async move |cx| {
764            let result = match args.command {
765                Command::Zeta1 {
766                    command: Zeta1Command::Context { context_args },
767                } => {
768                    let context = zeta1_context(context_args, &app_state, cx).await.unwrap();
769                    serde_json::to_string_pretty(&context.body).map_err(|err| anyhow::anyhow!(err))
770                }
771                Command::Zeta2 { command } => match command {
772                    Zeta2Command::Predict { example_path } => {
773                        let example = NamedExample::load(example_path).unwrap();
774                        zeta2_predict(example, &app_state, cx).await.unwrap();
775                        let _ = cx.update(|cx| cx.quit());
776                        return;
777                    }
778                    Zeta2Command::Syntax {
779                        args,
780                        syntax_args,
781                        command,
782                    } => match command {
783                        Zeta2SyntaxCommand::Context { context_args } => {
784                            zeta2_syntax_context(args, syntax_args, context_args, &app_state, cx)
785                                .await
786                        }
787                        Zeta2SyntaxCommand::Stats {
788                            worktree,
789                            extension,
790                            limit,
791                            skip,
792                        } => {
793                            retrieval_stats(
794                                worktree,
795                                app_state,
796                                extension,
797                                limit,
798                                skip,
799                                syntax_args_to_options(&args, &syntax_args, false),
800                                cx,
801                            )
802                            .await
803                        }
804                    },
805                    Zeta2Command::Llm { args, command } => match command {
806                        Zeta2LlmCommand::Context { context_args } => {
807                            zeta2_llm_context(args, context_args, &app_state, cx).await
808                        }
809                    },
810                },
811                Command::ConvertExample {
812                    path,
813                    output_format,
814                } => {
815                    let example = NamedExample::load(path).unwrap();
816                    example.write(output_format, io::stdout()).unwrap();
817                    let _ = cx.update(|cx| cx.quit());
818                    return;
819                }
820            };
821
822            match result {
823                Ok(output) => {
824                    println!("{}", output);
825                    let _ = cx.update(|cx| cx.quit());
826                }
827                Err(e) => {
828                    eprintln!("Failed: {:?}", e);
829                    exit(1);
830                }
831            }
832        })
833        .detach();
834    });
835}