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/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;
}
_ => {}
}
}