yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ep: Send V4 prediction requests (#59930)

Browse Source 
# Objective

- Describe the objective or issue this PR addresses.
- If you're fixing a specific issue, use "Fixes #X" for each issue as
[described in the GitHub
docs](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword).

## Solution

- Describe the solution used to achieve the objective above.

## Testing

- Did you test these changes? If so, how?
- Are there any parts that need more testing?
- How can other people (reviewers) test your changes? Is there anything
specific they need to know?
- If relevant, what platforms did you test these changes on, and are
there any important ones you can't test?

## Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content adheres to Zed's UI standards
([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
and
[icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md)
guidelines)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

## Showcase

> This section is optional. If this PR does not include a visual change
or does not add a new user-facing feature, you can delete this section.

- Help others understand the result of this PR by showcasing your
awesome work!
- If this PR includes a visual change, consider adding a screenshot,
GIF, or video
- A before/after comparison is very useful for changes to existing
features!

While a showcase should aim to be brief and digestible, you can use a
toggleable section to save space on longer showcases:

<details>
  <summary>Click to view showcase</summary>

My super cool demos here

</details>

---

Release Notes:

- N/A or Added/Fixed/Improved ...
This commit is contained in:
Ben Kunkle authored and GitHub committed 2026-06-26 15:47:51 +00:00
1 parent c2ba335b29
commit aa0cfcd190
5 files changed
+737 -125
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/edit_prediction/src/edit_prediction.rs
+95 -23
View File
@@ -13,6 +13,7 @@ use cloud_llm_client::predict_edits_v3::{
PREDICT_EDITS_TRIGGER_HEADER_NAME, PredictEditsMode, PredictEditsV3Request,
PredictEditsV3Response, RawCompletionRequest, RawCompletionResponse,
};
use cloud_llm_client::predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response};
use cloud_llm_client::{
EditPredictionRejectReason, EditPredictionRejection,
MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
@@ -59,7 +60,7 @@ use std::rc::Rc;
use text::{AnchorRangeExt, Edit};
use workspace::{AppState, Workspace};
use zeta_prompt::ContextSource;
use zeta_prompt::{Zeta2PromptInput, ZetaFormat};
use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput, ZetaFormat};
use std::mem;
use std::ops::Range;
@@ -190,7 +191,6 @@ pub enum EditPredictionModel {
Mercury,
}
#[derive(Clone)]
pub struct EditPredictionModelInput {
project: Entity<Project>,
buffer: Entity<Buffer>,
@@ -198,12 +198,14 @@ pub struct EditPredictionModelInput {
position: Anchor,
events: Vec<Arc<zeta_prompt::Event>>,
related_files: Vec<RelatedFile>,
editable_context: Option<Task<anyhow::Result<Vec<RelatedFile>>>>,
mode: PredictEditsMode,
trigger: PredictEditsRequestTrigger,
diagnostic_search_range: Range<Point>,
debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,
can_collect_data: bool,
is_open_source: bool,
allow_jump: bool,
}
#[derive(Debug)]
@@ -1107,6 +1109,7 @@ impl EditPredictionStore {
let client = self.client.clone();
let llm_token = self.llm_token.clone();
let app_version = AppVersion::global(cx);
let is_jumps_api = cx.has_flag::<EditPredictionJumpsFeatureFlag>();
let organization_id = self
.user_store
.read(cx)
@@ -1118,9 +1121,10 @@ impl EditPredictionStore {
.background_spawn(async move {
let organization_id =
organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
let url = client
.http_client()
.build_zed_llm_url("/edit_prediction_experiments", &[])?;
let url = client.http_client().build_zed_llm_url(
"/edit_prediction_experiments",
&[("is_jumps_api", if is_jumps_api { "true" } else { "false" })],
)?;
let mut response = client
.authenticated_llm_request(&llm_token, organization_id, |token| {
Ok(http_client::Request::builder()
@@ -2740,23 +2744,27 @@ impl EditPredictionStore {
}
self.get_or_init_project(&project, cx);
let project_state = self.projects.get(&project.entity_id()).unwrap();
let stored_events = project_state.events(cx);
let prompt_history_boundary = Some(PromptHistoryBoundary {
first_event_seq: project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
let (stored_events, prompt_history_boundary, debug_tx) = {
let project_state = self.projects.get(&project.entity_id()).unwrap();
(
project_state.events(cx),
Some(PromptHistoryBoundary {
first_event_seq: project_state
.last_event
.as_ref()
.map_or(project_state.next_last_event_seq, |last_event| {
last_event.seq
}),
snapshot: project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone()),
}),
snapshot: project_state
.last_event
.as_ref()
.map(|last_event| last_event.new_snapshot.clone()),
});
project_state.debug_tx.clone(),
)
};
let events: Vec<Arc<zeta_prompt::Event>> =
stored_events.iter().map(|e| e.event.clone()).collect();
let debug_tx = project_state.debug_tx.clone();
let snapshot = active_buffer.read(cx).snapshot();
let cursor_point = position.to_point(&snapshot);
@@ -2766,6 +2774,7 @@ impl EditPredictionStore {
Point::new(diagnostic_search_start, 0)..Point::new(diagnostic_search_end, 0);
let related_files = self.context_for_project(&project, cx);
let allow_jump = is_cloud_zeta && cx.has_flag::<EditPredictionJumpsFeatureFlag>();
let mode = match all_language_settings(snapshot.file(), cx).edit_predictions_mode() {
EditPredictionsMode::Eager => PredictEditsMode::Eager,
EditPredictionsMode::Subtle => PredictEditsMode::Subtle,
@@ -2807,6 +2816,16 @@ impl EditPredictionStore {
&& is_open_source
&& self.is_data_collection_enabled(cx)
&& matches!(self.edit_prediction_model, EditPredictionModel::Zeta);
let editable_context = allow_jump.then(|| {
self.collect_editable_context(
project.clone(),
active_buffer.clone(),
position,
Vec::new(),
vec![ContextSource::CurrentFile, ContextSource::EditHistory],
cx,
)
});
let inputs = EditPredictionModelInput {
project: project.clone(),
buffer: active_buffer,
@@ -2814,12 +2833,14 @@ impl EditPredictionStore {
position,
events,
related_files,
editable_context,
mode,
trigger,
diagnostic_search_range,
debug_tx,
can_collect_data,
is_open_source,
allow_jump,
};
let task = match self.edit_prediction_model {
@@ -2906,11 +2927,62 @@ impl EditPredictionStore {
trigger: PredictEditsRequestTrigger,
mode: PredictEditsMode,
) -> Result<(PredictEditsV3Response, Option<EditPredictionUsage>)> {
let url = client
.http_client()
.build_zed_llm_url("/predict_edits/v3", &[])?;
let request = PredictEditsV3Request { input };
Self::send_predict_edits_request(
"/predict_edits/v3",
request,
preferred_experiment,
client,
llm_token,
organization_id,
app_version,
trigger,
mode,
)
.await
}
pub(crate) async fn send_v4_request(
input: Zeta3PromptInput,
preferred_experiment: Option<String>,
client: Arc<Client>,
llm_token: LlmApiToken,
organization_id: Option<OrganizationId>,
app_version: Version,
trigger: PredictEditsRequestTrigger,
mode: PredictEditsMode,
) -> Result<(PredictEditsV4Response, Option<EditPredictionUsage>)> {
let request = PredictEditsV4Request { input };
Self::send_predict_edits_request(
"/predict_edits/v4",
request,
preferred_experiment,
client,
llm_token,
organization_id,
app_version,
trigger,
mode,
)
.await
}
async fn send_predict_edits_request<Req, Res>(
path: &str,
request: Req,
preferred_experiment: Option<String>,
client: Arc<Client>,
llm_token: LlmApiToken,
organization_id: Option<OrganizationId>,
app_version: Version,
trigger: PredictEditsRequestTrigger,
mode: PredictEditsMode,
) -> Result<(Res, Option<EditPredictionUsage>)>
where
Req: serde::Serialize,
Res: serde::de::DeserializeOwned,
{
let url = client.http_client().build_zed_llm_url(path, &[])?;
let request_id = uuid::Uuid::new_v4().to_string();
let json_bytes = serde_json::to_vec(&request)?;
crates/edit_prediction/src/edit_prediction_tests.rs
+159 -13
View File
@@ -10,22 +10,24 @@ use cloud_llm_client::{
EditPredictionRejectReason, EditPredictionRejection, PredictEditsRequestTrigger,
RejectEditPredictionsBody,
predict_edits_v3::{PredictEditsV3Request, PredictEditsV3Response},
predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response},
};
use db::AppDatabase;
use edit_prediction_types::EditPredictionRequestTrigger;
use feature_flags::{FeatureFlag as _, FeatureFlagAppExt as _, FeatureFlagsSettings};
use futures::{
AsyncReadExt, FutureExt, StreamExt,
channel::{mpsc, oneshot},
};
use gpui::App;
use gpui::{
Entity, TestAppContext,
Entity, TestAppContext, UpdateGlobal,
http_client::{FakeHttpClient, Response},
};
use indoc::indoc;
use language::{
Anchor, Buffer, Capability, Diagnostic, DiagnosticEntry, DiagnosticSet, DiagnosticSeverity,
Point,
Point, unified_diff_with_offsets,
};
use lsp::LanguageServerId;
use parking_lot::Mutex;
@@ -43,6 +45,7 @@ use uuid::Uuid;
use workspace::{AppState, CollaboratorId, MultiWorkspace};
use zeta_prompt::Zeta2PromptInput;
use crate::prediction::EditPredictionInputs;
use crate::udiff::apply_diff_to_string;
use crate::{
BufferEditPrediction, EDIT_PREDICTION_SETTLED_QUIESCENCE, EditPredictionId,
@@ -54,6 +57,7 @@ use super::*;
#[gpui::test]
async fn test_current_state(cx: &mut TestAppContext) {
let (ep_store, mut requests) = init_test_with_fake_client(cx);
cx.update(|cx| set_jumps_feature_flag_override(cx, "on"));
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/root",
@@ -92,12 +96,12 @@ async fn test_current_state(cx: &mut TestAppContext) {
cx,
)
});
let (request, respond_tx) = requests.predict.next().await.unwrap();
let (_request, respond_tx) = requests.predict_v4.next().await.unwrap();
respond_tx
.send(model_response(
&request,
indoc! {r"
.send(PredictEditsV4Response {
request_id: Uuid::new_v4().to_string(),
patch: indoc! {r"
--- a/root/1.txt
+++ b/root/1.txt
@@ ... @@
@@ -105,8 +109,10 @@ async fn test_current_state(cx: &mut TestAppContext) {
-How
+How are you?
Bye
"},
))
"}
.to_string(),
model_version: None,
})
.unwrap();
cx.run_until_parked();
@@ -2501,11 +2507,8 @@ fn prompt_from_request(request: &PredictEditsV3Request) -> String {
.expect("default zeta prompt formatting should succeed in edit prediction tests")
}
fn assert_no_predict_request_ready(
requests: &mut mpsc::UnboundedReceiver<(
PredictEditsV3Request,
oneshot::Sender<PredictEditsV3Response>,
)>,
fn assert_no_predict_request_ready<Request, Response>(
requests: &mut mpsc::UnboundedReceiver<(Request, oneshot::Sender<Response>)>,
) {
if requests.next().now_or_never().flatten().is_some() {
panic!("Unexpected prediction request while throttled.");
@@ -2517,6 +2520,10 @@ struct RequestChannels {
PredictEditsV3Request,
oneshot::Sender<PredictEditsV3Response>,
)>,
predict_v4: mpsc::UnboundedReceiver<(
PredictEditsV4Request,
oneshot::Sender<PredictEditsV4Response>,
)>,
reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>,
settled: mpsc::UnboundedReceiver<SettledEditPrediction>,
}
@@ -2535,6 +2542,7 @@ fn init_test_with_fake_client_and_legacy_data_collection(
cx.set_global(AppDatabase::test_new());
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
disable_jumps_feature_flag(cx);
zlog::init_test();
if let Some(legacy_data_collection_choice) = legacy_data_collection_choice {
@@ -2549,6 +2557,7 @@ fn init_test_with_fake_client_and_legacy_data_collection(
}
let (predict_req_tx, predict_req_rx) = mpsc::unbounded();
let (predict_v4_req_tx, predict_v4_req_rx) = mpsc::unbounded();
let (reject_req_tx, reject_req_rx) = mpsc::unbounded();
let (settled_req_tx, settled_req_rx) = mpsc::unbounded();
@@ -2562,6 +2571,7 @@ fn init_test_with_fake_client_and_legacy_data_collection(
.map(str::to_owned);
let mut body = req.into_body();
let predict_req_tx = predict_req_tx.clone();
let predict_v4_req_tx = predict_v4_req_tx.clone();
let reject_req_tx = reject_req_tx.clone();
let settled_req_tx = settled_req_tx.clone();
async move {
@@ -2591,6 +2601,27 @@ fn init_test_with_fake_client_and_legacy_data_collection(
}
serde_json::to_string(&response).unwrap()
}
"/predict_edits/v4" => {
let mut buf = Vec::new();
body.read_to_end(&mut buf).await.ok();
let decompressed = zstd::decode_all(&buf[..]).unwrap();
let req = serde_json::from_slice(&decompressed).unwrap();
let (res_tx, res_rx) = oneshot::channel::<PredictEditsV4Response>();
predict_v4_req_tx.unbounded_send((req, res_tx)).unwrap();
let response = res_rx.await?;
if response.request_id == REQUEST_TIMEOUT_RESPONSE_ID {
return Ok(Response::builder()
.status(http_client::http::StatusCode::REQUEST_TIMEOUT)
.body(
http_client::http::StatusCode::REQUEST_TIMEOUT
.as_str()
.into(),
)
.unwrap());
}
serde_json::to_string(&response).unwrap()
}
"/predict_edits/reject" => {
let mut buf = Vec::new();
body.read_to_end(&mut buf).await.ok();
@@ -2638,6 +2669,7 @@ fn init_test_with_fake_client_and_legacy_data_collection(
user_store,
RequestChannels {
predict: predict_req_rx,
predict_v4: predict_v4_req_rx,
reject: reject_req_rx,
settled: settled_req_rx,
},
@@ -2867,6 +2899,63 @@ async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
);
}
#[gpui::test]
async fn test_edit_prediction_v4_end_of_buffer(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/project",
json!({
"file.txt": "lorem\n"
}),
)
.await;
let project = Project::test(fs, [path!("/project").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
let path = project
.find_project_path(path!("/project/file.txt"), cx)
.unwrap();
project.open_buffer(path, cx)
})
.await
.unwrap();
let (ep_store, response) = make_test_ep_store(&project, cx).await;
*response.lock() = "lorem\nipsum\n".to_string();
let position = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
ep_store.update(cx, |ep_store, cx| {
ep_store.register_project(&project, cx);
ep_store.register_buffer(&buffer, &project, cx);
ep_store.refresh_prediction_from_buffer(
project.clone(),
buffer.clone(),
position,
EditPredictionRequestTrigger::Other,
cx,
);
});
cx.run_until_parked();
let edits = ep_store.update(cx, |ep_store, cx| {
let prediction = ep_store
.prediction_at(&buffer, None, &project, cx)
.expect("should have prediction");
let prediction = match prediction {
BufferEditPrediction::Local { prediction }
| BufferEditPrediction::Jump { prediction } => prediction,
};
assert!(prediction.editable_range.is_some());
prediction.edits.iter().cloned().collect::<Vec<_>>()
});
buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
buffer.read_with(cx, |buffer, _| {
assert_eq!(buffer.text(), "lorem\nipsum\n");
});
}
#[gpui::test]
async fn test_edit_prediction_no_spurious_trailing_newline(cx: &mut TestAppContext) {
// Test that zeta2's newline normalization logic doesn't insert spurious newlines.
@@ -3013,6 +3102,26 @@ fn init_test(cx: &mut TestAppContext) {
cx.set_global(AppDatabase::test_new());
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
disable_jumps_feature_flag(cx);
});
}
fn disable_jumps_feature_flag(cx: &mut App) {
SettingsStore::update_global(cx, |store, _| {
store.register_setting::<FeatureFlagsSettings>();
});
set_jumps_feature_flag_override(cx, "off");
cx.update_flags(false, vec![]);
}
fn set_jumps_feature_flag_override(cx: &mut App, value: &str) {
SettingsStore::update_global(cx, |store, cx| {
store.update_user_settings(cx, |content| {
content.feature_flags.get_or_insert_default().insert(
EditPredictionJumpsFeatureFlag::NAME.to_string(),
value.to_string(),
);
});
});
}
@@ -3105,6 +3214,43 @@ async fn make_test_ep_store(
)
.unwrap())
}
(Method::POST, "/predict_edits/v4") => {
let mut buf = Vec::new();
body.read_to_end(&mut buf).await.ok();
let decompressed = zstd::decode_all(&buf[..]).unwrap();
let req: PredictEditsV4Request =
serde_json::from_slice(&decompressed).unwrap();
next_request_id += 1;
let output = completion_response.lock().clone();
let file = req
.input
.editable_context
.iter()
.find(|file| file.path == req.input.cursor_path)
.or_else(|| req.input.editable_context.first())
.expect("V4 requests should include editable context");
let excerpt = file
.excerpts
.first()
.expect("V4 editable context should include an excerpt");
let diff = unified_diff_with_offsets(
&excerpt.text,
&output,
excerpt.row_range.start,
excerpt.row_range.start,
);
let path = file.path.to_string_lossy();
let response = PredictEditsV4Response {
request_id: format!("request-{next_request_id}"),
patch: format!("--- a/{path}\n+++ b/{path}\n{diff}"),
model_version: None,
};
Ok(http_client::Response::builder()
.status(200)
.body(serde_json::to_string(&response).unwrap().into())
.unwrap())
}
_ => Ok(http_client::Response::builder()
.status(404)
.body("Not Found".to_string().into())
crates/edit_prediction/src/mercury.rs
+5 -3
View File
@@ -1,7 +1,9 @@
use crate::{
DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionInputs,
EditPredictionModelInput, EditPredictionStartedDebugEvent, EditPredictionStore,
open_ai_response::text_from_response, prediction::EditPredictionResult, zeta::compute_edits,
DebugEvent, EditPredictionFinishedDebugEvent, EditPredictionId, EditPredictionModelInput,
EditPredictionStartedDebugEvent, EditPredictionStore,
open_ai_response::text_from_response,
prediction::{EditPredictionInputs, EditPredictionResult},
zeta::compute_edits,
};
use anyhow::{Context as _, Result};
use cloud_llm_client::EditPredictionRejectReason;
crates/edit_prediction/src/udiff.rs
+231 -19
View File
@@ -2,14 +2,18 @@ use std::{mem, ops::Range, path::Path, path::PathBuf, sync::Arc};
use anyhow::{Context as _, Result, anyhow};
use collections::{HashMap, hash_map::Entry};
use edit_prediction_types::PredictedCursorPosition;
use gpui::{AsyncApp, Entity};
use language::{Anchor, Buffer, OffsetRangeExt as _, TextBufferSnapshot, text_diff};
use language::{
Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, TextBufferSnapshot, text_diff,
};
use postage::stream::Stream as _;
use project::Project;
use util::{paths::PathStyle, rel_path::RelPath};
use worktree::Worktree;
use zeta_prompt::udiff::{
DiffEvent, DiffParser, FileStatus, Hunk, disambiguate_by_line_number, find_context_candidates,
DiffEvent, DiffParser, FileStatus, Hunk, INLINE_CURSOR_MARKER, disambiguate_by_line_number,
find_context_candidates,
};
pub use zeta_prompt::udiff::{
@@ -30,6 +34,83 @@ impl OpenedBuffers {
}
}
pub async fn prediction_edits_for_single_file_diff(
diff_str: &str,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<
Option<(
Entity<Buffer>,
BufferSnapshot,
Vec<(Range<Anchor>, Arc<str>)>,
Option<PredictedCursorPosition>,
)>,
> {
let mut diff = DiffParser::new(diff_str);
let mut target_file = None;
let mut edits = Vec::new();
let mut cursor_position = None;
while let Some(event) = diff.next()? {
match event {
DiffEvent::Hunk { path, hunk, status } => {
anyhow::ensure!(
status == FileStatus::Modified,
"V4 edit predictions only support modifying existing files"
);
let path = path.to_string();
if let Some((target_path, _, _)) = &target_file {
anyhow::ensure!(
target_path == &path,
"V4 edit predictions only support one file"
);
} else {
let project_path = project
.update(cx, |project, cx| {
project.find_project_path(Path::new(&path), cx)
})
.with_context(|| format!("no such path: {path}"))?;
let buffer = project
.update(cx, |project, cx| project.open_buffer(project_path, cx))
.await?;
let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
target_file = Some((path, buffer, snapshot));
}
let (_, _, snapshot) = target_file.as_ref().context("missing target file")?;
for (range, text) in resolve_hunk_edits_in_buffer(
hunk,
snapshot,
&[Anchor::min_max_range_for_buffer(snapshot.remote_id())],
status,
)? {
if let Some(marker_offset) = text.find(INLINE_CURSOR_MARKER) {
cursor_position.get_or_insert_with(|| {
PredictedCursorPosition::new(range.start, marker_offset)
});
let text = text.replace(INLINE_CURSOR_MARKER, "");
if range.start != range.end || !text.is_empty() {
edits.push((range, text.into()));
}
} else {
edits.push((range, text));
}
}
}
DiffEvent::FileEnd { renamed_to } => {
anyhow::ensure!(
renamed_to.is_none(),
"V4 edit predictions do not support renames"
);
}
}
}
Ok(target_file.map(|(_, buffer, snapshot)| (buffer, snapshot, edits, cursor_position)))
}
#[must_use]
pub async fn apply_diff(
diff_str: &str,
@@ -243,7 +324,7 @@ fn resolve_hunk_edits_in_buffer(
buffer: &TextBufferSnapshot,
ranges: &[Range<Anchor>],
status: FileStatus,
) -> Result<impl Iterator<Item = (Range<Anchor>, Arc<str>)>, anyhow::Error> {
) -> Result<Vec<(Range<Anchor>, Arc<str>)>, anyhow::Error> {
let context_offset = if status == FileStatus::Created || hunk.context.is_empty() {
0
} else {
@@ -272,22 +353,26 @@ fn resolve_hunk_edits_in_buffer(
return Err(anyhow!("Edit range {:?} exceeds buffer length", edit.range));
}
let iter = hunk.edits.into_iter().flat_map(move |edit| {
let old_text = buffer
.text_for_range(context_offset + edit.range.start..context_offset + edit.range.end)
.collect::<String>();
let edits_within_hunk = language::text_diff(&old_text, &edit.text);
edits_within_hunk
.into_iter()
.map(move |(inner_range, inner_text)| {
(
buffer.anchor_after(context_offset + edit.range.start + inner_range.start)
..buffer.anchor_before(context_offset + edit.range.start + inner_range.end),
inner_text,
)
})
});
Ok(iter)
Ok(hunk
.edits
.into_iter()
.flat_map(move |edit| {
let old_text = buffer
.text_for_range(context_offset + edit.range.start..context_offset + edit.range.end)
.collect::<String>();
let edits_within_hunk = language::text_diff(&old_text, &edit.text);
edits_within_hunk
.into_iter()
.map(move |(inner_range, inner_text)| {
(
buffer.anchor_after(context_offset + edit.range.start + inner_range.start)
..buffer
.anchor_before(context_offset + edit.range.start + inner_range.end),
inner_text,
)
})
})
.collect())
}
#[cfg(test)]
@@ -295,10 +380,12 @@ mod tests {
use super::*;
use gpui::TestAppContext;
use indoc::indoc;
use language::ToOffset as _;
use pretty_assertions::assert_eq;
use project::{FakeFs, Project};
use serde_json::json;
use settings::SettingsStore;
use std::path::Path;
use util::path;
#[test]
@@ -363,6 +450,131 @@ mod tests {
assert_eq!(text, "REPLACED");
}
#[gpui::test]
async fn test_prediction_edits_for_single_file_diff_can_target_project_file(
cx: &mut TestAppContext,
) {
let fs = init_test(cx);
fs.insert_tree(
path!("/root"),
json!({
"file1": "Hello!\nHow\nBye\n",
"file2": "Hola!\nComo\nAdios\n",
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/file2
+++ b/file2
@@ ... @@
Hola!
-Como
+Como estas?
Adios
"#};
let (buffer, snapshot, edits, cursor_position) =
prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())
.await
.unwrap()
.unwrap();
assert!(cursor_position.is_none());
buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
assert_eq!(
snapshot.file().unwrap().path().as_std_path(),
Path::new("file2")
);
buffer.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), "Hola!\nComo estas?\nAdios\n");
});
}
#[gpui::test]
async fn test_prediction_edits_for_single_file_diff_strips_inline_cursor_marker(
cx: &mut TestAppContext,
) {
let fs = init_test(cx);
fs.insert_tree(
path!("/root"),
json!({
"file": "Hello!\nHow\nBye\n",
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/file
+++ b/file
@@ ... @@
Hello!
-How
+How are <|user_cursor|>you?
Bye
"#};
let (buffer, snapshot, edits, cursor_position) =
prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())
.await
.unwrap()
.unwrap();
assert!(
edits
.iter()
.all(|(_, text)| !text.contains(INLINE_CURSOR_MARKER))
);
let cursor_position = cursor_position.unwrap();
assert_eq!(
cursor_position.anchor.to_offset(&snapshot),
"Hello!\nHow".len()
);
assert_eq!(cursor_position.offset, " are ".len());
buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
buffer.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), "Hello!\nHow are you?\nBye\n");
});
}
#[gpui::test]
async fn test_prediction_edits_for_single_file_diff_does_not_treat_completed_literal_marker_as_cursor(
cx: &mut TestAppContext,
) {
let fs = init_test(cx);
fs.insert_tree(
path!("/root"),
json!({
"file": "text <|user_cursor\n",
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let diff = indoc! {r#"
--- a/file
+++ b/file
@@ ... @@
-text <|user_cursor
+text <|user_cursor|>
"#};
let (buffer, _, edits, cursor_position) =
prediction_edits_for_single_file_diff(diff, &project, &mut cx.to_async())
.await
.unwrap()
.unwrap();
assert!(cursor_position.is_none());
buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
buffer.read_with(cx, |buffer, _cx| {
assert_eq!(buffer.text(), "text <|user_cursor|>\n");
});
}
#[gpui::test]
async fn test_apply_diff_successful(cx: &mut TestAppContext) {
let fs = init_test(cx);
crates/edit_prediction/src/zeta.rs
+247 -67
View File
@@ -6,8 +6,9 @@ use crate::{
cursor_excerpt::{self, compute_cursor_excerpt, compute_syntax_ranges},
data_collection::CapturedPredictionContext,
prediction::EditPredictionResult,
udiff::prediction_edits_for_single_file_diff,
};
use anyhow::Result;
use anyhow::{Context as _, Result};
use cloud_llm_client::{AcceptEditPredictionBody, predict_edits_v3::RawCompletionRequest};
use edit_prediction_types::PredictedCursorPosition;
use gpui::{App, AppContext as _, Entity, Task, TaskExt, WeakEntity, prelude::*};
@@ -15,6 +16,7 @@ use language::{
Buffer, BufferSnapshot, DiagnosticSeverity, EditPredictionPromptFormat, OffsetRangeExt as _,
ToOffset as _, ZetaVersion, language_settings::all_language_settings, text_diff,
};
use project::Project;
use release_channel::AppVersion;
use text::{Anchor, Bias, Point};
use ui::SharedString;
@@ -24,8 +26,9 @@ use workspace::workspace_error::{ErrorAction, ErrorSeverity, WorkspaceError};
use std::{ops::Range, path::Path, sync::Arc};
use zeta_prompt::{
ParsedOutput, Zeta2PromptInput, ZetaFormat, excerpt_ranges_for_format, format_zeta_prompt,
get_prefill, parse_zeta2_model_output, stop_tokens_for_format,
FilePosition, ParsedOutput, Zeta2PromptInput, Zeta3PromptInput, ZetaFormat,
excerpt_ranges_for_format, format_zeta_prompt, get_prefill, parse_zeta2_model_output,
stop_tokens_for_format,
zeta1::{self, EDITABLE_REGION_END_MARKER},
};
@@ -40,6 +43,7 @@ pub(crate) fn request_prediction_with_zeta(
snapshot,
position,
related_files,
editable_context,
events,
debug_tx,
mode,
@@ -48,6 +52,7 @@ pub(crate) fn request_prediction_with_zeta(
diagnostic_search_range,
can_collect_data,
is_open_source,
allow_jump,
..
}: EditPredictionModelInput,
context_task: Option<Task<Result<CapturedPredictionContext>>>,
@@ -83,15 +88,25 @@ pub(crate) fn request_prediction_with_zeta(
.map(|organization| organization.id.clone());
let app_version = AppVersion::global(cx);
struct Prediction {
prompt_input: Zeta2PromptInput,
buffer: Entity<Buffer>,
snapshot: BufferSnapshot,
edits: Vec<(Range<Anchor>, Arc<str>)>,
cursor_position: Option<PredictedCursorPosition>,
editable_range_in_buffer: Range<usize>,
enum PendingPrediction {
Resolved {
inputs: EditPredictionInputs,
buffer: Entity<Buffer>,
snapshot: BufferSnapshot,
edits: Vec<(Range<Anchor>, Arc<str>)>,
cursor_position: Option<PredictedCursorPosition>,
editable_range_in_buffer: Option<Range<usize>>,
},
Patch {
inputs: EditPredictionInputs,
project: Entity<Project>,
buffer: Entity<Buffer>,
snapshot: BufferSnapshot,
patch: String,
},
}
let project_for_request = project.clone();
let request_task = cx.background_spawn({
async move {
let local_zeta_version = custom_server_settings
@@ -119,27 +134,78 @@ pub(crate) fn request_prediction_with_zeta(
})
.unwrap_or_default();
let cursor_offset = position.to_offset(&snapshot);
let (full_context_offset_range, prompt_input) = zeta2_prompt_input(
&snapshot,
related_files,
events,
diagnostic_search_range,
excerpt_path,
cursor_offset,
is_open_source,
can_collect_data,
repo_url,
);
enum RequestInput {
V3 {
full_context_offset_range: Range<usize>,
prompt_input: Zeta2PromptInput,
formatted_prompt: Option<String>,
},
V4(Zeta3PromptInput),
}
let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format);
let cursor_offset = position.to_offset(&snapshot);
let request_input =
if allow_jump && custom_server_settings.is_none() && raw_config.is_none() {
let cursor_point = snapshot.offset_to_point(cursor_offset);
let editable_context = editable_context
.context("missing editable context task")?
.await?;
let active_buffer_diagnostics = active_buffer_diagnostics(
&snapshot,
diagnostic_search_range,
cursor_point.row,
ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
);
let syntax_ranges =
compute_syntax_ranges(&snapshot, cursor_offset, &(0..snapshot.len()));
RequestInput::V4(Zeta3PromptInput {
cursor_path: excerpt_path,
cursor_position: FilePosition {
row: cursor_point.row,
column: cursor_point.column,
},
events,
editable_context,
syntax_ranges,
active_buffer_diagnostics,
in_open_source_repo: is_open_source,
can_collect_data,
repo_url,
})
} else {
let (full_context_offset_range, prompt_input) = zeta2_prompt_input(
&snapshot,
related_files,
events,
diagnostic_search_range,
excerpt_path,
cursor_offset,
is_open_source,
can_collect_data,
repo_url,
);
let formatted_prompt = format_zeta_prompt(&prompt_input, zeta_format);
RequestInput::V3 {
full_context_offset_range,
prompt_input,
formatted_prompt,
}
};
if let Some(debug_tx) = &debug_tx {
let prompt = match &request_input {
RequestInput::V3 {
formatted_prompt, ..
} => formatted_prompt.clone(),
RequestInput::V4(_) => None,
};
debug_tx
.unbounded_send(DebugEvent::EditPredictionStarted(
EditPredictionStartedDebugEvent {
buffer: buffer.downgrade(),
prompt: formatted_prompt.clone(),
prompt,
position,
},
))
@@ -148,6 +214,45 @@ pub(crate) fn request_prediction_with_zeta(
log::trace!("Sending edit prediction request");
let (full_context_offset_range, prompt_input, formatted_prompt) = match request_input {
RequestInput::V3 {
full_context_offset_range,
prompt_input,
formatted_prompt,
} => (full_context_offset_range, prompt_input, formatted_prompt),
RequestInput::V4(v4_prompt_input) => {
let inputs = EditPredictionInputs::V3(v4_prompt_input.clone());
let (response, usage) = EditPredictionStore::send_v4_request(
v4_prompt_input,
preferred_experiment.clone(),
client,
llm_token,
organization_id,
app_version,
trigger,
mode,
)
.await?;
let request_id = EditPredictionId(response.request_id.into());
let model_version = response.model_version;
return anyhow::Ok((
Some((
request_id,
Some(PendingPrediction::Patch {
inputs,
project: project_for_request,
buffer,
snapshot: snapshot.clone(),
patch: response.patch,
}),
model_version,
)),
usage,
));
}
};
let Some((request_id, output, model_version, usage)) =
(if let Some(custom_settings) = &custom_server_settings {
let max_tokens = custom_settings.max_output_tokens * 4;
@@ -272,7 +377,6 @@ pub(crate) fn request_prediction_with_zeta(
Some((request_id, output, None, usage))
} else {
// Use V3 endpoint - server handles model/version selection and suffix stripping
let (response, usage) = EditPredictionStore::send_v3_request(
prompt_input.clone(),
preferred_experiment.clone(),
@@ -357,14 +461,14 @@ pub(crate) fn request_prediction_with_zeta(
})
.unwrap_or_default();
let prediction = Prediction {
prompt_input,
let prediction = Some(PendingPrediction::Resolved {
inputs: EditPredictionInputs::V2(prompt_input),
buffer,
snapshot: snapshot.clone(),
edits,
cursor_position,
editable_range_in_buffer,
};
editable_range_in_buffer: Some(editable_range_in_buffer),
});
anyhow::Ok((Some((request_id, prediction, model_version)), usage))
}
@@ -378,25 +482,95 @@ pub(crate) fn request_prediction_with_zeta(
};
let request_duration = cx.background_executor().now() - request_start;
let Prediction {
prompt_input: inputs,
buffer: edited_buffer,
snapshot: edited_buffer_snapshot,
let Some(prediction) = prediction else {
return Ok(None);
};
let (
inputs,
edited_buffer,
edited_buffer_snapshot,
edits,
cursor_position,
editable_range_in_buffer,
} = prediction;
) = match prediction {
PendingPrediction::Resolved {
inputs,
buffer,
snapshot,
edits,
cursor_position,
editable_range_in_buffer,
} => (
inputs,
buffer,
snapshot,
edits,
cursor_position,
editable_range_in_buffer,
),
PendingPrediction::Patch {
inputs,
project,
buffer: fallback_buffer,
snapshot: fallback_snapshot,
patch,
} => {
let Some((buffer, snapshot, edits, cursor_position)) =
prediction_edits_for_single_file_diff(&patch, &project, cx).await?
else {
return Ok(Some(
EditPredictionResult::new(
id,
&fallback_buffer,
&fallback_snapshot,
Arc::new([]),
None,
None,
inputs,
model_version,
trigger,
request_duration,
cx,
)
.await,
));
};
let editable_range_in_buffer =
edits
.iter()
.fold(None, |range: Option<Range<usize>>, (edit_range, _)| {
let edit_range = edit_range.start.to_offset(&snapshot)
..edit_range.end.to_offset(&snapshot);
Some(match range {
Some(range) => {
range.start.min(edit_range.start)..range.end.max(edit_range.end)
}
None => edit_range,
})
});
(
inputs,
buffer,
snapshot,
edits,
cursor_position,
editable_range_in_buffer,
)
}
};
let editable_anchor_range =
edited_buffer_snapshot.anchor_range_inside(editable_range_in_buffer.clone());
let result = EditPredictionResult::new(
id,
&edited_buffer,
&edited_buffer_snapshot,
edits.into(),
cursor_position,
Some(editable_anchor_range),
EditPredictionInputs::V2(inputs),
editable_range_in_buffer
.clone()
.map(|range| edited_buffer_snapshot.anchor_range_inside(range)),
inputs,
model_version,
trigger,
request_duration,
@@ -404,34 +578,35 @@ pub(crate) fn request_prediction_with_zeta(
)
.await;
let prediction = &result.prediction;
let weak_this = this.clone();
let request_id = prediction.id.clone();
let edited_buffer = edited_buffer.clone();
let edited_buffer_snapshot = edited_buffer_snapshot.clone();
let editable_range_in_buffer = editable_range_in_buffer.clone();
let edit_preview = prediction.edit_preview.clone();
let model_version = prediction.model_version.clone();
cx.spawn(async move |cx| {
weak_this
.update(cx, |this, cx| {
this.enqueue_settled_prediction(
request_id.clone(),
&project,
&edited_buffer,
&edited_buffer_snapshot,
editable_range_in_buffer,
&edit_preview,
context_task,
prompt_history_boundary,
model_version,
request_duration,
cx,
);
})
.ok();
})
.detach();
if let Some(editable_range_in_buffer) = editable_range_in_buffer.clone() {
let prediction = &result.prediction;
let weak_this = this.clone();
let request_id = prediction.id.clone();
let edited_buffer = edited_buffer.clone();
let edited_buffer_snapshot = edited_buffer_snapshot.clone();
let edit_preview = prediction.edit_preview.clone();
let model_version = prediction.model_version.clone();
cx.spawn(async move |cx| {
weak_this
.update(cx, |this, cx| {
this.enqueue_settled_prediction(
request_id.clone(),
&project,
&edited_buffer,
&edited_buffer_snapshot,
editable_range_in_buffer,
&edit_preview,
context_task,
prompt_history_boundary,
model_version,
request_duration,
cx,
);
})
.ok();
})
.detach();
}
Ok(Some(result))
})
@@ -594,7 +769,7 @@ pub(crate) fn active_buffer_diagnostics(
pub fn zeta2_prompt_input(
snapshot: &language::BufferSnapshot,
related_files: Vec<zeta_prompt::RelatedFile>,
mut related_files: Vec<zeta_prompt::RelatedFile>,
events: Vec<Arc<zeta_prompt::Event>>,
diagnostic_search_range: Range<Point>,
excerpt_path: Arc<Path>,
@@ -623,6 +798,11 @@ pub fn zeta2_prompt_input(
snapshot.offset_to_point(cursor_offset).row,
ACTIVE_BUFFER_DIAGNOSTIC_ADDITIONAL_CONTEXT_TOKEN_COUNT,
);
for file in &mut related_files {
for excerpt in &mut file.excerpts {
excerpt.context_source = zeta_prompt::ContextSource::Lsp;
}
}
let prompt_input = zeta_prompt::Zeta2PromptInput {
cursor_path: excerpt_path,