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, Instant};
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
415 let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
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 let refresh_task = zeta.update(cx, |zeta, cx| {
422 zeta.register_buffer(&cursor_buffer, &project, cx);
423 zeta.refresh_context(project.clone(), cursor_buffer.clone(), cursor_anchor, cx)
424 })?;
425
426 let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info())?;
427 let mut context_retrieval_started_at = None;
428 let mut context_retrieval_finished_at = None;
429 let mut search_queries_generated_at = None;
430 let mut search_queries_executed_at = None;
431 let mut prediction_started_at = None;
432 let mut prediction_finished_at = None;
433 let mut excerpts_text = String::new();
434 let mut prediction_task = None;
435 while let Some(event) = debug_rx.next().await {
436 match event {
437 zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
438 context_retrieval_started_at = Some(info.timestamp);
439 }
440 zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
441 search_queries_generated_at = Some(info.timestamp);
442 }
443 zeta2::ZetaDebugInfo::SearchQueriesExecuted(info) => {
444 search_queries_executed_at = Some(info.timestamp);
445 }
446 zeta2::ZetaDebugInfo::ContextRetrievalFinished(info) => {
447 context_retrieval_finished_at = Some(info.timestamp);
448
449 prediction_task = Some(zeta.update(cx, |zeta, cx| {
450 zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx)
451 })?);
452 }
453 zeta2::ZetaDebugInfo::EditPredicted(request) => {
454 prediction_started_at = Some(Instant::now());
455 request.response_rx.await?.map_err(|err| anyhow!(err))?;
456 prediction_finished_at = Some(Instant::now());
457
458 for included_file in request.request.included_files {
459 let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
460 write_codeblock(
461 &included_file.path,
462 included_file.excerpts.iter(),
463 if included_file.path == request.request.excerpt_path {
464 &insertions
465 } else {
466 &[]
467 },
468 included_file.max_row,
469 false,
470 &mut excerpts_text,
471 );
472 }
473 break;
474 }
475 _ => {}
476 }
477 }
478
479 refresh_task.await.context("context retrieval failed")?;
480 let prediction = prediction_task.unwrap().await?.context("No prediction")?;
481
482 println!("## Excerpts\n");
483 println!("{excerpts_text}");
484
485 let old_text = prediction.snapshot.text();
486 let new_text = prediction.buffer.update(cx, |buffer, cx| {
487 buffer.edit(prediction.edits.iter().cloned(), None, cx);
488 buffer.text()
489 })?;
490 let diff = language::unified_diff(&old_text, &new_text);
491
492 println!("## Prediction\n");
493 println!("{diff}");
494
495 println!("## Time\n");
496
497 let planning_search_time =
498 search_queries_generated_at.unwrap() - context_retrieval_started_at.unwrap();
499
500 println!("Planning searches: {}ms", planning_search_time.as_millis());
501 println!(
502 "Running searches: {}ms",
503 (search_queries_executed_at.unwrap() - search_queries_generated_at.unwrap()).as_millis()
504 );
505
506 let filtering_search_time =
507 context_retrieval_finished_at.unwrap() - search_queries_executed_at.unwrap();
508 println!(
509 "Filtering context results: {}ms",
510 filtering_search_time.as_millis()
511 );
512
513 let prediction_time = prediction_finished_at.unwrap() - prediction_started_at.unwrap();
514 println!("Making Prediction: {}ms", prediction_time.as_millis());
515
516 println!("-------------------");
517 let total_time =
518 (prediction_finished_at.unwrap() - context_retrieval_started_at.unwrap()).as_millis();
519 println!("Total: {}ms", total_time);
520
521 let inference_time =
522 (planning_search_time + filtering_search_time + prediction_time).as_millis();
523 println!(
524 "Inference: {}ms ({:.2}%)",
525 inference_time,
526 (inference_time as f64 / total_time as f64) * 100.
527 );
528
529 anyhow::Ok(())
530}
531
532async fn zeta2_syntax_context(
533 zeta2_args: Zeta2Args,
534 syntax_args: Zeta2SyntaxArgs,
535 args: ContextArgs,
536 app_state: &Arc<ZetaCliAppState>,
537 cx: &mut AsyncApp,
538) -> Result<String> {
539 let LoadedContext {
540 worktree,
541 project,
542 buffer,
543 clipped_cursor,
544 ..
545 } = load_context(&args, app_state, cx).await?;
546
547 // wait for worktree scan before starting zeta2 so that wait_for_initial_indexing waits for
548 // the whole worktree.
549 worktree
550 .read_with(cx, |worktree, _cx| {
551 worktree.as_local().unwrap().scan_complete()
552 })?
553 .await;
554 let output = cx
555 .update(|cx| {
556 let zeta = cx.new(|cx| {
557 zeta2::Zeta::new(app_state.client.clone(), app_state.user_store.clone(), cx)
558 });
559 let indexing_done_task = zeta.update(cx, |zeta, cx| {
560 zeta.set_options(syntax_args_to_options(&zeta2_args, &syntax_args, true));
561 zeta.register_buffer(&buffer, &project, cx);
562 zeta.wait_for_initial_indexing(&project, cx)
563 });
564 cx.spawn(async move |cx| {
565 indexing_done_task.await?;
566 let request = zeta
567 .update(cx, |zeta, cx| {
568 let cursor = buffer.read(cx).snapshot().anchor_before(clipped_cursor);
569 zeta.cloud_request_for_zeta_cli(&project, &buffer, cursor, cx)
570 })?
571 .await?;
572
573 let (prompt_string, section_labels) = cloud_zeta2_prompt::build_prompt(&request)?;
574
575 match zeta2_args.output_format {
576 OutputFormat::Prompt => anyhow::Ok(prompt_string),
577 OutputFormat::Request => anyhow::Ok(serde_json::to_string_pretty(&request)?),
578 OutputFormat::Full => anyhow::Ok(serde_json::to_string_pretty(&json!({
579 "request": request,
580 "prompt": prompt_string,
581 "section_labels": section_labels,
582 }))?),
583 }
584 })
585 })?
586 .await?;
587
588 Ok(output)
589}
590
591async fn zeta2_llm_context(
592 zeta2_args: Zeta2Args,
593 context_args: ContextArgs,
594 app_state: &Arc<ZetaCliAppState>,
595 cx: &mut AsyncApp,
596) -> Result<String> {
597 let LoadedContext {
598 buffer,
599 clipped_cursor,
600 snapshot: cursor_snapshot,
601 project,
602 ..
603 } = load_context(&context_args, app_state, cx).await?;
604
605 let cursor_position = cursor_snapshot.anchor_after(clipped_cursor);
606
607 cx.update(|cx| {
608 LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
609 registry
610 .provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
611 .unwrap()
612 .authenticate(cx)
613 })
614 })?
615 .await?;
616
617 let edit_history_unified_diff = match context_args.edit_history {
618 Some(events) => events.read_to_string().await?,
619 None => String::new(),
620 };
621
622 let (debug_tx, mut debug_rx) = mpsc::unbounded();
623
624 let excerpt_options = EditPredictionExcerptOptions {
625 max_bytes: zeta2_args.max_excerpt_bytes,
626 min_bytes: zeta2_args.min_excerpt_bytes,
627 target_before_cursor_over_total_bytes: zeta2_args.target_before_cursor_over_total_bytes,
628 };
629
630 let related_excerpts = cx
631 .update(|cx| {
632 zeta2::related_excerpts::find_related_excerpts(
633 buffer,
634 cursor_position,
635 &project,
636 edit_history_unified_diff,
637 &LlmContextOptions {
638 excerpt: excerpt_options.clone(),
639 },
640 Some(debug_tx),
641 cx,
642 )
643 })?
644 .await?;
645
646 let cursor_excerpt = EditPredictionExcerpt::select_from_buffer(
647 clipped_cursor,
648 &cursor_snapshot,
649 &excerpt_options,
650 None,
651 )
652 .context("line didn't fit")?;
653
654 #[derive(Serialize)]
655 struct Output {
656 excerpts: Vec<OutputExcerpt>,
657 formatted_excerpts: String,
658 meta: OutputMeta,
659 }
660
661 #[derive(Default, Serialize)]
662 struct OutputMeta {
663 search_prompt: String,
664 search_queries: Vec<SearchToolQuery>,
665 }
666
667 #[derive(Serialize)]
668 struct OutputExcerpt {
669 path: PathBuf,
670 #[serde(flatten)]
671 excerpt: Excerpt,
672 }
673
674 let mut meta = OutputMeta::default();
675
676 while let Some(debug_info) = debug_rx.next().await {
677 match debug_info {
678 zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
679 meta.search_prompt = info.search_prompt;
680 }
681 zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
682 meta.search_queries = info.queries
683 }
684 _ => {}
685 }
686 }
687
688 cx.update(|cx| {
689 let mut excerpts = Vec::new();
690 let mut formatted_excerpts = String::new();
691
692 let cursor_insertions = [(
693 predict_edits_v3::Point {
694 line: Line(clipped_cursor.row),
695 column: clipped_cursor.column,
696 },
697 CURSOR_MARKER,
698 )];
699
700 let mut cursor_excerpt_added = false;
701
702 for (buffer, ranges) in related_excerpts {
703 let excerpt_snapshot = buffer.read(cx).snapshot();
704
705 let mut line_ranges = ranges
706 .into_iter()
707 .map(|range| {
708 let point_range = range.to_point(&excerpt_snapshot);
709 Line(point_range.start.row)..Line(point_range.end.row)
710 })
711 .collect::<Vec<_>>();
712
713 let Some(file) = excerpt_snapshot.file() else {
714 continue;
715 };
716 let path = file.full_path(cx);
717
718 let is_cursor_file = path == cursor_snapshot.file().unwrap().full_path(cx);
719 if is_cursor_file {
720 let insertion_ix = line_ranges
721 .binary_search_by(|probe| {
722 probe
723 .start
724 .cmp(&cursor_excerpt.line_range.start)
725 .then(cursor_excerpt.line_range.end.cmp(&probe.end))
726 })
727 .unwrap_or_else(|ix| ix);
728 line_ranges.insert(insertion_ix, cursor_excerpt.line_range.clone());
729 cursor_excerpt_added = true;
730 }
731
732 let merged_excerpts =
733 zeta2::merge_excerpts::merge_excerpts(&excerpt_snapshot, line_ranges)
734 .into_iter()
735 .map(|excerpt| OutputExcerpt {
736 path: path.clone(),
737 excerpt,
738 });
739
740 let excerpt_start_ix = excerpts.len();
741 excerpts.extend(merged_excerpts);
742
743 write_codeblock(
744 &path,
745 excerpts[excerpt_start_ix..].iter().map(|e| &e.excerpt),
746 if is_cursor_file {
747 &cursor_insertions
748 } else {
749 &[]
750 },
751 Line(excerpt_snapshot.max_point().row),
752 true,
753 &mut formatted_excerpts,
754 );
755 }
756
757 if !cursor_excerpt_added {
758 write_codeblock(
759 &cursor_snapshot.file().unwrap().full_path(cx),
760 &[Excerpt {
761 start_line: cursor_excerpt.line_range.start,
762 text: cursor_excerpt.text(&cursor_snapshot).body.into(),
763 }],
764 &cursor_insertions,
765 Line(cursor_snapshot.max_point().row),
766 true,
767 &mut formatted_excerpts,
768 );
769 }
770
771 let output = Output {
772 excerpts,
773 formatted_excerpts,
774 meta,
775 };
776
777 Ok(serde_json::to_string_pretty(&output)?)
778 })
779 .unwrap()
780}
781
782async fn zeta1_context(
783 args: ContextArgs,
784 app_state: &Arc<ZetaCliAppState>,
785 cx: &mut AsyncApp,
786) -> Result<zeta::GatherContextOutput> {
787 let LoadedContext {
788 full_path_str,
789 snapshot,
790 clipped_cursor,
791 ..
792 } = load_context(&args, app_state, cx).await?;
793
794 let events = match args.edit_history {
795 Some(events) => events.read_to_string().await?,
796 None => String::new(),
797 };
798
799 let prompt_for_events = move || (events, 0);
800 cx.update(|cx| {
801 zeta::gather_context(
802 full_path_str,
803 &snapshot,
804 clipped_cursor,
805 prompt_for_events,
806 cx,
807 )
808 })?
809 .await
810}
811
812fn main() {
813 zlog::init();
814 zlog::init_output_stderr();
815 let args = ZetaCliArgs::parse();
816 let http_client = Arc::new(ReqwestClient::new());
817 let app = Application::headless().with_http_client(http_client);
818
819 app.run(move |cx| {
820 let app_state = Arc::new(headless::init(cx));
821 cx.spawn(async move |cx| {
822 let result = match args.command {
823 Command::Zeta1 {
824 command: Zeta1Command::Context { context_args },
825 } => {
826 let context = zeta1_context(context_args, &app_state, cx).await.unwrap();
827 serde_json::to_string_pretty(&context.body).map_err(|err| anyhow::anyhow!(err))
828 }
829 Command::Zeta2 { command } => match command {
830 Zeta2Command::Predict { example_path } => {
831 let example = NamedExample::load(example_path).unwrap();
832 zeta2_predict(example, &app_state, cx).await.unwrap();
833 let _ = cx.update(|cx| cx.quit());
834 return;
835 }
836 Zeta2Command::Syntax {
837 args,
838 syntax_args,
839 command,
840 } => match command {
841 Zeta2SyntaxCommand::Context { context_args } => {
842 zeta2_syntax_context(args, syntax_args, context_args, &app_state, cx)
843 .await
844 }
845 Zeta2SyntaxCommand::Stats {
846 worktree,
847 extension,
848 limit,
849 skip,
850 } => {
851 retrieval_stats(
852 worktree,
853 app_state,
854 extension,
855 limit,
856 skip,
857 syntax_args_to_options(&args, &syntax_args, false),
858 cx,
859 )
860 .await
861 }
862 },
863 Zeta2Command::Llm { args, command } => match command {
864 Zeta2LlmCommand::Context { context_args } => {
865 zeta2_llm_context(args, context_args, &app_state, cx).await
866 }
867 },
868 },
869 Command::ConvertExample {
870 path,
871 output_format,
872 } => {
873 let example = NamedExample::load(path).unwrap();
874 example.write(output_format, io::stdout()).unwrap();
875 let _ = cx.update(|cx| cx.quit());
876 return;
877 }
878 };
879
880 match result {
881 Ok(output) => {
882 println!("{}", output);
883 let _ = cx.update(|cx| cx.quit());
884 }
885 Err(e) => {
886 eprintln!("Failed: {:?}", e);
887 exit(1);
888 }
889 }
890 })
891 .detach();
892 });
893}