oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

zeta2: Build edit prediction prompt and process model output in client (#41870)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>
Co-authored-by: Ben Kunkle <ben@zed.dev>
Co-authored-by: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
This commit is contained in:
Max Brunsfeld
2025-11-06 18:36:58 -05:00
committed by GitHub
co-authored by Agus Zubiaga Ben Kunkle Piotr Osiewicz
parent fb87972f44
commit 784fdcaee3
32 changed files with 2198 additions and 2392 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/zeta_cli/src/evaluate.rs
+4 -5
View File
@@ -7,13 +7,14 @@ use std::{
use anyhow::Result;
use clap::Args;
use cloud_llm_client::udiff::DiffLine;
use collections::HashSet;
use gpui::AsyncApp;
use zeta2::udiff::DiffLine;
use crate::{
example::{Example, NamedExample},
headless::ZetaCliAppState,
paths::CACHE_DIR,
predict::{PredictionDetails, zeta2_predict},
};
@@ -54,10 +55,8 @@ pub async fn run_evaluate_one(
app_state: Arc<ZetaCliAppState>,
cx: &mut AsyncApp,
) -> Result<EvaluationResult> {
let cache_dir = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap_or_default())
.join("../../target/zeta-prediction-cache");
let example = NamedExample::load(&example_path).unwrap();
let example_cache_path = cache_dir.join(&example_path.file_name().unwrap());
let example_cache_path = CACHE_DIR.join(&example_path.file_name().unwrap());
let predictions = if !re_run && example_cache_path.exists() {
let file_contents = fs::read_to_string(&example_cache_path)?;
@@ -74,7 +73,7 @@ pub async fn run_evaluate_one(
};
if !example_cache_path.exists() {
fs::create_dir_all(&cache_dir).unwrap();
fs::create_dir_all(&*CACHE_DIR).unwrap();
fs::write(
example_cache_path,
serde_json::to_string(&predictions).unwrap(),
crates/zeta_cli/src/example.rs
+64 -415
View File
@@ -1,28 +1,31 @@
use std::{
borrow::Cow,
cell::RefCell,
env,
fmt::{self, Display},
fs,
io::Write,
mem,
ops::Range,
path::{Path, PathBuf},
sync::Arc,
};
use anyhow::{Context as _, Result};
use anyhow::{Context as _, Result, anyhow};
use clap::ValueEnum;
use collections::{HashMap, HashSet};
use cloud_zeta2_prompt::CURSOR_MARKER;
use collections::HashMap;
use futures::{
AsyncWriteExt as _,
lock::{Mutex, OwnedMutexGuard},
};
use gpui::{AsyncApp, Entity, http_client::Url};
use language::Buffer;
use language::{Anchor, Buffer};
use project::{Project, ProjectPath};
use pulldown_cmark::CowStr;
use serde::{Deserialize, Serialize};
use util::{paths::PathStyle, rel_path::RelPath};
use zeta2::udiff::OpenedBuffers;
use crate::paths::{REPOS_DIR, WORKTREES_DIR};
const UNCOMMITTED_DIFF_HEADING: &str = "Uncommitted Diff";
const EDIT_HISTORY_HEADING: &str = "Edit History";
@@ -215,12 +218,10 @@ impl NamedExample {
let (repo_owner, repo_name) = self.repo_name()?;
let file_name = self.file_name();
let worktrees_dir = env::current_dir()?.join("target").join("zeta-worktrees");
let repos_dir = env::current_dir()?.join("target").join("zeta-repos");
fs::create_dir_all(&repos_dir)?;
fs::create_dir_all(&worktrees_dir)?;
fs::create_dir_all(&*REPOS_DIR)?;
fs::create_dir_all(&*WORKTREES_DIR)?;
let repo_dir = repos_dir.join(repo_owner.as_ref()).join(repo_name.as_ref());
let repo_dir = REPOS_DIR.join(repo_owner.as_ref()).join(repo_name.as_ref());
let repo_lock = lock_repo(&repo_dir).await;
if !repo_dir.is_dir() {
@@ -251,7 +252,7 @@ impl NamedExample {
};
// Create the worktree for this example if needed.
let worktree_path = worktrees_dir.join(&file_name);
let worktree_path = WORKTREES_DIR.join(&file_name);
if worktree_path.is_dir() {
run_git(&worktree_path, &["clean", "--force", "-d"]).await?;
run_git(&worktree_path, &["reset", "--hard", "HEAD"]).await?;
@@ -309,7 +310,6 @@ impl NamedExample {
.collect()
}
#[allow(unused)]
fn repo_name(&self) -> Result<(Cow<'_, str>, Cow<'_, str>)> {
// git@github.com:owner/repo.git
if self.example.repository_url.contains('@') {
@@ -344,13 +344,63 @@ impl NamedExample {
}
}
pub async fn cursor_position(
&self,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<(Entity<Buffer>, Anchor)> {
let worktree = project.read_with(cx, |project, cx| {
project.visible_worktrees(cx).next().unwrap()
})?;
let cursor_path = RelPath::new(&self.example.cursor_path, PathStyle::Posix)?.into_arc();
let cursor_buffer = project
.update(cx, |project, cx| {
project.open_buffer(
ProjectPath {
worktree_id: worktree.read(cx).id(),
path: cursor_path,
},
cx,
)
})?
.await?;
let cursor_offset_within_excerpt = self
.example
.cursor_position
.find(CURSOR_MARKER)
.ok_or_else(|| anyhow!("missing cursor marker"))?;
let mut cursor_excerpt = self.example.cursor_position.clone();
cursor_excerpt.replace_range(
cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
"",
);
let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
let text = buffer.text();
let mut matches = text.match_indices(&cursor_excerpt);
let Some((excerpt_offset, _)) = matches.next() else {
anyhow::bail!(
"Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
);
};
assert!(matches.next().is_none());
Ok(excerpt_offset)
})??;
let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
let cursor_anchor =
cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
Ok((cursor_buffer, cursor_anchor))
}
#[must_use]
pub async fn apply_edit_history(
&self,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<HashSet<Entity<Buffer>>> {
apply_diff(&self.example.edit_history, project, cx).await
) -> Result<OpenedBuffers<'_>> {
zeta2::udiff::apply_diff(&self.example.edit_history, project, cx).await
}
}
@@ -446,404 +496,3 @@ pub async fn lock_repo(path: impl AsRef<Path>) -> OwnedMutexGuard<()> {
.lock_owned()
.await
}
#[must_use]
pub async fn apply_diff(
diff: &str,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<HashSet<Entity<Buffer>>> {
use cloud_llm_client::udiff::DiffLine;
use std::fmt::Write;
#[derive(Debug, Default)]
struct HunkState {
context: String,
edits: Vec<Edit>,
}
#[derive(Debug)]
struct Edit {
range: Range<usize>,
text: String,
}
let mut old_path = None;
let mut new_path = None;
let mut hunk = HunkState::default();
let mut diff_lines = diff.lines().map(DiffLine::parse).peekable();
let mut open_buffers = HashSet::default();
while let Some(diff_line) = diff_lines.next() {
match diff_line {
DiffLine::OldPath { path } => old_path = Some(path),
DiffLine::NewPath { path } => {
if old_path.is_none() {
anyhow::bail!(
"Found a new path header (`+++`) before an (`---`) old path header"
);
}
new_path = Some(path)
}
DiffLine::Context(ctx) => {
writeln!(&mut hunk.context, "{ctx}")?;
}
DiffLine::Deletion(del) => {
let range = hunk.context.len()..hunk.context.len() + del.len() + '\n'.len_utf8();
if let Some(last_edit) = hunk.edits.last_mut()
&& last_edit.range.end == range.start
{
last_edit.range.end = range.end;
} else {
hunk.edits.push(Edit {
range,
text: String::new(),
});
}
writeln!(&mut hunk.context, "{del}")?;
}
DiffLine::Addition(add) => {
let range = hunk.context.len()..hunk.context.len();
if let Some(last_edit) = hunk.edits.last_mut()
&& last_edit.range.end == range.start
{
writeln!(&mut last_edit.text, "{add}").unwrap();
} else {
hunk.edits.push(Edit {
range,
text: format!("{add}\n"),
});
}
}
DiffLine::HunkHeader(_) | DiffLine::Garbage(_) => {}
}
let at_hunk_end = match diff_lines.peek() {
Some(DiffLine::OldPath { .. }) | Some(DiffLine::HunkHeader(_)) | None => true,
_ => false,
};
if at_hunk_end {
let hunk = mem::take(&mut hunk);
let Some(old_path) = old_path.as_deref() else {
anyhow::bail!("Missing old path (`---`) header")
};
let Some(new_path) = new_path.as_deref() else {
anyhow::bail!("Missing new path (`+++`) header")
};
let buffer = project
.update(cx, |project, cx| {
let project_path = project
.find_project_path(old_path, cx)
.context("Failed to find old_path in project")?;
anyhow::Ok(project.open_buffer(project_path, cx))
})??
.await?;
open_buffers.insert(buffer.clone());
if old_path != new_path {
project
.update(cx, |project, cx| {
let project_file = project::File::from_dyn(buffer.read(cx).file()).unwrap();
let new_path = ProjectPath {
worktree_id: project_file.worktree_id(cx),
path: project_file.path.clone(),
};
project.rename_entry(project_file.entry_id.unwrap(), new_path, cx)
})?
.await?;
}
// TODO is it worth using project search?
buffer.update(cx, |buffer, cx| {
let context_offset = if hunk.context.is_empty() {
0
} else {
let text = buffer.text();
if let Some(offset) = text.find(&hunk.context) {
if text[offset + 1..].contains(&hunk.context) {
anyhow::bail!("Context is not unique enough:\n{}", hunk.context);
}
offset
} else {
anyhow::bail!(
"Failed to match context:\n{}\n\nBuffer:\n{}",
hunk.context,
text
);
}
};
buffer.edit(
hunk.edits.into_iter().map(|edit| {
(
context_offset + edit.range.start..context_offset + edit.range.end,
edit.text,
)
}),
None,
cx,
);
anyhow::Ok(())
})??;
}
}
anyhow::Ok(open_buffers)
}
#[cfg(test)]
mod tests {
use super::*;
use ::fs::FakeFs;
use gpui::TestAppContext;
use indoc::indoc;
use pretty_assertions::assert_eq;
use project::Project;
use serde_json::json;
use settings::SettingsStore;
use util::path;
#[gpui::test]
async fn test_apply_diff_successful(cx: &mut TestAppContext) {
let buffer_1_text = indoc! {r#"
one
two
three
four
five
"# };
let buffer_1_text_final = indoc! {r#"
3
4
5
"# };
let buffer_2_text = indoc! {r#"
six
seven
eight
nine
ten
"# };
let buffer_2_text_final = indoc! {r#"
5
six
seven
7.5
eight
nine
ten
11
"# };
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
Project::init_settings(cx);
language::init(cx);
});
let fs = FakeFs::new(cx.background_executor.clone());
fs.insert_tree(
path!("/root"),
json!({
"file1": buffer_1_text,
"file2": buffer_2_text,
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/root/file1
+++ b/root/file1
one
two
-three
+3
four
five
--- a/root/file1
+++ b/root/file1
3
-four
-five
+4
+5
--- a/root/file1
+++ b/root/file1
-one
-two
3
4
--- a/root/file2
+++ b/root/file2
+5
six
--- a/root/file2
+++ b/root/file2
seven
+7.5
eight
--- a/root/file2
+++ b/root/file2
ten
+11
"#};
let _buffers = apply_diff(diff, &project, &mut cx.to_async())
.await
.unwrap();
let buffer_1 = project
.update(cx, |project, cx| {
let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
project.open_buffer(project_path, cx)
})
.await
.unwrap();
buffer_1.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), buffer_1_text_final);
});
let buffer_2 = project
.update(cx, |project, cx| {
let project_path = project.find_project_path(path!("/root/file2"), cx).unwrap();
project.open_buffer(project_path, cx)
})
.await
.unwrap();
buffer_2.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), buffer_2_text_final);
});
}
#[gpui::test]
async fn test_apply_diff_non_unique(cx: &mut TestAppContext) {
let buffer_1_text = indoc! {r#"
one
two
three
four
five
one
two
three
four
five
"# };
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
Project::init_settings(cx);
language::init(cx);
});
let fs = FakeFs::new(cx.background_executor.clone());
fs.insert_tree(
path!("/root"),
json!({
"file1": buffer_1_text,
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/root/file1
+++ b/root/file1
one
two
-three
+3
four
five
"#};
apply_diff(diff, &project, &mut cx.to_async())
.await
.expect_err("Non-unique edits should fail");
}
#[gpui::test]
async fn test_apply_diff_unique_via_previous_context(cx: &mut TestAppContext) {
let start = indoc! {r#"
one
two
three
four
five
four
five
"# };
let end = indoc! {r#"
one
two
3
four
5
four
five
"# };
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
Project::init_settings(cx);
language::init(cx);
});
let fs = FakeFs::new(cx.background_executor.clone());
fs.insert_tree(
path!("/root"),
json!({
"file1": start,
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/root/file1
+++ b/root/file1
one
two
-three
+3
four
-five
+5
"#};
let _buffers = apply_diff(diff, &project, &mut cx.to_async())
.await
.unwrap();
let buffer_1 = project
.update(cx, |project, cx| {
let project_path = project.find_project_path(path!("/root/file1"), cx).unwrap();
project.open_buffer(project_path, cx)
})
.await
.unwrap();
buffer_1.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), end);
});
}
}
crates/zeta_cli/src/main.rs
+6 -223
View File
@@ -1,6 +1,7 @@
mod evaluate;
mod example;
mod headless;
mod paths;
mod predict;
mod source_location;
mod syntax_retrieval_stats;
@@ -10,28 +11,22 @@ use crate::evaluate::{EvaluateArguments, run_evaluate};
use crate::example::{ExampleFormat, NamedExample};
use crate::predict::{PredictArguments, run_zeta2_predict};
use crate::syntax_retrieval_stats::retrieval_stats;
use ::serde::Serialize;
use ::util::paths::PathStyle;
use anyhow::{Context as _, Result, anyhow};
use anyhow::{Result, anyhow};
use clap::{Args, Parser, Subcommand};
use cloud_llm_client::predict_edits_v3::{self, Excerpt};
use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
use cloud_llm_client::predict_edits_v3;
use edit_prediction_context::{
EditPredictionContextOptions, EditPredictionExcerpt, EditPredictionExcerptOptions,
EditPredictionScoreOptions, Line,
EditPredictionContextOptions, EditPredictionExcerptOptions, EditPredictionScoreOptions,
};
use futures::StreamExt as _;
use futures::channel::mpsc;
use gpui::{Application, AsyncApp, Entity, prelude::*};
use language::{Bias, Buffer, BufferSnapshot, OffsetRangeExt, Point};
use language_model::LanguageModelRegistry;
use language::{Bias, Buffer, BufferSnapshot, Point};
use project::{Project, Worktree};
use reqwest_client::ReqwestClient;
use serde_json::json;
use std::io::{self};
use std::time::Duration;
use std::{collections::HashSet, path::PathBuf, str::FromStr, sync::Arc};
use zeta2::{ContextMode, LlmContextOptions, SearchToolQuery};
use zeta2::ContextMode;
use crate::headless::ZetaCliAppState;
use crate::source_location::SourceLocation;
@@ -79,12 +74,6 @@ enum Zeta2Command {
#[command(subcommand)]
command: Zeta2SyntaxCommand,
},
Llm {
#[clap(flatten)]
args: Zeta2Args,
#[command(subcommand)]
command: Zeta2LlmCommand,
},
Predict(PredictArguments),
Eval(EvaluateArguments),
}
@@ -107,14 +96,6 @@ enum Zeta2SyntaxCommand {
},
}
#[derive(Subcommand, Debug)]
enum Zeta2LlmCommand {
Context {
#[clap(flatten)]
context_args: ContextArgs,
},
}
#[derive(Debug, Args)]
#[group(requires = "worktree")]
struct ContextArgs {
@@ -388,197 +369,6 @@ async fn zeta2_syntax_context(
Ok(output)
}
async fn zeta2_llm_context(
zeta2_args: Zeta2Args,
context_args: ContextArgs,
app_state: &Arc<ZetaCliAppState>,
cx: &mut AsyncApp,
) -> Result<String> {
let LoadedContext {
buffer,
clipped_cursor,
snapshot: cursor_snapshot,
project,
..
} = load_context(&context_args, app_state, cx).await?;
let cursor_position = cursor_snapshot.anchor_after(clipped_cursor);
cx.update(|cx| {
LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
registry
.provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
.unwrap()
.authenticate(cx)
})
})?
.await?;
let edit_history_unified_diff = match context_args.edit_history {
Some(events) => events.read_to_string().await?,
None => String::new(),
};
let (debug_tx, mut debug_rx) = mpsc::unbounded();
let excerpt_options = EditPredictionExcerptOptions {
max_bytes: zeta2_args.max_excerpt_bytes,
min_bytes: zeta2_args.min_excerpt_bytes,
target_before_cursor_over_total_bytes: zeta2_args.target_before_cursor_over_total_bytes,
};
let related_excerpts = cx
.update(|cx| {
zeta2::related_excerpts::find_related_excerpts(
buffer,
cursor_position,
&project,
edit_history_unified_diff,
&LlmContextOptions {
excerpt: excerpt_options.clone(),
},
Some(debug_tx),
cx,
)
})?
.await?;
let cursor_excerpt = EditPredictionExcerpt::select_from_buffer(
clipped_cursor,
&cursor_snapshot,
&excerpt_options,
None,
)
.context("line didn't fit")?;
#[derive(Serialize)]
struct Output {
excerpts: Vec<OutputExcerpt>,
formatted_excerpts: String,
meta: OutputMeta,
}
#[derive(Default, Serialize)]
struct OutputMeta {
search_prompt: String,
search_queries: Vec<SearchToolQuery>,
}
#[derive(Serialize)]
struct OutputExcerpt {
path: PathBuf,
#[serde(flatten)]
excerpt: Excerpt,
}
let mut meta = OutputMeta::default();
while let Some(debug_info) = debug_rx.next().await {
match debug_info {
zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
meta.search_prompt = info.search_prompt;
}
zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
meta.search_queries = info.queries
}
_ => {}
}
}
cx.update(|cx| {
let mut excerpts = Vec::new();
let mut formatted_excerpts = String::new();
let cursor_insertions = [(
predict_edits_v3::Point {
line: Line(clipped_cursor.row),
column: clipped_cursor.column,
},
CURSOR_MARKER,
)];
let mut cursor_excerpt_added = false;
for (buffer, ranges) in related_excerpts {
let excerpt_snapshot = buffer.read(cx).snapshot();
let mut line_ranges = ranges
.into_iter()
.map(|range| {
let point_range = range.to_point(&excerpt_snapshot);
Line(point_range.start.row)..Line(point_range.end.row)
})
.collect::<Vec<_>>();
let Some(file) = excerpt_snapshot.file() else {
continue;
};
let path = file.full_path(cx);
let is_cursor_file = path == cursor_snapshot.file().unwrap().full_path(cx);
if is_cursor_file {
let insertion_ix = line_ranges
.binary_search_by(|probe| {
probe
.start
.cmp(&cursor_excerpt.line_range.start)
.then(cursor_excerpt.line_range.end.cmp(&probe.end))
})
.unwrap_or_else(|ix| ix);
line_ranges.insert(insertion_ix, cursor_excerpt.line_range.clone());
cursor_excerpt_added = true;
}
let merged_excerpts =
zeta2::merge_excerpts::merge_excerpts(&excerpt_snapshot, line_ranges)
.into_iter()
.map(|excerpt| OutputExcerpt {
path: path.clone(),
excerpt,
});
let excerpt_start_ix = excerpts.len();
excerpts.extend(merged_excerpts);
write_codeblock(
&path,
excerpts[excerpt_start_ix..].iter().map(|e| &e.excerpt),
if is_cursor_file {
&cursor_insertions
} else {
&[]
},
Line(excerpt_snapshot.max_point().row),
true,
&mut formatted_excerpts,
);
}
if !cursor_excerpt_added {
write_codeblock(
&cursor_snapshot.file().unwrap().full_path(cx),
&[Excerpt {
start_line: cursor_excerpt.line_range.start,
text: cursor_excerpt.text(&cursor_snapshot).body.into(),
}],
&cursor_insertions,
Line(cursor_snapshot.max_point().row),
true,
&mut formatted_excerpts,
);
}
let output = Output {
excerpts,
formatted_excerpts,
meta,
};
Ok(serde_json::to_string_pretty(&output)?)
})
.unwrap()
}
async fn zeta1_context(
args: ContextArgs,
app_state: &Arc<ZetaCliAppState>,
@@ -670,13 +460,6 @@ fn main() {
};
println!("{}", result.unwrap());
}
Zeta2Command::Llm { args, command } => match command {
Zeta2LlmCommand::Context { context_args } => {
let result =
zeta2_llm_context(args, context_args, &app_state, cx).await;
println!("{}", result.unwrap());
}
},
},
Command::ConvertExample {
path,
crates/zeta_cli/src/paths.rs
+8
View File
@@ -0,0 +1,8 @@
use std::{env, path::PathBuf, sync::LazyLock};
static TARGET_DIR: LazyLock<PathBuf> = LazyLock::new(|| env::current_dir().unwrap().join("target"));
pub static CACHE_DIR: LazyLock<PathBuf> =
LazyLock::new(|| TARGET_DIR.join("zeta-prediction-cache"));
pub static REPOS_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_DIR.join("zeta-repos"));
pub static WORKTREES_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_DIR.join("zeta-worktrees"));
pub static LOGS_DIR: LazyLock<PathBuf> = LazyLock::new(|| TARGET_DIR.join("zeta-logs"));
crates/zeta_cli/src/predict.rs
+41 -64
View File
@@ -1,22 +1,20 @@
use crate::example::{ActualExcerpt, NamedExample};
use crate::headless::ZetaCliAppState;
use crate::paths::LOGS_DIR;
use ::serde::Serialize;
use ::util::paths::PathStyle;
use anyhow::{Context as _, Result, anyhow};
use clap::Args;
use cloud_zeta2_prompt::{CURSOR_MARKER, write_codeblock};
use futures::StreamExt as _;
use gpui::AsyncApp;
use language_model::LanguageModelRegistry;
use project::{Project, ProjectPath};
use project::Project;
use serde::Deserialize;
use std::cell::Cell;
use std::fs;
use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use util::rel_path::RelPath;
#[derive(Debug, Args)]
pub struct PredictArguments {
@@ -50,21 +48,12 @@ pub async fn zeta2_predict(
app_state: &Arc<ZetaCliAppState>,
cx: &mut AsyncApp,
) -> Result<PredictionDetails> {
fs::create_dir_all(&*LOGS_DIR)?;
let worktree_path = example.setup_worktree().await?;
if !AUTHENTICATED.get() {
AUTHENTICATED.set(true);
cx.update(|cx| {
LanguageModelRegistry::global(cx).update(cx, |registry, cx| {
registry
.provider(&zeta2::related_excerpts::MODEL_PROVIDER_ID)
.unwrap()
.authenticate(cx)
})
})?
.await?;
app_state
.client
.sign_in_with_optional_connect(true, cx)
@@ -83,6 +72,8 @@ pub async fn zeta2_predict(
)
})?;
let buffer_store = project.read_with(cx, |project, _| project.buffer_store().clone())?;
let worktree = project
.update(cx, |project, cx| {
project.create_worktree(&worktree_path, true, cx)
@@ -94,58 +85,30 @@ pub async fn zeta2_predict(
})?
.await;
let _edited_buffers = example.apply_edit_history(&project, cx).await?;
let cursor_path = RelPath::new(&example.example.cursor_path, PathStyle::Posix)?.into_arc();
let cursor_buffer = project
.update(cx, |project, cx| {
project.open_buffer(
ProjectPath {
worktree_id: worktree.read(cx).id(),
path: cursor_path,
},
cx,
)
})?
.await?;
let cursor_offset_within_excerpt = example
.example
.cursor_position
.find(CURSOR_MARKER)
.ok_or_else(|| anyhow!("missing cursor marker"))?;
let mut cursor_excerpt = example.example.cursor_position.clone();
cursor_excerpt.replace_range(
cursor_offset_within_excerpt..(cursor_offset_within_excerpt + CURSOR_MARKER.len()),
"",
);
let excerpt_offset = cursor_buffer.read_with(cx, |buffer, _cx| {
let text = buffer.text();
let mut matches = text.match_indices(&cursor_excerpt);
let Some((excerpt_offset, _)) = matches.next() else {
anyhow::bail!(
"Cursor excerpt did not exist in buffer.\nExcerpt:\n\n{cursor_excerpt}\nBuffer text:\n{text}\n"
);
};
assert!(matches.next().is_none());
Ok(excerpt_offset)
})??;
let cursor_offset = excerpt_offset + cursor_offset_within_excerpt;
let cursor_anchor =
cursor_buffer.read_with(cx, |buffer, _| buffer.anchor_after(cursor_offset))?;
let zeta = cx.update(|cx| zeta2::Zeta::global(&app_state.client, &app_state.user_store, cx))?;
cx.subscribe(&buffer_store, {
let project = project.clone();
move |_, event, cx| match event {
project::buffer_store::BufferStoreEvent::BufferAdded(buffer) => {
zeta2::Zeta::try_global(cx)
.unwrap()
.update(cx, |zeta, cx| zeta.register_buffer(&buffer, &project, cx));
}
_ => {}
}
})?
.detach();
let _edited_buffers = example.apply_edit_history(&project, cx).await?;
let (cursor_buffer, cursor_anchor) = example.cursor_position(&project, cx).await?;
let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info())?;
let refresh_task = zeta.update(cx, |zeta, cx| {
zeta.register_buffer(&cursor_buffer, &project, cx);
zeta.refresh_context(project.clone(), cursor_buffer.clone(), cursor_anchor, cx)
})?;
let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info())?;
let mut context_retrieval_started_at = None;
let mut context_retrieval_finished_at = None;
let mut search_queries_generated_at = None;
@@ -159,9 +122,14 @@ pub async fn zeta2_predict(
match event {
zeta2::ZetaDebugInfo::ContextRetrievalStarted(info) => {
context_retrieval_started_at = Some(info.timestamp);
fs::write(LOGS_DIR.join("search_prompt.md"), &info.search_prompt)?;
}
zeta2::ZetaDebugInfo::SearchQueriesGenerated(info) => {
search_queries_generated_at = Some(info.timestamp);
fs::write(
LOGS_DIR.join("search_queries.json"),
serde_json::to_string_pretty(&info.regex_by_glob).unwrap(),
)?;
}
zeta2::ZetaDebugInfo::SearchQueriesExecuted(info) => {
search_queries_executed_at = Some(info.timestamp);
@@ -173,11 +141,21 @@ pub async fn zeta2_predict(
zeta.request_prediction(&project, &cursor_buffer, cursor_anchor, cx)
})?);
}
zeta2::ZetaDebugInfo::EditPredicted(request) => {
zeta2::ZetaDebugInfo::EditPredictionRequested(request) => {
prediction_started_at = Some(Instant::now());
request.response_rx.await?.map_err(|err| anyhow!(err))?;
fs::write(
LOGS_DIR.join("prediction_prompt.md"),
&request.local_prompt.unwrap_or_default(),
)?;
let response = request.response_rx.await?.0.map_err(|err| anyhow!(err))?;
prediction_finished_at = Some(Instant::now());
fs::write(
LOGS_DIR.join("prediction_response.json"),
&serde_json::to_string_pretty(&response).unwrap(),
)?;
for included_file in request.request.included_files {
let insertions = vec![(request.request.cursor_point, CURSOR_MARKER)];
result
@@ -201,7 +179,6 @@ pub async fn zeta2_predict(
}
break;
}
_ => {}
}
}