1mod headless;
2mod retrieval_stats;
3mod source_location;
4mod util;
5
6use crate::retrieval_stats::retrieval_stats;
7use ::util::paths::PathStyle;
8use anyhow::{Result, anyhow};
9use clap::{Args, Parser, Subcommand};
10use cloud_llm_client::predict_edits_v3::{self};
11use edit_prediction_context::{
12 EditPredictionContextOptions, EditPredictionExcerptOptions, EditPredictionScoreOptions,
13};
14use gpui::{Application, AsyncApp, prelude::*};
15use language::Bias;
16use language_model::LlmApiToken;
17use project::Project;
18use release_channel::AppVersion;
19use reqwest_client::ReqwestClient;
20use serde_json::json;
21use std::{collections::HashSet, path::PathBuf, process::exit, str::FromStr, sync::Arc};
22use zeta::{PerformPredictEditsParams, Zeta};
23
24use crate::headless::ZetaCliAppState;
25use crate::source_location::SourceLocation;
26use crate::util::{open_buffer, open_buffer_with_language_server};
27
28#[derive(Parser, Debug)]
29#[command(name = "zeta")]
30struct ZetaCliArgs {
31 #[command(subcommand)]
32 command: Commands,
33}
34
35#[derive(Subcommand, Debug)]
36enum Commands {
37 Context(ContextArgs),
38 Zeta2Context {
39 #[clap(flatten)]
40 zeta2_args: Zeta2Args,
41 #[clap(flatten)]
42 context_args: ContextArgs,
43 },
44 Predict {
45 #[arg(long)]
46 predict_edits_body: Option<FileOrStdin>,
47 #[clap(flatten)]
48 context_args: Option<ContextArgs>,
49 },
50 RetrievalStats {
51 #[clap(flatten)]
52 zeta2_args: Zeta2Args,
53 #[arg(long)]
54 worktree: PathBuf,
55 #[arg(long)]
56 extension: Option<String>,
57 #[arg(long)]
58 limit: Option<usize>,
59 #[arg(long)]
60 skip: Option<usize>,
61 },
62}
63
64#[derive(Debug, Args)]
65#[group(requires = "worktree")]
66struct ContextArgs {
67 #[arg(long)]
68 worktree: PathBuf,
69 #[arg(long)]
70 cursor: SourceLocation,
71 #[arg(long)]
72 use_language_server: bool,
73 #[arg(long)]
74 events: Option<FileOrStdin>,
75}
76
77#[derive(Debug, Args)]
78struct Zeta2Args {
79 #[arg(long, default_value_t = 8192)]
80 max_prompt_bytes: usize,
81 #[arg(long, default_value_t = 2048)]
82 max_excerpt_bytes: usize,
83 #[arg(long, default_value_t = 1024)]
84 min_excerpt_bytes: usize,
85 #[arg(long, default_value_t = 0.66)]
86 target_before_cursor_over_total_bytes: f32,
87 #[arg(long, default_value_t = 1024)]
88 max_diagnostic_bytes: usize,
89 #[arg(long, value_enum, default_value_t = PromptFormat::default())]
90 prompt_format: PromptFormat,
91 #[arg(long, value_enum, default_value_t = Default::default())]
92 output_format: OutputFormat,
93 #[arg(long, default_value_t = 42)]
94 file_indexing_parallelism: usize,
95 #[arg(long, default_value_t = false)]
96 disable_imports_gathering: bool,
97}
98
99#[derive(clap::ValueEnum, Default, Debug, Clone)]
100enum PromptFormat {
101 #[default]
102 MarkedExcerpt,
103 LabeledSections,
104 OnlySnippets,
105}
106
107impl Into<predict_edits_v3::PromptFormat> for PromptFormat {
108 fn into(self) -> predict_edits_v3::PromptFormat {
109 match self {
110 Self::MarkedExcerpt => predict_edits_v3::PromptFormat::MarkedExcerpt,
111 Self::LabeledSections => predict_edits_v3::PromptFormat::LabeledSections,
112 Self::OnlySnippets => predict_edits_v3::PromptFormat::OnlySnippets,
113 }
114 }
115}
116
117#[derive(clap::ValueEnum, Default, Debug, Clone)]
118enum OutputFormat {
119 #[default]
120 Prompt,
121 Request,
122 Full,
123}
124
125#[derive(Debug, Clone)]
126enum FileOrStdin {
127 File(PathBuf),
128 Stdin,
129}
130
131impl FileOrStdin {
132 async fn read_to_string(&self) -> Result<String, std::io::Error> {
133 match self {
134 FileOrStdin::File(path) => smol::fs::read_to_string(path).await,
135 FileOrStdin::Stdin => smol::unblock(|| std::io::read_to_string(std::io::stdin())).await,
136 }
137 }
138}
139
140impl FromStr for FileOrStdin {
141 type Err = <PathBuf as FromStr>::Err;
142
143 fn from_str(s: &str) -> Result<Self, Self::Err> {
144 match s {
145 "-" => Ok(Self::Stdin),
146 _ => Ok(Self::File(PathBuf::from_str(s)?)),
147 }
148 }
149}
150
151enum GetContextOutput {
152 Zeta1(zeta::GatherContextOutput),
153 Zeta2(String),
154}
155
156async fn get_context(
157 zeta2_args: Option<Zeta2Args>,
158 args: ContextArgs,
159 app_state: &Arc<ZetaCliAppState>,
160 cx: &mut AsyncApp,
161) -> Result<GetContextOutput> {
162 let ContextArgs {
163 worktree: worktree_path,
164 cursor,
165 use_language_server,
166 events,
167 } = args;
168
169 let worktree_path = worktree_path.canonicalize()?;
170
171 let project = cx.update(|cx| {
172 Project::local(
173 app_state.client.clone(),
174 app_state.node_runtime.clone(),
175 app_state.user_store.clone(),
176 app_state.languages.clone(),
177 app_state.fs.clone(),
178 None,
179 cx,
180 )
181 })?;
182
183 let worktree = project
184 .update(cx, |project, cx| {
185 project.create_worktree(&worktree_path, true, cx)
186 })?
187 .await?;
188
189 let mut ready_languages = HashSet::default();
190 let (_lsp_open_handle, buffer) = if use_language_server {
191 let (lsp_open_handle, _, buffer) = open_buffer_with_language_server(
192 project.clone(),
193 worktree.clone(),
194 cursor.path.clone(),
195 &mut ready_languages,
196 cx,
197 )
198 .await?;
199 (Some(lsp_open_handle), buffer)
200 } else {
201 let buffer =
202 open_buffer(project.clone(), worktree.clone(), cursor.path.clone(), cx).await?;
203 (None, buffer)
204 };
205
206 let full_path_str = worktree
207 .read_with(cx, |worktree, _| worktree.root_name().join(&cursor.path))?
208 .display(PathStyle::local())
209 .to_string();
210
211 let snapshot = cx.update(|cx| buffer.read(cx).snapshot())?;
212 let clipped_cursor = snapshot.clip_point(cursor.point, Bias::Left);
213 if clipped_cursor != cursor.point {
214 let max_row = snapshot.max_point().row;
215 if cursor.point.row < max_row {
216 return Err(anyhow!(
217 "Cursor position {:?} is out of bounds (line length is {})",
218 cursor.point,
219 snapshot.line_len(cursor.point.row)
220 ));
221 } else {
222 return Err(anyhow!(
223 "Cursor position {:?} is out of bounds (max row is {})",
224 cursor.point,
225 max_row
226 ));
227 }
228 }
229
230 let events = match events {
231 Some(events) => events.read_to_string().await?,
232 None => String::new(),
233 };
234
235 if let Some(zeta2_args) = zeta2_args {
236 // wait for worktree scan before starting zeta2 so that wait_for_initial_indexing waits for
237 // the whole worktree.
238 worktree
239 .read_with(cx, |worktree, _cx| {
240 worktree.as_local().unwrap().scan_complete()
241 })?
242 .await;
243 let output = cx
244 .update(|cx| {
245 let zeta = cx.new(|cx| {
246 zeta2::Zeta::new(app_state.client.clone(), app_state.user_store.clone(), cx)
247 });
248 let indexing_done_task = zeta.update(cx, |zeta, cx| {
249 zeta.set_options(zeta2_args.to_options(true));
250 zeta.register_buffer(&buffer, &project, cx);
251 zeta.wait_for_initial_indexing(&project, cx)
252 });
253 cx.spawn(async move |cx| {
254 indexing_done_task.await?;
255 let request = zeta
256 .update(cx, |zeta, cx| {
257 let cursor = buffer.read(cx).snapshot().anchor_before(clipped_cursor);
258 zeta.cloud_request_for_zeta_cli(&project, &buffer, cursor, cx)
259 })?
260 .await?;
261
262 let planned_prompt = cloud_zeta2_prompt::PlannedPrompt::populate(&request)?;
263 let (prompt_string, section_labels) = planned_prompt.to_prompt_string()?;
264
265 match zeta2_args.output_format {
266 OutputFormat::Prompt => anyhow::Ok(prompt_string),
267 OutputFormat::Request => {
268 anyhow::Ok(serde_json::to_string_pretty(&request)?)
269 }
270 OutputFormat::Full => anyhow::Ok(serde_json::to_string_pretty(&json!({
271 "request": request,
272 "prompt": prompt_string,
273 "section_labels": section_labels,
274 }))?),
275 }
276 })
277 })?
278 .await?;
279 Ok(GetContextOutput::Zeta2(output))
280 } else {
281 let prompt_for_events = move || (events, 0);
282 Ok(GetContextOutput::Zeta1(
283 cx.update(|cx| {
284 zeta::gather_context(
285 full_path_str,
286 &snapshot,
287 clipped_cursor,
288 prompt_for_events,
289 cx,
290 )
291 })?
292 .await?,
293 ))
294 }
295}
296
297impl Zeta2Args {
298 fn to_options(&self, omit_excerpt_overlaps: bool) -> zeta2::ZetaOptions {
299 zeta2::ZetaOptions {
300 context: EditPredictionContextOptions {
301 use_imports: !self.disable_imports_gathering,
302 excerpt: EditPredictionExcerptOptions {
303 max_bytes: self.max_excerpt_bytes,
304 min_bytes: self.min_excerpt_bytes,
305 target_before_cursor_over_total_bytes: self
306 .target_before_cursor_over_total_bytes,
307 },
308 score: EditPredictionScoreOptions {
309 omit_excerpt_overlaps,
310 },
311 },
312 max_diagnostic_bytes: self.max_diagnostic_bytes,
313 max_prompt_bytes: self.max_prompt_bytes,
314 prompt_format: self.prompt_format.clone().into(),
315 file_indexing_parallelism: self.file_indexing_parallelism,
316 }
317 }
318}
319
320fn main() {
321 zlog::init();
322 zlog::init_output_stderr();
323 let args = ZetaCliArgs::parse();
324 let http_client = Arc::new(ReqwestClient::new());
325 let app = Application::headless().with_http_client(http_client);
326
327 app.run(move |cx| {
328 let app_state = Arc::new(headless::init(cx));
329 cx.spawn(async move |cx| {
330 let result = match args.command {
331 Commands::Zeta2Context {
332 zeta2_args,
333 context_args,
334 } => match get_context(Some(zeta2_args), context_args, &app_state, cx).await {
335 Ok(GetContextOutput::Zeta1 { .. }) => unreachable!(),
336 Ok(GetContextOutput::Zeta2(output)) => Ok(output),
337 Err(err) => Err(err),
338 },
339 Commands::Context(context_args) => {
340 match get_context(None, context_args, &app_state, cx).await {
341 Ok(GetContextOutput::Zeta1(output)) => {
342 Ok(serde_json::to_string_pretty(&output.body).unwrap())
343 }
344 Ok(GetContextOutput::Zeta2 { .. }) => unreachable!(),
345 Err(err) => Err(err),
346 }
347 }
348 Commands::Predict {
349 predict_edits_body,
350 context_args,
351 } => {
352 cx.spawn(async move |cx| {
353 let app_version = cx.update(|cx| AppVersion::global(cx))?;
354 app_state.client.sign_in(true, cx).await?;
355 let llm_token = LlmApiToken::default();
356 llm_token.refresh(&app_state.client).await?;
357
358 let predict_edits_body =
359 if let Some(predict_edits_body) = predict_edits_body {
360 serde_json::from_str(&predict_edits_body.read_to_string().await?)?
361 } else if let Some(context_args) = context_args {
362 match get_context(None, context_args, &app_state, cx).await? {
363 GetContextOutput::Zeta1(output) => output.body,
364 GetContextOutput::Zeta2 { .. } => unreachable!(),
365 }
366 } else {
367 return Err(anyhow!(
368 "Expected either --predict-edits-body-file \
369 or the required args of the `context` command."
370 ));
371 };
372
373 let (response, _usage) =
374 Zeta::perform_predict_edits(PerformPredictEditsParams {
375 client: app_state.client.clone(),
376 llm_token,
377 app_version,
378 body: predict_edits_body,
379 })
380 .await?;
381
382 Ok(response.output_excerpt)
383 })
384 .await
385 }
386 Commands::RetrievalStats {
387 zeta2_args,
388 worktree,
389 extension,
390 limit,
391 skip,
392 } => {
393 retrieval_stats(
394 worktree,
395 app_state,
396 extension,
397 limit,
398 skip,
399 (&zeta2_args).to_options(false),
400 cx,
401 )
402 .await
403 }
404 };
405 match result {
406 Ok(output) => {
407 println!("{}", output);
408 let _ = cx.update(|cx| cx.quit());
409 }
410 Err(e) => {
411 eprintln!("Failed: {:?}", e);
412 exit(1);
413 }
414 }
415 })
416 .detach();
417 });
418}