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