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 cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
376
377    let cursor_buffer = project
378        .update(cx, |project, cx| {
379            project.open_buffer(
380                ProjectPath {
381                    worktree_id: worktree.read(cx).id(),
382                    path: cursor_path,
383                },
384                cx,
385            )
386        })?
387        .await?;
388
389    let cursor_offset_within_excerpt = example
390        .example
391        .cursor_position
392        .find(CURSOR_MARKER)
393        .ok_or_else(|| anyhow!("missing cursor marker"))?;
394    let mut cursor_excerpt = example.example.cursor_position.clone();
395    cursor_excerpt.replace_range(
396        cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
397        "",
398    );
399    let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
400        let text = buffer.text();
401
402        let mut matches = text.match_indices(&cursor_excerpt);
403        let Some((excerpt_offset, _)) = matches.next() else {
404            anyhow::bail!(
405                "Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
406            );
407        };
408        assert!(matches.next().is_none());
409
410        Ok(excerpt_offset)
411    })??;
412    let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
413
414    let cursor_anchor =
415        cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
416
417    let zeta = cx.update(|cx| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
418
419    zeta.update(cx, |zeta, cx| {
420        zeta.register_buffer(&cursor_buffer, &project, cx);
421    })?;
422
423    example.apply_edit_history(&project, cx).await?;
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    prediction_task.await.context("No prediction")?;
432    let mut response = None;
433
434    let mut excerpts_text = String::new();
435    while let Some(event) = debug_rx.next().await {
436        match event {
437            zeta2::ZetaDebugInfo::EditPredicted(request) => {
438                response = Some(request.response_rx.await?);
439                for included_file in request.request.included_files {
440                    let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
441                    write_codeblock(
442                        &included_file.path,
443                        included_file.excerpts.iter(),
444                        if included_file.path == request.request.excerpt_path {
445                            &insertions
446                        } else {
447                            &[]
448                        },
449                        included_file.max_row,
450                        false,
451                        &mut excerpts_text,
452                    );
453                }
454            }
455            _ => {}
456        }
457    }
458
459    println!("## Excerpts\n");
460    println!("{excerpts_text}");
461
462    println!("## Prediction\n");
463    let response = response
464        .unwrap()
465        .map(|r| r.debug_info.unwrap().model_response)
466        .unwrap_or_else(|s| s);
467    println!("{response}");
468
469    anyhow::Ok(())
470}
471
472async fn zeta2_syntax_context(
473    zeta2_args: Zeta2Args,
474    syntax_args: Zeta2SyntaxArgs,
475    args: ContextArgs,
476    app_state: &Arc<ZetaCliAppState>,
477    cx: &mut AsyncApp,
478) -> Result<String> {
479    let LoadedContext {
480        worktree,
481        project,
482        buffer,
483        clipped_cursor,
484        ..
485    } = load_context(&args, app_state, cx).await?;
486
487    // wait for worktree scan before starting zeta2 so that wait_for_initial_indexing waits for
488    // the whole worktree.
489    worktree
490        .read_with(cx, |worktree, _cx| {
491            worktree.as_local().unwrap().scan_complete()
492        })?
493        .await;
494    let output = cx
495        .update(|cx| {
496            let zeta = cx.new(|cx| {
497                zeta2::Zeta::new(app_state.client.clone(), app_state.user_store.clone(), cx)
498            });
499            let indexing_done_task = zeta.update(cx, |zeta, cx| {
500                zeta.set_options(syntax_args_to_options(&zeta2_args, &syntax_args, true));
501                zeta.register_buffer(&buffer, &project, cx);
502                zeta.wait_for_initial_indexing(&project, cx)
503            });
504            cx.spawn(async move |cx| {
505                indexing_done_task.await?;
506                let request = zeta
507                    .update(cx, |zeta, cx| {
508                        let cursor = buffer.read(cx).snapshot().anchor_before(clipped_cursor);
509                        zeta.cloud_request_for_zeta_cli(&project, &buffer, cursor, cx)
510                    })?
511                    .await?;
512
513                let (prompt_string, section_labels) = cloud_zeta2_prompt::build_prompt(&request)?;
514
515                match zeta2_args.output_format {
516                    OutputFormat::Prompt => anyhow::Ok(prompt_string),
517                    OutputFormat::Request => anyhow::Ok(serde_json::to_string_pretty(&request)?),
518                    OutputFormat::Full => anyhow::Ok(serde_json::to_string_pretty(&json!({
519                        "request": request,
520                        "prompt": prompt_string,
521                        "section_labels": section_labels,
522                    }))?),
523                }
524            })
525        })?
526        .await?;
527
528    Ok(output)
529}
530
531async fn zeta2_llm_context(
532    zeta2_args: Zeta2Args,
533    context_args: ContextArgs,
534    app_state: &Arc<ZetaCliAppState>,
535    cx: &mut AsyncApp,
536) -> Result<String> {
537    let LoadedContext {
538        buffer,
539        clipped_cursor,
540        snapshot: cursor_snapshot,
541        project,
542        ..
543    } = load_context(&context_args, app_state, cx).await?;
544
545    let cursor_position = cursor_snapshot.anchor_after(clipped_cursor);
546
547    cx.update(|cx| {
548        LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
549            registry
550                .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
551                .unwrap()
552                .authenticate(cx)
553        })
554    })?
555    .await?;
556
557    let edit_history_unified_diff = match context_args.edit_history {
558        Some(events) => events.read_to_string().await?,
559        None => String::new(),
560    };
561
562    let (debug_tx, mut debug_rx) = mpsc::unbounded();
563
564    let excerpt_options = EditPredictionExcerptOptions {
565        max_bytes: zeta2_args.max_excerpt_bytes,
566        min_bytes: zeta2_args.min_excerpt_bytes,
567        target_before_cursor_over_total_bytes: zeta2_args.target_before_cursor_over_total_bytes,
568    };
569
570    let related_excerpts = cx
571        .update(|cx| {
572            zeta2::related_excerpts::find_related_excerpts(
573                buffer,
574                cursor_position,
575                &project,
576                edit_history_unified_diff,
577                &LlmContextOptions {
578                    excerpt: excerpt_options.clone(),
579                },
580                Some(debug_tx),
581                cx,
582            )
583        })?
584        .await?;
585
586    let cursor_excerpt = EditPredictionExcerpt::select_from_buffer(
587        clipped_cursor,
588        &cursor_snapshot,
589        &excerpt_options,
590        None,
591    )
592    .context("line didn't fit")?;
593
594    #[derive(Serialize)]
595    struct Output {
596        excerpts: Vec<OutputExcerpt>,
597        formatted_excerpts: String,
598        meta: OutputMeta,
599    }
600
601    #[derive(Default, Serialize)]
602    struct OutputMeta {
603        search_prompt: String,
604        search_queries: Vec<SearchToolQuery>,
605    }
606
607    #[derive(Serialize)]
608    struct OutputExcerpt {
609        path: PathBuf,
610        #[serde(flatten)]
611        excerpt: Excerpt,
612    }
613
614    let mut meta = OutputMeta::default();
615
616    while let Some(debug_info) = debug_rx.next().await {
617        match debug_info {
618            zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
619                meta.search_prompt = info.search_prompt;
620            }
621            zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
622                meta.search_queries = info.queries
623            }
624            _ => {}
625        }
626    }
627
628    cx.update(|cx| {
629        let mut excerpts = Vec::new();
630        let mut formatted_excerpts = String::new();
631
632        let cursor_insertions = [(
633            predict_edits_v3::Point {
634                line: Line(clipped_cursor.row),
635                column: clipped_cursor.column,
636            },
637            CURSOR_MARKER,
638        )];
639
640        let mut cursor_excerpt_added = false;
641
642        for (buffer, ranges) in related_excerpts {
643            let excerpt_snapshot = buffer.read(cx).snapshot();
644
645            let mut line_ranges = ranges
646                .into_iter()
647                .map(|range| {
648                    let point_range = range.to_point(&excerpt_snapshot);
649                    Line(point_range.start.row)..Line(point_range.end.row)
650                })
651                .collect::<Vec<_>>();
652
653            let Some(file) = excerpt_snapshot.file() else {
654                continue;
655            };
656            let path = file.full_path(cx);
657
658            let is_cursor_file = path == cursor_snapshot.file().unwrap().full_path(cx);
659            if is_cursor_file {
660                let insertion_ix = line_ranges
661                    .binary_search_by(|probe| {
662                        probe
663                            .start
664                            .cmp(&cursor_excerpt.line_range.start)
665                            .then(cursor_excerpt.line_range.end.cmp(&probe.end))
666                    })
667                    .unwrap_or_else(|ix| ix);
668                line_ranges.insert(insertion_ix, cursor_excerpt.line_range.clone());
669                cursor_excerpt_added = true;
670            }
671
672            let merged_excerpts =
673                zeta2::merge_excerpts::merge_excerpts(&excerpt_snapshot, line_ranges)
674                    .into_iter()
675                    .map(|excerpt| OutputExcerpt {
676                        path: path.clone(),
677                        excerpt,
678                    });
679
680            let excerpt_start_ix = excerpts.len();
681            excerpts.extend(merged_excerpts);
682
683            write_codeblock(
684                &path,
685                excerpts[excerpt_start_ix..].iter().map(|e| &e.excerpt),
686                if is_cursor_file {
687                    &cursor_insertions
688                } else {
689                    &[]
690                },
691                Line(excerpt_snapshot.max_point().row),
692                true,
693                &mut formatted_excerpts,
694            );
695        }
696
697        if !cursor_excerpt_added {
698            write_codeblock(
699                &cursor_snapshot.file().unwrap().full_path(cx),
700                &[Excerpt {
701                    start_line: cursor_excerpt.line_range.start,
702                    text: cursor_excerpt.text(&cursor_snapshot).body.into(),
703                }],
704                &cursor_insertions,
705                Line(cursor_snapshot.max_point().row),
706                true,
707                &mut formatted_excerpts,
708            );
709        }
710
711        let output = Output {
712            excerpts,
713            formatted_excerpts,
714            meta,
715        };
716
717        Ok(serde_json::to_string_pretty(&output)?)
718    })
719    .unwrap()
720}
721
722async fn zeta1_context(
723    args: ContextArgs,
724    app_state: &Arc<ZetaCliAppState>,
725    cx: &mut AsyncApp,
726) -> Result<zeta::GatherContextOutput> {
727    let LoadedContext {
728        full_path_str,
729        snapshot,
730        clipped_cursor,
731        ..
732    } = load_context(&args, app_state, cx).await?;
733
734    let events = match args.edit_history {
735        Some(events) => events.read_to_string().await?,
736        None => String::new(),
737    };
738
739    let prompt_for_events = move || (events, 0);
740    cx.update(|cx| {
741        zeta::gather_context(
742            full_path_str,
743            &snapshot,
744            clipped_cursor,
745            prompt_for_events,
746            cx,
747        )
748    })?
749    .await
750}
751
752fn main() {
753    zlog::init();
754    zlog::init_output_stderr();
755    let args = ZetaCliArgs::parse();
756    let http_client = Arc::new(ReqwestClient::new());
757    let app = Application::headless().with_http_client(http_client);
758
759    app.run(move |cx| {
760        let app_state = Arc::new(headless::init(cx));
761        cx.spawn(async move |cx| {
762            let result = match args.command {
763                Command::Zeta1 {
764                    command: Zeta1Command::Context { context_args },
765                } => {
766                    let context = zeta1_context(context_args, &app_state, cx).await.unwrap();
767                    serde_json::to_string_pretty(&context.body).map_err(|err| anyhow::anyhow!(err))
768                }
769                Command::Zeta2 { command } => match command {
770                    Zeta2Command::Predict { example_path } => {
771                        let example = NamedExample::load(example_path).unwrap();
772                        zeta2_predict(example, &app_state, cx).await.unwrap();
773                        return;
774                    }
775                    Zeta2Command::Syntax {
776                        args,
777                        syntax_args,
778                        command,
779                    } => match command {
780                        Zeta2SyntaxCommand::Context { context_args } => {
781                            zeta2_syntax_context(args, syntax_args, context_args, &app_state, cx)
782                                .await
783                        }
784                        Zeta2SyntaxCommand::Stats {
785                            worktree,
786                            extension,
787                            limit,
788                            skip,
789                        } => {
790                            retrieval_stats(
791                                worktree,
792                                app_state,
793                                extension,
794                                limit,
795                                skip,
796                                syntax_args_to_options(&args, &syntax_args, false),
797                                cx,
798                            )
799                            .await
800                        }
801                    },
802                    Zeta2Command::Llm { args, command } => match command {
803                        Zeta2LlmCommand::Context { context_args } => {
804                            zeta2_llm_context(args, context_args, &app_state, cx).await
805                        }
806                    },
807                },
808                Command::ConvertExample {
809                    path,
810                    output_format,
811                } => {
812                    let example = NamedExample::load(path).unwrap();
813                    example.write(output_format, io::stdout()).unwrap();
814                    let _ = cx.update(|cx| cx.quit());
815                    return;
816                }
817            };
818
819            match result {
820                Ok(output) => {
821                    println!("{}", output);
822                    let _ = cx.update(|cx| cx.quit());
823                }
824                Err(e) => {
825                    eprintln!("Failed: {:?}", e);
826                    exit(1);
827                }
828            }
829        })
830        .detach();
831    });
832}