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

Combine zeta and zeta2 edit prediction providers (#43284)

Browse Source 
We've realized that a lot of the logic within an
`EditPredictionProvider` is not specific to a particular edit prediction
model / service. Rather, it is just the generic state management
required to perform edit predictions at all in Zed. We want to move to a
setup where there's one "built-in" edit prediction provider in Zed,
which can be pointed at different edit prediction models. The only logic
that is different for different models is how we construct the prompt,
send the request, and parse the output.

This PR also changes the behavior of the staff-only `zeta2` feature flag
so that in only gates your *ability* to use Zeta2, but you can still use
your local settings file to choose between different edit prediction
models/services: zeta1, zeta2, and sweep.

This PR also makes zeta1's outcome reporting and prediction-rating
features work with all prediction models, not just zeta1.

To do:
* [x] remove duplicated logic around sending cloud requests between
zeta1 and zeta2
* [x] port the outcome reporting logic from zeta to zeta2.
* [x] get the "rate completions" modal working with all EP models
   * [x] display edit prediction diff
   * [x] show edit history events
* [x] remove the original `zeta` crate.

Release Notes:

- N/A

---------

Co-authored-by: Agus Zubiaga <agus@zed.dev>
Co-authored-by: Ben Kunkle <ben@zed.dev>
This commit is contained in:
Max Brunsfeld
2025-11-24 22:17:48 -08:00
committed by GitHub
co-authored by Agus Zubiaga Ben Kunkle
parent 17d7988ad4
commit 9122dd2d70
41 changed files with 5030 additions and 6225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/zeta/src/assemble_excerpts.rs
+173
View File
@@ -0,0 +1,173 @@
use cloud_llm_client::predict_edits_v3::Excerpt;
use edit_prediction_context::Line;
use language::{BufferSnapshot, Point};
use std::ops::Range;
pub fn assemble_excerpts(
buffer: &BufferSnapshot,
merged_line_ranges: impl IntoIterator<Item = Range<Line>>,
) -> Vec<Excerpt> {
let mut output = Vec::new();
let outline_items = buffer.outline_items_as_points_containing(0..buffer.len(), false, None);
let mut outline_items = outline_items.into_iter().peekable();
for range in merged_line_ranges {
let point_range = Point::new(range.start.0, 0)..Point::new(range.end.0, 0);
while let Some(outline_item) = outline_items.peek() {
if outline_item.range.start >= point_range.start {
break;
}
if outline_item.range.end > point_range.start {
let mut point_range = outline_item.source_range_for_text.clone();
point_range.start.column = 0;
point_range.end.column = buffer.line_len(point_range.end.row);
output.push(Excerpt {
start_line: Line(point_range.start.row),
text: buffer
.text_for_range(point_range.clone())
.collect::<String>()
.into(),
})
}
outline_items.next();
}
output.push(Excerpt {
start_line: Line(point_range.start.row),
text: buffer
.text_for_range(point_range.clone())
.collect::<String>()
.into(),
})
}
output
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
use cloud_llm_client::predict_edits_v3;
use gpui::{TestAppContext, prelude::*};
use indoc::indoc;
use language::{Buffer, Language, LanguageConfig, LanguageMatcher, OffsetRangeExt};
use pretty_assertions::assert_eq;
use util::test::marked_text_ranges;
#[gpui::test]
fn test_rust(cx: &mut TestAppContext) {
let table = [
(
indoc! {r#"
struct User {
first_name: String,
« last_name: String,
ageˇ: u32,
» email: String,
create_at: Instant,
}
impl User {
pub fn first_name(&self) -> String {
self.first_name.clone()
}
pub fn full_name(&self) -> String {
« format!("{} {}", self.first_name, self.last_name)
» }
}
"#},
indoc! {r#"
1|struct User {
…
3| last_name: String,
4| age<|cursor|>: u32,
…
9|impl User {
…
14| pub fn full_name(&self) -> String {
15| format!("{} {}", self.first_name, self.last_name)
…
"#},
),
(
indoc! {r#"
struct User {
first_name: String,
« last_name: String,
age: u32,
}
»"#
},
indoc! {r#"
1|struct User {
…
3| last_name: String,
4| age: u32,
5|}
"#},
),
];
for (input, expected_output) in table {
let input_without_ranges = input.replace(['«', '»'], "");
let input_without_caret = input.replace('ˇ', "");
let cursor_offset = input_without_ranges.find('ˇ');
let (input, ranges) = marked_text_ranges(&input_without_caret, false);
let buffer =
cx.new(|cx| Buffer::local(input, cx).with_language(Arc::new(rust_lang()), cx));
buffer.read_with(cx, |buffer, _cx| {
let insertions = cursor_offset
.map(|offset| {
let point = buffer.offset_to_point(offset);
vec![(
predict_edits_v3::Point {
line: Line(point.row),
column: point.column,
},
"<|cursor|>",
)]
})
.unwrap_or_default();
let ranges: Vec<Range<Line>> = ranges
.into_iter()
.map(|range| {
let point_range = range.to_point(&buffer);
Line(point_range.start.row)..Line(point_range.end.row)
})
.collect();
let mut output = String::new();
cloud_zeta2_prompt::write_excerpts(
assemble_excerpts(&buffer.snapshot(), ranges).iter(),
&insertions,
Line(buffer.max_point().row),
true,
&mut output,
);
assert_eq!(output, expected_output);
});
}
}
fn rust_lang() -> Language {
Language::new(
LanguageConfig {
name: "Rust".into(),
matcher: LanguageMatcher {
path_suffixes: vec!["rs".to_string()],
..Default::default()
},
..Default::default()
},
Some(language::tree_sitter_rust::LANGUAGE.into()),
)
.with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
.unwrap()
}
}
crates/zeta/src/completion_diff_element.rs
-173
View File
@@ -1,173 +0,0 @@
use std::cmp;
use crate::EditPrediction;
use gpui::{
AnyElement, App, BorderStyle, Bounds, Corners, Edges, HighlightStyle, Hsla, StyledText,
TextLayout, TextStyle, point, prelude::*, quad, size,
};
use language::OffsetRangeExt;
use settings::Settings;
use theme::ThemeSettings;
use ui::prelude::*;
pub struct CompletionDiffElement {
element: AnyElement,
text_layout: TextLayout,
cursor_offset: usize,
}
impl CompletionDiffElement {
pub fn new(completion: &EditPrediction, cx: &App) -> Self {
let mut diff = completion
.snapshot
.text_for_range(completion.excerpt_range.clone())
.collect::<String>();
let mut cursor_offset_in_diff = None;
let mut delta = 0;
let mut diff_highlights = Vec::new();
for (old_range, new_text) in completion.edits.iter() {
let old_range = old_range.to_offset(&completion.snapshot);
if cursor_offset_in_diff.is_none() && completion.cursor_offset <= old_range.end {
cursor_offset_in_diff =
Some(completion.cursor_offset - completion.excerpt_range.start + delta);
}
let old_start_in_diff = old_range.start - completion.excerpt_range.start + delta;
let old_end_in_diff = old_range.end - completion.excerpt_range.start + delta;
if old_start_in_diff < old_end_in_diff {
diff_highlights.push((
old_start_in_diff..old_end_in_diff,
HighlightStyle {
background_color: Some(cx.theme().status().deleted_background),
strikethrough: Some(gpui::StrikethroughStyle {
thickness: px(1.),
color: Some(cx.theme().colors().text_muted),
}),
..Default::default()
},
));
}
if !new_text.is_empty() {
diff.insert_str(old_end_in_diff, new_text);
diff_highlights.push((
old_end_in_diff..old_end_in_diff + new_text.len(),
HighlightStyle {
background_color: Some(cx.theme().status().created_background),
..Default::default()
},
));
delta += new_text.len();
}
}
let cursor_offset_in_diff = cursor_offset_in_diff
.unwrap_or_else(|| completion.cursor_offset - completion.excerpt_range.start + delta);
let settings = ThemeSettings::get_global(cx).clone();
let text_style = TextStyle {
color: cx.theme().colors().editor_foreground,
font_size: settings.buffer_font_size(cx).into(),
font_family: settings.buffer_font.family,
font_features: settings.buffer_font.features,
font_fallbacks: settings.buffer_font.fallbacks,
line_height: relative(settings.buffer_line_height.value()),
font_weight: settings.buffer_font.weight,
font_style: settings.buffer_font.style,
..Default::default()
};
let element = StyledText::new(diff).with_default_highlights(&text_style, diff_highlights);
let text_layout = element.layout().clone();
CompletionDiffElement {
element: element.into_any_element(),
text_layout,
cursor_offset: cursor_offset_in_diff,
}
}
}
impl IntoElement for CompletionDiffElement {
type Element = Self;
fn into_element(self) -> Self {
self
}
}
impl Element for CompletionDiffElement {
type RequestLayoutState = ();
type PrepaintState = ();
fn id(&self) -> Option<ElementId> {
None
}
fn source_location(&self) -> Option<&'static core::panic::Location<'static>> {
None
}
fn request_layout(
&mut self,
_id: Option<&gpui::GlobalElementId>,
_inspector_id: Option<&gpui::InspectorElementId>,
window: &mut Window,
cx: &mut App,
) -> (gpui::LayoutId, Self::RequestLayoutState) {
(self.element.request_layout(window, cx), ())
}
fn prepaint(
&mut self,
_id: Option<&gpui::GlobalElementId>,
_inspector_id: Option<&gpui::InspectorElementId>,
_bounds: gpui::Bounds<Pixels>,
_request_layout: &mut Self::RequestLayoutState,
window: &mut Window,
cx: &mut App,
) -> Self::PrepaintState {
self.element.prepaint(window, cx);
}
fn paint(
&mut self,
_id: Option<&gpui::GlobalElementId>,
_inspector_id: Option<&gpui::InspectorElementId>,
_bounds: gpui::Bounds<Pixels>,
_request_layout: &mut Self::RequestLayoutState,
_prepaint: &mut Self::PrepaintState,
window: &mut Window,
cx: &mut App,
) {
if let Some(position) = self.text_layout.position_for_index(self.cursor_offset) {
let bounds = self.text_layout.bounds();
let line_height = self.text_layout.line_height();
let line_width = self
.text_layout
.line_layout_for_index(self.cursor_offset)
.map_or(bounds.size.width, |layout| layout.width());
window.paint_quad(quad(
Bounds::new(
point(bounds.origin.x, position.y),
size(cmp::max(bounds.size.width, line_width), line_height),
),
Corners::default(),
cx.theme().colors().editor_active_line_background,
Edges::default(),
Hsla::transparent_black(),
BorderStyle::default(),
));
self.element.paint(window, cx);
window.paint_quad(quad(
Bounds::new(position, size(px(2.), line_height)),
Corners::default(),
cx.theme().players().local().cursor,
Edges::default(),
Hsla::transparent_black(),
BorderStyle::default(),
));
}
}
}
crates/zeta/src/init.rs
-110
View File
@@ -1,110 +0,0 @@
use std::any::{Any, TypeId};
use command_palette_hooks::CommandPaletteFilter;
use feature_flags::{FeatureFlagAppExt as _, PredictEditsRateCompletionsFeatureFlag};
use gpui::actions;
use language::language_settings::EditPredictionProvider;
use project::DisableAiSettings;
use settings::{Settings, SettingsStore, update_settings_file};
use ui::App;
use workspace::Workspace;
use crate::{RateCompletionModal, onboarding_modal::ZedPredictModal};
actions!(
edit_prediction,
[
/// Resets the edit prediction onboarding state.
ResetOnboarding,
/// Opens the rate completions modal.
RateCompletions
]
);
pub fn init(cx: &mut App) {
feature_gate_predict_edits_actions(cx);
cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
workspace.register_action(|workspace, _: &RateCompletions, window, cx| {
if cx.has_flag::<PredictEditsRateCompletionsFeatureFlag>() {
RateCompletionModal::toggle(workspace, window, cx);
}
});
workspace.register_action(
move |workspace, _: &zed_actions::OpenZedPredictOnboarding, window, cx| {
ZedPredictModal::toggle(
workspace,
workspace.user_store().clone(),
workspace.client().clone(),
window,
cx,
)
},
);
workspace.register_action(|workspace, _: &ResetOnboarding, _window, cx| {
update_settings_file(workspace.app_state().fs.clone(), cx, move |settings, _| {
settings
.project
.all_languages
.features
.get_or_insert_default()
.edit_prediction_provider = Some(EditPredictionProvider::None)
});
});
})
.detach();
}
fn feature_gate_predict_edits_actions(cx: &mut App) {
let rate_completion_action_types = [TypeId::of::<RateCompletions>()];
let reset_onboarding_action_types = [TypeId::of::<ResetOnboarding>()];
let zeta_all_action_types = [
TypeId::of::<RateCompletions>(),
TypeId::of::<ResetOnboarding>(),
zed_actions::OpenZedPredictOnboarding.type_id(),
TypeId::of::<crate::ClearHistory>(),
TypeId::of::<crate::ThumbsUpActiveCompletion>(),
TypeId::of::<crate::ThumbsDownActiveCompletion>(),
TypeId::of::<crate::NextEdit>(),
TypeId::of::<crate::PreviousEdit>(),
];
CommandPaletteFilter::update_global(cx, |filter, _cx| {
filter.hide_action_types(&rate_completion_action_types);
filter.hide_action_types(&reset_onboarding_action_types);
filter.hide_action_types(&[zed_actions::OpenZedPredictOnboarding.type_id()]);
});
cx.observe_global::<SettingsStore>(move |cx| {
let is_ai_disabled = DisableAiSettings::get_global(cx).disable_ai;
let has_feature_flag = cx.has_flag::<PredictEditsRateCompletionsFeatureFlag>();
CommandPaletteFilter::update_global(cx, |filter, _cx| {
if is_ai_disabled {
filter.hide_action_types(&zeta_all_action_types);
} else if has_feature_flag {
filter.show_action_types(&rate_completion_action_types);
} else {
filter.hide_action_types(&rate_completion_action_types);
}
});
})
.detach();
cx.observe_flag::<PredictEditsRateCompletionsFeatureFlag, _>(move |is_enabled, cx| {
if !DisableAiSettings::get_global(cx).disable_ai {
if is_enabled {
CommandPaletteFilter::update_global(cx, |filter, _cx| {
filter.show_action_types(&rate_completion_action_types);
});
} else {
CommandPaletteFilter::update_global(cx, |filter, _cx| {
filter.hide_action_types(&rate_completion_action_types);
});
}
}
})
.detach();
}
crates/zeta/src/onboarding_modal.rs
+11 -1
View File
@@ -1,6 +1,6 @@
use std::sync::Arc;
use crate::{ZedPredictUpsell, onboarding_event};
use crate::ZedPredictUpsell;
use ai_onboarding::EditPredictionOnboarding;
use client::{Client, UserStore};
use db::kvp::Dismissable;
@@ -14,6 +14,16 @@ use settings::update_settings_file;
use ui::{Vector, VectorName, prelude::*};
use workspace::{ModalView, Workspace};
#[macro_export]
macro_rules! onboarding_event {
($name:expr) => {
telemetry::event!($name, source = "Edit Prediction Onboarding");
};
($name:expr, $($key:ident $(= $value:expr)?),+ $(,)?) => {
telemetry::event!($name, source = "Edit Prediction Onboarding", $($key $(= $value)?),+);
};
}
/// Introduces user to Zed's Edit Prediction feature
pub struct ZedPredictModal {
onboarding: Entity<EditPredictionOnboarding>,
crates/zeta/src/onboarding_telemetry.rs
-9
View File
@@ -1,9 +0,0 @@
#[macro_export]
macro_rules! onboarding_event {
($name:expr) => {
telemetry::event!($name, source = "Edit Prediction Onboarding");
};
($name:expr, $($key:ident $(= $value:expr)?),+ $(,)?) => {
telemetry::event!($name, source = "Edit Prediction Onboarding", $($key $(= $value)?),+);
};
}
crates/zeta/src/prediction.rs
+308
View File
@@ -0,0 +1,308 @@
use std::{
ops::Range,
path::Path,
sync::Arc,
time::{Duration, Instant},
};
use gpui::{AsyncApp, Entity, SharedString};
use language::{Anchor, Buffer, BufferSnapshot, EditPreview, OffsetRangeExt, TextBufferSnapshot};
use serde::Serialize;
#[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
pub struct EditPredictionId(pub SharedString);
impl From<EditPredictionId> for gpui::ElementId {
fn from(value: EditPredictionId) -> Self {
gpui::ElementId::Name(value.0)
}
}
impl std::fmt::Display for EditPredictionId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Clone)]
pub struct EditPrediction {
pub id: EditPredictionId,
pub edits: Arc<[(Range<Anchor>, Arc<str>)]>,
pub snapshot: BufferSnapshot,
pub edit_preview: EditPreview,
// We keep a reference to the buffer so that we do not need to reload it from disk when applying the prediction.
pub buffer: Entity<Buffer>,
pub buffer_snapshotted_at: Instant,
pub response_received_at: Instant,
pub inputs: EditPredictionInputs,
}
#[derive(Debug, Clone, Serialize)]
pub struct EditPredictionInputs {
pub events: Vec<Arc<cloud_llm_client::predict_edits_v3::Event>>,
pub included_files: Vec<cloud_llm_client::predict_edits_v3::IncludedFile>,
pub cursor_point: cloud_llm_client::predict_edits_v3::Point,
pub cursor_path: Arc<Path>,
}
impl EditPrediction {
pub async fn new(
id: EditPredictionId,
edited_buffer: &Entity<Buffer>,
edited_buffer_snapshot: &BufferSnapshot,
edits: Arc<[(Range<Anchor>, Arc<str>)]>,
buffer_snapshotted_at: Instant,
response_received_at: Instant,
inputs: EditPredictionInputs,
cx: &mut AsyncApp,
) -> Option<Self> {
let (edits, snapshot, edit_preview_task) = edited_buffer
.read_with(cx, |buffer, cx| {
let new_snapshot = buffer.snapshot();
let edits: Arc<[_]> =
interpolate_edits(&edited_buffer_snapshot, &new_snapshot, edits)?.into();
Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
})
.ok()??;
let edit_preview = edit_preview_task.await;
Some(EditPrediction {
id,
edits,
snapshot,
edit_preview,
inputs,
buffer: edited_buffer.clone(),
buffer_snapshotted_at,
response_received_at,
})
}
pub fn interpolate(
&self,
new_snapshot: &TextBufferSnapshot,
) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
interpolate_edits(&self.snapshot, new_snapshot, self.edits.clone())
}
pub fn targets_buffer(&self, buffer: &Buffer) -> bool {
self.snapshot.remote_id() == buffer.remote_id()
}
pub fn latency(&self) -> Duration {
self.response_received_at - self.buffer_snapshotted_at
}
}
impl std::fmt::Debug for EditPrediction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EditPrediction")
.field("id", &self.id)
.field("edits", &self.edits)
.finish()
}
}
pub fn interpolate_edits(
old_snapshot: &TextBufferSnapshot,
new_snapshot: &TextBufferSnapshot,
current_edits: Arc<[(Range<Anchor>, Arc<str>)]>,
) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
let mut edits = Vec::new();
let mut model_edits = current_edits.iter().peekable();
for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
while let Some((model_old_range, _)) = model_edits.peek() {
let model_old_range = model_old_range.to_offset(old_snapshot);
if model_old_range.end < user_edit.old.start {
let (model_old_range, model_new_text) = model_edits.next().unwrap();
edits.push((model_old_range.clone(), model_new_text.clone()));
} else {
break;
}
}
if let Some((model_old_range, model_new_text)) = model_edits.peek() {
let model_old_offset_range = model_old_range.to_offset(old_snapshot);
if user_edit.old == model_old_offset_range {
let user_new_text = new_snapshot
.text_for_range(user_edit.new.clone())
.collect::<String>();
if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
if !model_suffix.is_empty() {
let anchor = old_snapshot.anchor_after(user_edit.old.end);
edits.push((anchor..anchor, model_suffix.into()));
}
model_edits.next();
continue;
}
}
}
return None;
}
edits.extend(model_edits.cloned());
if edits.is_empty() { None } else { Some(edits) }
}
#[cfg(test)]
mod tests {
use super::*;
use gpui::{App, Entity, TestAppContext, prelude::*};
use language::{Buffer, ToOffset as _};
#[gpui::test]
async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
to_prediction_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
});
let edit_preview = cx
.read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
.await;
let prediction = EditPrediction {
id: EditPredictionId("prediction-1".into()),
edits,
snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
buffer: buffer.clone(),
edit_preview,
inputs: EditPredictionInputs {
events: vec![],
included_files: vec![],
cursor_point: cloud_llm_client::predict_edits_v3::Point {
line: cloud_llm_client::predict_edits_v3::Line(0),
column: 0,
},
cursor_path: Path::new("path.txt").into(),
},
buffer_snapshotted_at: Instant::now(),
response_received_at: Instant::now(),
};
cx.update(|cx| {
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..5, "REM".into()), (9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..2, "REM".into()), (6..8, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.undo(cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..5, "REM".into()), (9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(3..3, "EM".into()), (7..9, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into()), (8..10, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into()), (8..10, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
assert_eq!(
from_prediction_edits(
&prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
})
}
fn to_prediction_edits(
iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
buffer: &Entity<Buffer>,
cx: &App,
) -> Vec<(Range<Anchor>, Arc<str>)> {
let buffer = buffer.read(cx);
iterator
.into_iter()
.map(|(range, text)| {
(
buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
text,
)
})
.collect()
}
fn from_prediction_edits(
editor_edits: &[(Range<Anchor>, Arc<str>)],
buffer: &Entity<Buffer>,
cx: &App,
) -> Vec<(Range<usize>, Arc<str>)> {
let buffer = buffer.read(cx);
editor_edits
.iter()
.map(|(range, text)| {
(
range.start.to_offset(buffer)..range.end.to_offset(buffer),
text.clone(),
)
})
.collect()
}
}
crates/zeta/src/provider.rs
+213
View File
@@ -0,0 +1,213 @@
use std::{cmp, sync::Arc, time::Duration};
use client::{Client, UserStore};
use edit_prediction::{DataCollectionState, Direction, EditPredictionProvider};
use gpui::{App, Entity, prelude::*};
use language::ToPoint as _;
use project::Project;
use crate::{BufferEditPrediction, Zeta, ZetaEditPredictionModel};
pub struct ZetaEditPredictionProvider {
zeta: Entity<Zeta>,
project: Entity<Project>,
}
impl ZetaEditPredictionProvider {
pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
pub fn new(
project: Entity<Project>,
client: &Arc<Client>,
user_store: &Entity<UserStore>,
cx: &mut Context<Self>,
) -> Self {
let zeta = Zeta::global(client, user_store, cx);
zeta.update(cx, |zeta, cx| {
zeta.register_project(&project, cx);
});
cx.observe(&zeta, |_this, _zeta, cx| {
cx.notify();
})
.detach();
Self {
project: project,
zeta,
}
}
}
impl EditPredictionProvider for ZetaEditPredictionProvider {
fn name() -> &'static str {
"zed-predict2"
}
fn display_name() -> &'static str {
"Zed's Edit Predictions 2"
}
fn show_completions_in_menu() -> bool {
true
}
fn show_tab_accept_marker() -> bool {
true
}
fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
// TODO [zeta2]
DataCollectionState::Unsupported
}
fn toggle_data_collection(&mut self, _cx: &mut App) {
// TODO [zeta2]
}
fn usage(&self, cx: &App) -> Option<client::EditPredictionUsage> {
self.zeta.read(cx).usage(cx)
}
fn is_enabled(
&self,
_buffer: &Entity<language::Buffer>,
_cursor_position: language::Anchor,
cx: &App,
) -> bool {
let zeta = self.zeta.read(cx);
if zeta.edit_prediction_model == ZetaEditPredictionModel::Sweep {
zeta.sweep_api_token.is_some()
} else {
true
}
}
fn is_refreshing(&self, cx: &App) -> bool {
self.zeta.read(cx).is_refreshing(&self.project)
}
fn refresh(
&mut self,
buffer: Entity<language::Buffer>,
cursor_position: language::Anchor,
_debounce: bool,
cx: &mut Context<Self>,
) {
let zeta = self.zeta.read(cx);
if zeta.user_store.read_with(cx, |user_store, _cx| {
user_store.account_too_young() || user_store.has_overdue_invoices()
}) {
return;
}
if let Some(current) = zeta.current_prediction_for_buffer(&buffer, &self.project, cx)
&& let BufferEditPrediction::Local { prediction } = current
&& prediction.interpolate(buffer.read(cx)).is_some()
{
return;
}
self.zeta.update(cx, |zeta, cx| {
zeta.refresh_context_if_needed(&self.project, &buffer, cursor_position, cx);
zeta.refresh_prediction_from_buffer(self.project.clone(), buffer, cursor_position, cx)
});
}
fn cycle(
&mut self,
_buffer: Entity<language::Buffer>,
_cursor_position: language::Anchor,
_direction: Direction,
_cx: &mut Context<Self>,
) {
}
fn accept(&mut self, cx: &mut Context<Self>) {
self.zeta.update(cx, |zeta, cx| {
zeta.accept_current_prediction(&self.project, cx);
});
}
fn discard(&mut self, cx: &mut Context<Self>) {
self.zeta.update(cx, |zeta, cx| {
zeta.discard_current_prediction(&self.project, cx);
});
}
fn did_show(&mut self, cx: &mut Context<Self>) {
self.zeta.update(cx, |zeta, cx| {
zeta.did_show_current_prediction(&self.project, cx);
});
}
fn suggest(
&mut self,
buffer: &Entity<language::Buffer>,
cursor_position: language::Anchor,
cx: &mut Context<Self>,
) -> Option<edit_prediction::EditPrediction> {
let prediction =
self.zeta
.read(cx)
.current_prediction_for_buffer(buffer, &self.project, cx)?;
let prediction = match prediction {
BufferEditPrediction::Local { prediction } => prediction,
BufferEditPrediction::Jump { prediction } => {
return Some(edit_prediction::EditPrediction::Jump {
id: Some(prediction.id.to_string().into()),
snapshot: prediction.snapshot.clone(),
target: prediction.edits.first().unwrap().0.start,
});
}
};
let buffer = buffer.read(cx);
let snapshot = buffer.snapshot();
let Some(edits) = prediction.interpolate(&snapshot) else {
self.zeta.update(cx, |zeta, cx| {
zeta.discard_current_prediction(&self.project, cx);
});
return None;
};
let cursor_row = cursor_position.to_point(&snapshot).row;
let (closest_edit_ix, (closest_edit_range, _)) =
edits.iter().enumerate().min_by_key(|(_, (range, _))| {
let distance_from_start = cursor_row.abs_diff(range.start.to_point(&snapshot).row);
let distance_from_end = cursor_row.abs_diff(range.end.to_point(&snapshot).row);
cmp::min(distance_from_start, distance_from_end)
})?;
let mut edit_start_ix = closest_edit_ix;
for (range, _) in edits[..edit_start_ix].iter().rev() {
let distance_from_closest_edit = closest_edit_range.start.to_point(&snapshot).row
- range.end.to_point(&snapshot).row;
if distance_from_closest_edit <= 1 {
edit_start_ix -= 1;
} else {
break;
}
}
let mut edit_end_ix = closest_edit_ix + 1;
for (range, _) in &edits[edit_end_ix..] {
let distance_from_closest_edit =
range.start.to_point(buffer).row - closest_edit_range.end.to_point(&snapshot).row;
if distance_from_closest_edit <= 1 {
edit_end_ix += 1;
} else {
break;
}
}
Some(edit_prediction::EditPrediction::Local {
id: Some(prediction.id.to_string().into()),
edits: edits[edit_start_ix..edit_end_ix].to_vec(),
edit_preview: Some(prediction.edit_preview.clone()),
})
}
}
crates/zeta/src/rate_completion_modal.rs → crates/zeta/src/rate_prediction_modal.rs
+298 -126
View File
@@ -1,8 +1,18 @@
use crate::{CompletionDiffElement, EditPrediction, EditPredictionRating, Zeta};
use editor::Editor;
use gpui::{App, DismissEvent, Entity, EventEmitter, FocusHandle, Focusable, actions, prelude::*};
use language::language_settings;
use crate::{EditPrediction, EditPredictionRating, Zeta};
use buffer_diff::{BufferDiff, BufferDiffSnapshot};
use cloud_zeta2_prompt::write_codeblock;
use editor::{Editor, ExcerptRange, MultiBuffer};
use gpui::{
App, BorderStyle, DismissEvent, EdgesRefinement, Entity, EventEmitter, FocusHandle, Focusable,
Length, StyleRefinement, TextStyleRefinement, Window, actions, prelude::*,
};
use language::{LanguageRegistry, Point, language_settings};
use markdown::{Markdown, MarkdownStyle};
use settings::Settings as _;
use std::fmt::Write;
use std::sync::Arc;
use std::time::Duration;
use theme::ThemeSettings;
use ui::{KeyBinding, List, ListItem, ListItemSpacing, Tooltip, prelude::*};
use workspace::{ModalView, Workspace};
@@ -10,41 +20,44 @@ actions!(
zeta,
[
/// Rates the active completion with a thumbs up.
ThumbsUpActiveCompletion,
ThumbsUpActivePrediction,
/// Rates the active completion with a thumbs down.
ThumbsDownActiveCompletion,
ThumbsDownActivePrediction,
/// Navigates to the next edit in the completion history.
NextEdit,
/// Navigates to the previous edit in the completion history.
PreviousEdit,
/// Focuses on the completions list.
FocusCompletions,
FocusPredictions,
/// Previews the selected completion.
PreviewCompletion,
PreviewPrediction,
]
);
pub struct RateCompletionModal {
pub struct RatePredictionsModal {
zeta: Entity<Zeta>,
active_completion: Option<ActiveCompletion>,
language_registry: Arc<LanguageRegistry>,
active_prediction: Option<ActivePrediction>,
selected_index: usize,
diff_editor: Entity<Editor>,
focus_handle: FocusHandle,
_subscription: gpui::Subscription,
current_view: RateCompletionView,
current_view: RatePredictionView,
}
struct ActiveCompletion {
completion: EditPrediction,
struct ActivePrediction {
prediction: EditPrediction,
feedback_editor: Entity<Editor>,
formatted_inputs: Entity<Markdown>,
}
#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)]
enum RateCompletionView {
enum RatePredictionView {
SuggestedEdits,
RawInput,
}
impl RateCompletionView {
impl RatePredictionView {
pub fn name(&self) -> &'static str {
match self {
Self::SuggestedEdits => "Suggested Edits",
@@ -53,25 +66,42 @@ impl RateCompletionView {
}
}
impl RateCompletionModal {
impl RatePredictionsModal {
pub fn toggle(workspace: &mut Workspace, window: &mut Window, cx: &mut Context<Workspace>) {
if let Some(zeta) = Zeta::global(cx) {
workspace.toggle_modal(window, cx, |_window, cx| RateCompletionModal::new(zeta, cx));
if let Some(zeta) = Zeta::try_global(cx) {
let language_registry = workspace.app_state().languages.clone();
workspace.toggle_modal(window, cx, |window, cx| {
RatePredictionsModal::new(zeta, language_registry, window, cx)
});
telemetry::event!("Rate Completion Modal Open", source = "Edit Prediction");
telemetry::event!("Rate Prediction Modal Open", source = "Edit Prediction");
}
}
pub fn new(zeta: Entity<Zeta>, cx: &mut Context<Self>) -> Self {
pub fn new(
zeta: Entity<Zeta>,
language_registry: Arc<LanguageRegistry>,
window: &mut Window,
cx: &mut Context<Self>,
) -> Self {
let subscription = cx.observe(&zeta, |_, _, cx| cx.notify());
Self {
zeta,
language_registry,
selected_index: 0,
focus_handle: cx.focus_handle(),
active_completion: None,
active_prediction: None,
_subscription: subscription,
current_view: RateCompletionView::SuggestedEdits,
diff_editor: cx.new(|cx| {
let multibuffer = cx.new(|_| MultiBuffer::new(language::Capability::ReadOnly));
let mut editor = Editor::for_multibuffer(multibuffer, None, window, cx);
editor.disable_inline_diagnostics();
editor.set_expand_all_diff_hunks(cx);
editor.set_show_git_diff_gutter(false, cx);
editor
}),
current_view: RatePredictionView::SuggestedEdits,
}
}
@@ -83,7 +113,7 @@ impl RateCompletionModal {
self.selected_index += 1;
self.selected_index = usize::min(
self.selected_index,
self.zeta.read(cx).shown_completions().count(),
self.zeta.read(cx).shown_predictions().count(),
);
cx.notify();
}
@@ -102,7 +132,7 @@ impl RateCompletionModal {
let next_index = self
.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.skip(self.selected_index)
.enumerate()
.skip(1) // Skip straight to the next item
@@ -122,7 +152,7 @@ impl RateCompletionModal {
let prev_index = self
.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.rev()
.skip((completions_len - 1) - self.selected_index)
.enumerate()
@@ -149,14 +179,14 @@ impl RateCompletionModal {
pub fn thumbs_up_active(
&mut self,
_: &ThumbsUpActiveCompletion,
_: &ThumbsUpActivePrediction,
window: &mut Window,
cx: &mut Context<Self>,
) {
self.zeta.update(cx, |zeta, cx| {
if let Some(active) = &self.active_completion {
zeta.rate_completion(
&active.completion,
if let Some(active) = &self.active_prediction {
zeta.rate_prediction(
&active.prediction,
EditPredictionRating::Positive,
active.feedback_editor.read(cx).text(cx),
cx,
@@ -165,9 +195,9 @@ impl RateCompletionModal {
});
let current_completion = self
.active_completion
.active_prediction
.as_ref()
.map(|completion| completion.completion.clone());
.map(|completion| completion.prediction.clone());
self.select_completion(current_completion, false, window, cx);
self.select_next_edit(&Default::default(), window, cx);
self.confirm(&Default::default(), window, cx);
@@ -177,18 +207,18 @@ impl RateCompletionModal {
pub fn thumbs_down_active(
&mut self,
_: &ThumbsDownActiveCompletion,
_: &ThumbsDownActivePrediction,
window: &mut Window,
cx: &mut Context<Self>,
) {
if let Some(active) = &self.active_completion {
if let Some(active) = &self.active_prediction {
if active.feedback_editor.read(cx).text(cx).is_empty() {
return;
}
self.zeta.update(cx, |zeta, cx| {
zeta.rate_completion(
&active.completion,
zeta.rate_prediction(
&active.prediction,
EditPredictionRating::Negative,
active.feedback_editor.read(cx).text(cx),
cx,
@@ -197,9 +227,9 @@ impl RateCompletionModal {
}
let current_completion = self
.active_completion
.active_prediction
.as_ref()
.map(|completion| completion.completion.clone());
.map(|completion| completion.prediction.clone());
self.select_completion(current_completion, false, window, cx);
self.select_next_edit(&Default::default(), window, cx);
self.confirm(&Default::default(), window, cx);
@@ -209,7 +239,7 @@ impl RateCompletionModal {
fn focus_completions(
&mut self,
_: &FocusCompletions,
_: &FocusPredictions,
window: &mut Window,
cx: &mut Context<Self>,
) {
@@ -219,14 +249,14 @@ impl RateCompletionModal {
fn preview_completion(
&mut self,
_: &PreviewCompletion,
_: &PreviewPrediction,
window: &mut Window,
cx: &mut Context<Self>,
) {
let completion = self
.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.skip(self.selected_index)
.take(1)
.next()
@@ -239,7 +269,7 @@ impl RateCompletionModal {
let completion = self
.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.skip(self.selected_index)
.take(1)
.next()
@@ -250,54 +280,145 @@ impl RateCompletionModal {
pub fn select_completion(
&mut self,
completion: Option<EditPrediction>,
prediction: Option<EditPrediction>,
focus: bool,
window: &mut Window,
cx: &mut Context<Self>,
) {
// Avoid resetting completion rating if it's already selected.
if let Some(completion) = completion.as_ref() {
if let Some(prediction) = prediction {
self.selected_index = self
.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.enumerate()
.find(|(_, completion_b)| completion.id == completion_b.id)
.find(|(_, completion_b)| prediction.id == completion_b.id)
.map(|(ix, _)| ix)
.unwrap_or(self.selected_index);
cx.notify();
if let Some(prev_completion) = self.active_completion.as_ref()
&& completion.id == prev_completion.completion.id
if let Some(prev_prediction) = self.active_prediction.as_ref()
&& prediction.id == prev_prediction.prediction.id
{
if focus {
window.focus(&prev_completion.feedback_editor.focus_handle(cx));
window.focus(&prev_prediction.feedback_editor.focus_handle(cx));
}
return;
}
self.diff_editor.update(cx, |editor, cx| {
let new_buffer = prediction.edit_preview.build_result_buffer(cx);
let new_buffer_snapshot = new_buffer.read(cx).snapshot();
let old_buffer_snapshot = prediction.snapshot.clone();
let new_buffer_id = new_buffer_snapshot.remote_id();
let range = prediction
.edit_preview
.compute_visible_range(&prediction.edits)
.unwrap_or(Point::zero()..Point::zero());
let start = Point::new(range.start.row.saturating_sub(5), 0);
let end = Point::new(range.end.row + 5, 0).min(new_buffer_snapshot.max_point());
let diff = cx.new::<BufferDiff>(|cx| {
let diff_snapshot = BufferDiffSnapshot::new_with_base_buffer(
new_buffer_snapshot.text.clone(),
Some(old_buffer_snapshot.text().into()),
old_buffer_snapshot.clone(),
cx,
);
let diff = BufferDiff::new(&new_buffer_snapshot, cx);
cx.spawn(async move |diff, cx| {
let diff_snapshot = diff_snapshot.await;
diff.update(cx, |diff, cx| {
diff.set_snapshot(diff_snapshot, &new_buffer_snapshot.text, cx);
})
})
.detach();
diff
});
editor.disable_header_for_buffer(new_buffer_id, cx);
editor.buffer().update(cx, |multibuffer, cx| {
multibuffer.clear(cx);
multibuffer.push_excerpts(
new_buffer,
vec![ExcerptRange {
context: start..end,
primary: start..end,
}],
cx,
);
multibuffer.add_diff(diff, cx);
});
});
let mut formatted_inputs = String::new();
write!(&mut formatted_inputs, "## Events\n\n").unwrap();
for event in &prediction.inputs.events {
write!(&mut formatted_inputs, "```diff\n{event}```\n\n").unwrap();
}
write!(&mut formatted_inputs, "## Included files\n\n").unwrap();
for included_file in &prediction.inputs.included_files {
let cursor_insertions = &[(prediction.inputs.cursor_point, "<|CURSOR|>")];
write!(
&mut formatted_inputs,
"### {}\n\n",
included_file.path.display()
)
.unwrap();
write_codeblock(
&included_file.path,
&included_file.excerpts,
if included_file.path == prediction.inputs.cursor_path {
cursor_insertions
} else {
&[]
},
included_file.max_row,
false,
&mut formatted_inputs,
);
}
self.active_prediction = Some(ActivePrediction {
prediction,
feedback_editor: cx.new(|cx| {
let mut editor = Editor::multi_line(window, cx);
editor.disable_scrollbars_and_minimap(window, cx);
editor.set_soft_wrap_mode(language_settings::SoftWrap::EditorWidth, cx);
editor.set_show_line_numbers(false, cx);
editor.set_show_git_diff_gutter(false, cx);
editor.set_show_code_actions(false, cx);
editor.set_show_runnables(false, cx);
editor.set_show_breakpoints(false, cx);
editor.set_show_wrap_guides(false, cx);
editor.set_show_indent_guides(false, cx);
editor.set_show_edit_predictions(Some(false), window, cx);
editor.set_placeholder_text("Add your feedback…", window, cx);
if focus {
cx.focus_self(window);
}
editor
}),
formatted_inputs: cx.new(|cx| {
Markdown::new(
formatted_inputs.into(),
Some(self.language_registry.clone()),
None,
cx,
)
}),
});
} else {
self.active_prediction = None;
}
self.active_completion = completion.map(|completion| ActiveCompletion {
completion,
feedback_editor: cx.new(|cx| {
let mut editor = Editor::multi_line(window, cx);
editor.disable_scrollbars_and_minimap(window, cx);
editor.set_soft_wrap_mode(language_settings::SoftWrap::EditorWidth, cx);
editor.set_show_line_numbers(false, cx);
editor.set_show_git_diff_gutter(false, cx);
editor.set_show_code_actions(false, cx);
editor.set_show_runnables(false, cx);
editor.set_show_breakpoints(false, cx);
editor.set_show_wrap_guides(false, cx);
editor.set_show_indent_guides(false, cx);
editor.set_show_edit_predictions(Some(false), window, cx);
editor.set_placeholder_text("Add your feedback…", window, cx);
if focus {
cx.focus_self(window);
}
editor
}),
});
cx.notify();
}
@@ -312,33 +433,31 @@ impl RateCompletionModal {
.child(
Button::new(
ElementId::Name("suggested-edits".into()),
RateCompletionView::SuggestedEdits.name(),
RatePredictionView::SuggestedEdits.name(),
)
.label_size(LabelSize::Small)
.on_click(cx.listener(move |this, _, _window, cx| {
this.current_view = RateCompletionView::SuggestedEdits;
this.current_view = RatePredictionView::SuggestedEdits;
cx.notify();
}))
.toggle_state(self.current_view == RateCompletionView::SuggestedEdits),
.toggle_state(self.current_view == RatePredictionView::SuggestedEdits),
)
.child(
Button::new(
ElementId::Name("raw-input".into()),
RateCompletionView::RawInput.name(),
RatePredictionView::RawInput.name(),
)
.label_size(LabelSize::Small)
.on_click(cx.listener(move |this, _, _window, cx| {
this.current_view = RateCompletionView::RawInput;
this.current_view = RatePredictionView::RawInput;
cx.notify();
}))
.toggle_state(self.current_view == RateCompletionView::RawInput),
.toggle_state(self.current_view == RatePredictionView::RawInput),
)
}
fn render_suggested_edits(&self, cx: &mut Context<Self>) -> Option<gpui::Stateful<Div>> {
let active_completion = self.active_completion.as_ref()?;
let bg_color = cx.theme().colors().editor_background;
Some(
div()
.id("diff")
@@ -347,14 +466,18 @@ impl RateCompletionModal {
.bg(bg_color)
.overflow_scroll()
.whitespace_nowrap()
.child(CompletionDiffElement::new(
&active_completion.completion,
cx,
)),
.child(self.diff_editor.clone()),
)
}
fn render_raw_input(&self, cx: &mut Context<Self>) -> Option<gpui::Stateful<Div>> {
fn render_raw_input(
&self,
window: &mut Window,
cx: &mut Context<Self>,
) -> Option<gpui::Stateful<Div>> {
let theme_settings = ThemeSettings::get_global(cx);
let buffer_font_size = theme_settings.buffer_font_size(cx);
Some(
v_flex()
.size_full()
@@ -368,30 +491,81 @@ impl RateCompletionModal {
.size_full()
.bg(cx.theme().colors().editor_background)
.overflow_scroll()
.child(if let Some(active_completion) = &self.active_completion {
format!(
"{}\n{}",
active_completion.completion.input_events,
active_completion.completion.input_excerpt
.child(if let Some(active_prediction) = &self.active_prediction {
markdown::MarkdownElement::new(
active_prediction.formatted_inputs.clone(),
MarkdownStyle {
base_text_style: window.text_style(),
syntax: cx.theme().syntax().clone(),
code_block: StyleRefinement {
text: Some(TextStyleRefinement {
font_family: Some(
theme_settings.buffer_font.family.clone(),
),
font_size: Some(buffer_font_size.into()),
..Default::default()
}),
padding: EdgesRefinement {
top: Some(DefiniteLength::Absolute(
AbsoluteLength::Pixels(px(8.)),
)),
left: Some(DefiniteLength::Absolute(
AbsoluteLength::Pixels(px(8.)),
)),
right: Some(DefiniteLength::Absolute(
AbsoluteLength::Pixels(px(8.)),
)),
bottom: Some(DefiniteLength::Absolute(
AbsoluteLength::Pixels(px(8.)),
)),
},
margin: EdgesRefinement {
top: Some(Length::Definite(px(8.).into())),
left: Some(Length::Definite(px(0.).into())),
right: Some(Length::Definite(px(0.).into())),
bottom: Some(Length::Definite(px(12.).into())),
},
border_style: Some(BorderStyle::Solid),
border_widths: EdgesRefinement {
top: Some(AbsoluteLength::Pixels(px(1.))),
left: Some(AbsoluteLength::Pixels(px(1.))),
right: Some(AbsoluteLength::Pixels(px(1.))),
bottom: Some(AbsoluteLength::Pixels(px(1.))),
},
border_color: Some(cx.theme().colors().border_variant),
background: Some(
cx.theme().colors().editor_background.into(),
),
..Default::default()
},
..Default::default()
},
)
.into_any_element()
} else {
"No active completion".to_string()
div()
.child("No active completion".to_string())
.into_any_element()
}),
)
.id("raw-input-view"),
)
}
fn render_active_completion(&mut self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
let active_completion = self.active_completion.as_ref()?;
let completion_id = active_completion.completion.id;
fn render_active_completion(
&mut self,
window: &mut Window,
cx: &mut Context<Self>,
) -> Option<impl IntoElement> {
let active_prediction = self.active_prediction.as_ref()?;
let completion_id = active_prediction.prediction.id.clone();
let focus_handle = &self.focus_handle(cx);
let border_color = cx.theme().colors().border;
let bg_color = cx.theme().colors().editor_background;
let rated = self.zeta.read(cx).is_completion_rated(completion_id);
let feedback_empty = active_completion
let rated = self.zeta.read(cx).is_prediction_rated(&completion_id);
let feedback_empty = active_prediction
.feedback_editor
.read(cx)
.text(cx)
@@ -412,10 +586,10 @@ impl RateCompletionModal {
.child(self.render_view_nav(cx))
.when_some(
match self.current_view {
RateCompletionView::SuggestedEdits => {
RatePredictionView::SuggestedEdits => {
self.render_suggested_edits(cx)
}
RateCompletionView::RawInput => self.render_raw_input(cx),
RatePredictionView::RawInput => self.render_raw_input(window, cx),
},
|this, element| this.child(element),
),
@@ -450,7 +624,7 @@ impl RateCompletionModal {
.h_40()
.pt_1()
.bg(bg_color)
.child(active_completion.feedback_editor.clone()),
.child(active_prediction.feedback_editor.clone()),
)
})
.child(
@@ -472,7 +646,7 @@ impl RateCompletionModal {
)
.child(Label::new("Rated completion.").color(Color::Muted)),
)
} else if active_completion.completion.edits.is_empty() {
} else if active_prediction.prediction.edits.is_empty() {
Some(
label_container
.child(
@@ -489,7 +663,7 @@ impl RateCompletionModal {
h_flex()
.gap_1()
.child(
Button::new("bad", "Bad Completion")
Button::new("bad", "Bad Prediction")
.icon(IconName::ThumbsDown)
.icon_size(IconSize::Small)
.icon_position(IconPosition::Start)
@@ -500,14 +674,14 @@ impl RateCompletionModal {
))
})
.key_binding(KeyBinding::for_action_in(
&ThumbsDownActiveCompletion,
&ThumbsDownActivePrediction,
focus_handle,
cx,
))
.on_click(cx.listener(move |this, _, window, cx| {
if this.active_completion.is_some() {
if this.active_prediction.is_some() {
this.thumbs_down_active(
&ThumbsDownActiveCompletion,
&ThumbsDownActivePrediction,
window,
cx,
);
@@ -515,20 +689,20 @@ impl RateCompletionModal {
})),
)
.child(
Button::new("good", "Good Completion")
Button::new("good", "Good Prediction")
.icon(IconName::ThumbsUp)
.icon_size(IconSize::Small)
.icon_position(IconPosition::Start)
.disabled(rated)
.key_binding(KeyBinding::for_action_in(
&ThumbsUpActiveCompletion,
&ThumbsUpActivePrediction,
focus_handle,
cx,
))
.on_click(cx.listener(move |this, _, window, cx| {
if this.active_completion.is_some() {
if this.active_prediction.is_some() {
this.thumbs_up_active(
&ThumbsUpActiveCompletion,
&ThumbsUpActivePrediction,
window,
cx,
);
@@ -543,34 +717,32 @@ impl RateCompletionModal {
fn render_shown_completions(&self, cx: &Context<Self>) -> impl Iterator<Item = ListItem> {
self.zeta
.read(cx)
.shown_completions()
.shown_predictions()
.cloned()
.enumerate()
.map(|(index, completion)| {
let selected = self
.active_completion
.active_prediction
.as_ref()
.is_some_and(|selected| selected.completion.id == completion.id);
let rated = self.zeta.read(cx).is_completion_rated(completion.id);
.is_some_and(|selected| selected.prediction.id == completion.id);
let rated = self.zeta.read(cx).is_prediction_rated(&completion.id);
let (icon_name, icon_color, tooltip_text) =
match (rated, completion.edits.is_empty()) {
(true, _) => (IconName::Check, Color::Success, "Rated Completion"),
(true, _) => (IconName::Check, Color::Success, "Rated Prediction"),
(false, true) => (IconName::File, Color::Muted, "No Edits Produced"),
(false, false) => (IconName::FileDiff, Color::Accent, "Edits Available"),
};
let file_name = completion
.path
.file_name()
.map(|f| f.to_string_lossy().into_owned())
.unwrap_or("untitled".to_string());
let file_path = completion
.path
.parent()
.map(|p| p.to_string_lossy().into_owned());
let file = completion.buffer.read(cx).file();
let file_name = file
.as_ref()
.map_or(SharedString::new_static("untitled"), |file| {
file.file_name(cx).to_string().into()
});
let file_path = file.map(|file| file.path().as_unix_str().to_string());
ListItem::new(completion.id)
ListItem::new(completion.id.clone())
.inset(true)
.spacing(ListItemSpacing::Sparse)
.focused(index == self.selected_index)
@@ -615,12 +787,12 @@ impl RateCompletionModal {
}
}
impl Render for RateCompletionModal {
impl Render for RatePredictionsModal {
fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
let border_color = cx.theme().colors().border;
h_flex()
.key_context("RateCompletionModal")
.key_context("RatePredictionModal")
.track_focus(&self.focus_handle)
.on_action(cx.listener(Self::dismiss))
.on_action(cx.listener(Self::confirm))
@@ -688,20 +860,20 @@ impl Render for RateCompletionModal {
),
),
)
.children(self.render_active_completion(cx))
.children(self.render_active_completion(window, cx))
.on_mouse_down_out(cx.listener(|_, _, _, cx| cx.emit(DismissEvent)))
}
}
impl EventEmitter<DismissEvent> for RateCompletionModal {}
impl EventEmitter<DismissEvent> for RatePredictionsModal {}
impl Focusable for RateCompletionModal {
impl Focusable for RatePredictionsModal {
fn focus_handle(&self, _cx: &App) -> FocusHandle {
self.focus_handle.clone()
}
}
impl ModalView for RateCompletionModal {}
impl ModalView for RatePredictionsModal {}
fn format_time_ago(elapsed: Duration) -> String {
let seconds = elapsed.as_secs();
crates/zeta/src/retrieval_search.rs
+643
View File
@@ -0,0 +1,643 @@
use anyhow::Result;
use cloud_zeta2_prompt::retrieval_prompt::SearchToolQuery;
use collections::HashMap;
use futures::{
StreamExt,
channel::mpsc::{self, UnboundedSender},
};
use gpui::{AppContext, AsyncApp, Entity};
use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt, Point, ToOffset, ToPoint};
use project::{
Project, WorktreeSettings,
search::{SearchQuery, SearchResult},
};
use smol::channel;
use std::ops::Range;
use util::{
ResultExt as _,
paths::{PathMatcher, PathStyle},
};
use workspace::item::Settings as _;
#[cfg(feature = "eval-support")]
type CachedSearchResults = std::collections::BTreeMap<std::path::PathBuf, Vec<Range<usize>>>;
pub async fn run_retrieval_searches(
queries: Vec<SearchToolQuery>,
project: Entity<Project>,
#[cfg(feature = "eval-support")] eval_cache: Option<std::sync::Arc<dyn crate::EvalCache>>,
cx: &mut AsyncApp,
) -> Result<HashMap<Entity<Buffer>, Vec<Range<Anchor>>>> {
#[cfg(feature = "eval-support")]
let cache = if let Some(eval_cache) = eval_cache {
use crate::EvalCacheEntryKind;
use anyhow::Context;
use collections::FxHasher;
use std::hash::{Hash, Hasher};
let mut hasher = FxHasher::default();
project.read_with(cx, |project, cx| {
let mut worktrees = project.worktrees(cx);
let Some(worktree) = worktrees.next() else {
panic!("Expected a single worktree in eval project. Found none.");
};
assert!(
worktrees.next().is_none(),
"Expected a single worktree in eval project. Found more than one."
);
worktree.read(cx).abs_path().hash(&mut hasher);
})?;
queries.hash(&mut hasher);
let key = (EvalCacheEntryKind::Search, hasher.finish());
if let Some(cached_results) = eval_cache.read(key) {
let file_results = serde_json::from_str::<CachedSearchResults>(&cached_results)
.context("Failed to deserialize cached search results")?;
let mut results = HashMap::default();
for (path, ranges) in file_results {
let buffer = project
.update(cx, |project, cx| {
let project_path = project.find_project_path(path, cx).unwrap();
project.open_buffer(project_path, cx)
})?
.await?;
let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
let mut ranges: Vec<_> = ranges
.into_iter()
.map(|range| {
snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end)
})
.collect();
merge_anchor_ranges(&mut ranges, &snapshot);
results.insert(buffer, ranges);
}
return Ok(results);
}
Some((eval_cache, serde_json::to_string_pretty(&queries)?, key))
} else {
None
};
let (exclude_matcher, path_style) = project.update(cx, |project, cx| {
let global_settings = WorktreeSettings::get_global(cx);
let exclude_patterns = global_settings
.file_scan_exclusions
.sources()
.iter()
.chain(global_settings.private_files.sources().iter());
let path_style = project.path_style(cx);
anyhow::Ok((PathMatcher::new(exclude_patterns, path_style)?, path_style))
})??;
let (results_tx, mut results_rx) = mpsc::unbounded();
for query in queries {
let exclude_matcher = exclude_matcher.clone();
let results_tx = results_tx.clone();
let project = project.clone();
cx.spawn(async move |cx| {
run_query(
query,
results_tx.clone(),
path_style,
exclude_matcher,
&project,
cx,
)
.await
.log_err();
})
.detach()
}
drop(results_tx);
#[cfg(feature = "eval-support")]
let cache = cache.clone();
cx.background_spawn(async move {
let mut results: HashMap<Entity<Buffer>, Vec<Range<Anchor>>> = HashMap::default();
let mut snapshots = HashMap::default();
let mut total_bytes = 0;
'outer: while let Some((buffer, snapshot, excerpts)) = results_rx.next().await {
snapshots.insert(buffer.entity_id(), snapshot);
let existing = results.entry(buffer).or_default();
existing.reserve(excerpts.len());
for (range, size) in excerpts {
// Blunt trimming of the results until we have a proper algorithmic filtering step
if (total_bytes + size) > MAX_RESULTS_LEN {
log::trace!("Combined results reached limit of {MAX_RESULTS_LEN}B");
break 'outer;
}
total_bytes += size;
existing.push(range);
}
}
#[cfg(feature = "eval-support")]
if let Some((cache, queries, key)) = cache {
let cached_results: CachedSearchResults = results
.iter()
.filter_map(|(buffer, ranges)| {
let snapshot = snapshots.get(&buffer.entity_id())?;
let path = snapshot.file().map(|f| f.path());
let mut ranges = ranges
.iter()
.map(|range| range.to_offset(&snapshot))
.collect::<Vec<_>>();
ranges.sort_unstable_by_key(|range| (range.start, range.end));
Some((path?.as_std_path().to_path_buf(), ranges))
})
.collect();
cache.write(
key,
&queries,
&serde_json::to_string_pretty(&cached_results)?,
);
}
for (buffer, ranges) in results.iter_mut() {
if let Some(snapshot) = snapshots.get(&buffer.entity_id()) {
merge_anchor_ranges(ranges, snapshot);
}
}
Ok(results)
})
.await
}
pub(crate) fn merge_anchor_ranges(ranges: &mut Vec<Range<Anchor>>, snapshot: &BufferSnapshot) {
ranges.sort_unstable_by(|a, b| {
a.start
.cmp(&b.start, snapshot)
.then(b.end.cmp(&a.end, snapshot))
});
let mut index = 1;
while index < ranges.len() {
if ranges[index - 1]
.end
.cmp(&ranges[index].start, snapshot)
.is_ge()
{
let removed = ranges.remove(index);
if removed.end.cmp(&ranges[index - 1].end, snapshot).is_gt() {
ranges[index - 1].end = removed.end;
}
} else {
index += 1;
}
}
}
const MAX_EXCERPT_LEN: usize = 768;
const MAX_RESULTS_LEN: usize = MAX_EXCERPT_LEN * 5;
struct SearchJob {
buffer: Entity<Buffer>,
snapshot: BufferSnapshot,
ranges: Vec<Range<usize>>,
query_ix: usize,
jobs_tx: channel::Sender<SearchJob>,
}
async fn run_query(
input_query: SearchToolQuery,
results_tx: UnboundedSender<(Entity<Buffer>, BufferSnapshot, Vec<(Range<Anchor>, usize)>)>,
path_style: PathStyle,
exclude_matcher: PathMatcher,
project: &Entity<Project>,
cx: &mut AsyncApp,
) -> Result<()> {
let include_matcher = PathMatcher::new(vec![input_query.glob], path_style)?;
let make_search = |regex: &str| -> Result<SearchQuery> {
SearchQuery::regex(
regex,
false,
true,
false,
true,
include_matcher.clone(),
exclude_matcher.clone(),
true,
None,
)
};
if let Some(outer_syntax_regex) = input_query.syntax_node.first() {
let outer_syntax_query = make_search(outer_syntax_regex)?;
let nested_syntax_queries = input_query
.syntax_node
.into_iter()
.skip(1)
.map(|query| make_search(&query))
.collect::<Result<Vec<_>>>()?;
let content_query = input_query
.content
.map(|regex| make_search(&regex))
.transpose()?;
let (jobs_tx, jobs_rx) = channel::unbounded();
let outer_search_results_rx =
project.update(cx, |project, cx| project.search(outer_syntax_query, cx))?;
let outer_search_task = cx.spawn(async move |cx| {
futures::pin_mut!(outer_search_results_rx);
while let Some(SearchResult::Buffer { buffer, ranges }) =
outer_search_results_rx.next().await
{
buffer
.read_with(cx, |buffer, _| buffer.parsing_idle())?
.await;
let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
let expanded_ranges: Vec<_> = ranges
.into_iter()
.filter_map(|range| expand_to_parent_range(&range, &snapshot))
.collect();
jobs_tx
.send(SearchJob {
buffer,
snapshot,
ranges: expanded_ranges,
query_ix: 0,
jobs_tx: jobs_tx.clone(),
})
.await?;
}
anyhow::Ok(())
});
let n_workers = cx.background_executor().num_cpus();
let search_job_task = cx.background_executor().scoped(|scope| {
for _ in 0..n_workers {
scope.spawn(async {
while let Ok(job) = jobs_rx.recv().await {
process_nested_search_job(
&results_tx,
&nested_syntax_queries,
&content_query,
job,
)
.await;
}
});
}
});
search_job_task.await;
outer_search_task.await?;
} else if let Some(content_regex) = &input_query.content {
let search_query = make_search(&content_regex)?;
let results_rx = project.update(cx, |project, cx| project.search(search_query, cx))?;
futures::pin_mut!(results_rx);
while let Some(SearchResult::Buffer { buffer, ranges }) = results_rx.next().await {
let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
let ranges = ranges
.into_iter()
.map(|range| {
let range = range.to_offset(&snapshot);
let range = expand_to_entire_lines(range, &snapshot);
let size = range.len();
let range =
snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end);
(range, size)
})
.collect();
let send_result = results_tx.unbounded_send((buffer.clone(), snapshot.clone(), ranges));
if let Err(err) = send_result
&& !err.is_disconnected()
{
log::error!("{err}");
}
}
} else {
log::warn!("Context gathering model produced a glob-only search");
}
anyhow::Ok(())
}
async fn process_nested_search_job(
results_tx: &UnboundedSender<(Entity<Buffer>, BufferSnapshot, Vec<(Range<Anchor>, usize)>)>,
queries: &Vec<SearchQuery>,
content_query: &Option<SearchQuery>,
job: SearchJob,
) {
if let Some(search_query) = queries.get(job.query_ix) {
let mut subranges = Vec::new();
for range in job.ranges {
let start = range.start;
let search_results = search_query.search(&job.snapshot, Some(range)).await;
for subrange in search_results {
let subrange = start + subrange.start..start + subrange.end;
subranges.extend(expand_to_parent_range(&subrange, &job.snapshot));
}
}
job.jobs_tx
.send(SearchJob {
buffer: job.buffer,
snapshot: job.snapshot,
ranges: subranges,
query_ix: job.query_ix + 1,
jobs_tx: job.jobs_tx.clone(),
})
.await
.ok();
} else {
let ranges = if let Some(content_query) = content_query {
let mut subranges = Vec::new();
for range in job.ranges {
let start = range.start;
let search_results = content_query.search(&job.snapshot, Some(range)).await;
for subrange in search_results {
let subrange = start + subrange.start..start + subrange.end;
subranges.push(subrange);
}
}
subranges
} else {
job.ranges
};
let matches = ranges
.into_iter()
.map(|range| {
let snapshot = &job.snapshot;
let range = expand_to_entire_lines(range, snapshot);
let size = range.len();
let range = snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end);
(range, size)
})
.collect();
let send_result = results_tx.unbounded_send((job.buffer, job.snapshot, matches));
if let Err(err) = send_result
&& !err.is_disconnected()
{
log::error!("{err}");
}
}
}
fn expand_to_entire_lines(range: Range<usize>, snapshot: &BufferSnapshot) -> Range<usize> {
let mut point_range = range.to_point(snapshot);
point_range.start.column = 0;
if point_range.end.column > 0 {
point_range.end = snapshot.max_point().min(point_range.end + Point::new(1, 0));
}
point_range.to_offset(snapshot)
}
fn expand_to_parent_range<T: ToPoint + ToOffset>(
range: &Range<T>,
snapshot: &BufferSnapshot,
) -> Option<Range<usize>> {
let mut line_range = range.to_point(&snapshot);
line_range.start.column = snapshot.indent_size_for_line(line_range.start.row).len;
line_range.end.column = snapshot.line_len(line_range.end.row);
// TODO skip result if matched line isn't the first node line?
let node = snapshot.syntax_ancestor(line_range)?;
Some(node.byte_range())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::assemble_excerpts::assemble_excerpts;
use cloud_zeta2_prompt::write_codeblock;
use edit_prediction_context::Line;
use gpui::TestAppContext;
use indoc::indoc;
use language::{Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
use pretty_assertions::assert_eq;
use project::FakeFs;
use serde_json::json;
use settings::SettingsStore;
use std::path::Path;
use util::path;
#[gpui::test]
async fn test_retrieval(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
path!("/root"),
json!({
"user.rs": indoc!{"
pub struct Organization {
owner: Arc<User>,
}
pub struct User {
first_name: String,
last_name: String,
}
impl Organization {
pub fn owner(&self) -> Arc<User> {
self.owner.clone()
}
}
impl User {
pub fn new(first_name: String, last_name: String) -> Self {
Self {
first_name,
last_name
}
}
pub fn first_name(&self) -> String {
self.first_name.clone()
}
pub fn last_name(&self) -> String {
self.last_name.clone()
}
}
"},
"main.rs": indoc!{r#"
fn main() {
let user = User::new(FIRST_NAME.clone(), "doe".into());
println!("user {:?}", user);
}
"#},
}),
)
.await;
let project = Project::test(fs, vec![Path::new(path!("/root"))], cx).await;
project.update(cx, |project, _cx| {
project.languages().add(rust_lang().into())
});
assert_results(
&project,
SearchToolQuery {
glob: "user.rs".into(),
syntax_node: vec!["impl\\s+User".into(), "pub\\s+fn\\s+first_name".into()],
content: None,
},
indoc! {r#"
`````root/user.rs
…
impl User {
…
pub fn first_name(&self) -> String {
self.first_name.clone()
}
…
`````
"#},
cx,
)
.await;
assert_results(
&project,
SearchToolQuery {
glob: "user.rs".into(),
syntax_node: vec!["impl\\s+User".into()],
content: Some("\\.clone".into()),
},
indoc! {r#"
`````root/user.rs
…
impl User {
…
pub fn first_name(&self) -> String {
self.first_name.clone()
…
pub fn last_name(&self) -> String {
self.last_name.clone()
…
`````
"#},
cx,
)
.await;
assert_results(
&project,
SearchToolQuery {
glob: "*.rs".into(),
syntax_node: vec![],
content: Some("\\.clone".into()),
},
indoc! {r#"
`````root/main.rs
fn main() {
let user = User::new(FIRST_NAME.clone(), "doe".into());
…
`````
`````root/user.rs
…
impl Organization {
pub fn owner(&self) -> Arc<User> {
self.owner.clone()
…
impl User {
…
pub fn first_name(&self) -> String {
self.first_name.clone()
…
pub fn last_name(&self) -> String {
self.last_name.clone()
…
`````
"#},
cx,
)
.await;
}
async fn assert_results(
project: &Entity<Project>,
query: SearchToolQuery,
expected_output: &str,
cx: &mut TestAppContext,
) {
let results = run_retrieval_searches(
vec![query],
project.clone(),
#[cfg(feature = "eval-support")]
None,
&mut cx.to_async(),
)
.await
.unwrap();
let mut results = results.into_iter().collect::<Vec<_>>();
results.sort_by_key(|results| {
results
.0
.read_with(cx, |buffer, _| buffer.file().unwrap().path().clone())
});
let mut output = String::new();
for (buffer, ranges) in results {
buffer.read_with(cx, |buffer, cx| {
let excerpts = ranges.into_iter().map(|range| {
let point_range = range.to_point(buffer);
if point_range.end.column > 0 {
Line(point_range.start.row)..Line(point_range.end.row + 1)
} else {
Line(point_range.start.row)..Line(point_range.end.row)
}
});
write_codeblock(
&buffer.file().unwrap().full_path(cx),
assemble_excerpts(&buffer.snapshot(), excerpts).iter(),
&[],
Line(buffer.max_point().row),
false,
&mut output,
);
});
}
output.pop();
assert_eq!(output, expected_output);
}
fn rust_lang() -> Language {
Language::new(
LanguageConfig {
name: "Rust".into(),
matcher: LanguageMatcher {
path_suffixes: vec!["rs".to_string()],
..Default::default()
},
..Default::default()
},
Some(tree_sitter_rust::LANGUAGE.into()),
)
.with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
.unwrap()
}
fn init_test(cx: &mut TestAppContext) {
cx.update(move |cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
zlog::init_test();
});
}
}
crates/zeta/src/sweep_ai.rs
+127
View File
@@ -0,0 +1,127 @@
use std::fmt;
use std::{path::Path, sync::Arc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize)]
pub struct AutocompleteRequest {
pub debug_info: Arc<str>,
pub repo_name: String,
pub branch: Option<String>,
pub file_path: Arc<Path>,
pub file_contents: String,
pub recent_changes: String,
pub cursor_position: usize,
pub original_file_contents: String,
pub file_chunks: Vec<FileChunk>,
pub retrieval_chunks: Vec<RetrievalChunk>,
pub recent_user_actions: Vec<UserAction>,
pub multiple_suggestions: bool,
pub privacy_mode_enabled: bool,
pub changes_above_cursor: bool,
}
#[derive(Debug, Clone, Serialize)]
pub struct FileChunk {
pub file_path: String,
pub start_line: usize,
pub end_line: usize,
pub content: String,
pub timestamp: Option<u64>,
}
#[derive(Debug, Clone, Serialize)]
pub struct RetrievalChunk {
pub file_path: String,
pub start_line: usize,
pub end_line: usize,
pub content: String,
pub timestamp: u64,
}
#[derive(Debug, Clone, Serialize)]
pub struct UserAction {
pub action_type: ActionType,
pub line_number: usize,
pub offset: usize,
pub file_path: String,
pub timestamp: u64,
}
#[allow(dead_code)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum ActionType {
CursorMovement,
InsertChar,
DeleteChar,
InsertSelection,
DeleteSelection,
}
#[derive(Debug, Clone, Deserialize)]
pub struct AutocompleteResponse {
pub autocomplete_id: String,
pub start_index: usize,
pub end_index: usize,
pub completion: String,
#[allow(dead_code)]
pub confidence: f64,
#[allow(dead_code)]
pub logprobs: Option<serde_json::Value>,
#[allow(dead_code)]
pub finish_reason: Option<String>,
#[allow(dead_code)]
pub elapsed_time_ms: u64,
#[allow(dead_code)]
#[serde(default, rename = "completions")]
pub additional_completions: Vec<AdditionalCompletion>,
}
#[allow(dead_code)]
#[derive(Debug, Clone, Deserialize)]
pub struct AdditionalCompletion {
pub start_index: usize,
pub end_index: usize,
pub completion: String,
pub confidence: f64,
pub autocomplete_id: String,
pub logprobs: Option<serde_json::Value>,
pub finish_reason: Option<String>,
}
pub(crate) fn write_event(
event: &cloud_llm_client::predict_edits_v3::Event,
f: &mut impl fmt::Write,
) -> fmt::Result {
match event {
cloud_llm_client::predict_edits_v3::Event::BufferChange {
old_path,
path,
diff,
..
} => {
if old_path != path {
// TODO confirm how to do this for sweep
// writeln!(f, "User renamed {:?} to {:?}\n", old_path, new_path)?;
}
if !diff.is_empty() {
write!(f, "File: {}:\n{}\n", path.display(), diff)?
}
fmt::Result::Ok(())
}
}
}
pub(crate) fn debug_info(cx: &gpui::App) -> Arc<str> {
format!(
"Zed v{version} ({sha}) - OS: {os} - Zed v{version}",
version = release_channel::AppVersion::global(cx),
sha = release_channel::AppCommitSha::try_global(cx)
.map_or("unknown".to_string(), |sha| sha.full()),
os = client::telemetry::os_name(),
)
.into()
}
crates/zeta/src/udiff.rs
+1024
View File
File diff suppressed because it is too large Load Diff
crates/zeta/src/xml_edits.rs
+637
View File
@@ -0,0 +1,637 @@
use anyhow::{Context as _, Result};
use language::{Anchor, BufferSnapshot, OffsetRangeExt as _, Point};
use std::{cmp, ops::Range, path::Path, sync::Arc};
const EDITS_TAG_NAME: &'static str = "edits";
const OLD_TEXT_TAG_NAME: &'static str = "old_text";
const NEW_TEXT_TAG_NAME: &'static str = "new_text";
const XML_TAGS: &[&str] = &[EDITS_TAG_NAME, OLD_TEXT_TAG_NAME, NEW_TEXT_TAG_NAME];
pub async fn parse_xml_edits<'a>(
input: &'a str,
get_buffer: impl Fn(&Path) -> Option<(&'a BufferSnapshot, &'a [Range<Anchor>])> + Send,
) -> Result<(&'a BufferSnapshot, Vec<(Range<Anchor>, Arc<str>)>)> {
parse_xml_edits_inner(input, get_buffer)
.await
.with_context(|| format!("Failed to parse XML edits:\n{input}"))
}
async fn parse_xml_edits_inner<'a>(
input: &'a str,
get_buffer: impl Fn(&Path) -> Option<(&'a BufferSnapshot, &'a [Range<Anchor>])> + Send,
) -> Result<(&'a BufferSnapshot, Vec<(Range<Anchor>, Arc<str>)>)> {
let xml_edits = extract_xml_replacements(input)?;
let (buffer, context_ranges) = get_buffer(xml_edits.file_path.as_ref())
.with_context(|| format!("no buffer for file {}", xml_edits.file_path))?;
let mut all_edits = vec![];
for (old_text, new_text) in xml_edits.replacements {
let match_range = fuzzy_match_in_ranges(old_text, buffer, context_ranges)?;
let matched_old_text = buffer
.text_for_range(match_range.clone())
.collect::<String>();
let edits_within_hunk = language::text_diff(&matched_old_text, new_text);
all_edits.extend(
edits_within_hunk
.into_iter()
.map(move |(inner_range, inner_text)| {
(
buffer.anchor_after(match_range.start + inner_range.start)
..buffer.anchor_before(match_range.start + inner_range.end),
inner_text,
)
}),
);
}
Ok((buffer, all_edits))
}
fn fuzzy_match_in_ranges(
old_text: &str,
buffer: &BufferSnapshot,
context_ranges: &[Range<Anchor>],
) -> Result<Range<usize>> {
let mut state = FuzzyMatcher::new(buffer, old_text);
let mut best_match = None;
let mut tie_match_range = None;
for range in context_ranges {
let best_match_cost = best_match.as_ref().map(|(score, _)| *score);
match (best_match_cost, state.match_range(range.to_offset(buffer))) {
(Some(lowest_cost), Some((new_cost, new_range))) => {
if new_cost == lowest_cost {
tie_match_range = Some(new_range);
} else if new_cost < lowest_cost {
tie_match_range.take();
best_match = Some((new_cost, new_range));
}
}
(None, Some(new_match)) => {
best_match = Some(new_match);
}
(None, None) | (Some(_), None) => {}
};
}
if let Some((_, best_match_range)) = best_match {
if let Some(tie_match_range) = tie_match_range {
anyhow::bail!(
"Multiple ambiguous matches:\n{:?}:\n{}\n\n{:?}:\n{}",
best_match_range.clone(),
buffer.text_for_range(best_match_range).collect::<String>(),
tie_match_range.clone(),
buffer.text_for_range(tie_match_range).collect::<String>()
);
}
return Ok(best_match_range);
}
anyhow::bail!(
"Failed to fuzzy match `old_text`:\n{}\nin:\n```\n{}\n```",
old_text,
context_ranges
.iter()
.map(|range| buffer.text_for_range(range.clone()).collect::<String>())
.collect::<Vec<String>>()
.join("```\n```")
);
}
#[derive(Debug)]
struct XmlEdits<'a> {
file_path: &'a str,
/// Vec of (old_text, new_text) pairs
replacements: Vec<(&'a str, &'a str)>,
}
fn extract_xml_replacements(input: &str) -> Result<XmlEdits<'_>> {
let mut cursor = 0;
let (edits_body_start, edits_attrs) =
find_tag_open(input, &mut cursor, EDITS_TAG_NAME)?.context("No edits tag found")?;
let file_path = edits_attrs
.trim_start()
.strip_prefix("path")
.context("no path attribute on edits tag")?
.trim_end()
.strip_prefix('=')
.context("no value for path attribute")?
.trim()
.trim_start_matches('"')
.trim_end_matches('"');
cursor = edits_body_start;
let mut edits_list = Vec::new();
while let Some((old_body_start, _)) = find_tag_open(input, &mut cursor, OLD_TEXT_TAG_NAME)? {
let old_body_end = find_tag_close(input, &mut cursor)?;
let old_text = trim_surrounding_newlines(&input[old_body_start..old_body_end]);
let (new_body_start, _) = find_tag_open(input, &mut cursor, NEW_TEXT_TAG_NAME)?
.context("no new_text tag following old_text")?;
let new_body_end = find_tag_close(input, &mut cursor)?;
let new_text = trim_surrounding_newlines(&input[new_body_start..new_body_end]);
edits_list.push((old_text, new_text));
}
Ok(XmlEdits {
file_path,
replacements: edits_list,
})
}
/// Trims a single leading and trailing newline
fn trim_surrounding_newlines(input: &str) -> &str {
let start = input.strip_prefix('\n').unwrap_or(input);
let end = start.strip_suffix('\n').unwrap_or(start);
end
}
fn find_tag_open<'a>(
input: &'a str,
cursor: &mut usize,
expected_tag: &str,
) -> Result<Option<(usize, &'a str)>> {
let mut search_pos = *cursor;
while search_pos < input.len() {
let Some(tag_start) = input[search_pos..].find("<") else {
break;
};
let tag_start = search_pos + tag_start;
if !input[tag_start + 1..].starts_with(expected_tag) {
search_pos = search_pos + tag_start + 1;
continue;
};
let after_tag_name = tag_start + expected_tag.len() + 1;
let close_bracket = input[after_tag_name..]
.find('>')
.with_context(|| format!("missing > after <{}", expected_tag))?;
let attrs_end = after_tag_name + close_bracket;
let body_start = attrs_end + 1;
let attributes = input[after_tag_name..attrs_end].trim();
*cursor = body_start;
return Ok(Some((body_start, attributes)));
}
Ok(None)
}
fn find_tag_close(input: &str, cursor: &mut usize) -> Result<usize> {
let mut depth = 1;
let mut search_pos = *cursor;
while search_pos < input.len() && depth > 0 {
let Some(bracket_offset) = input[search_pos..].find('<') else {
break;
};
let bracket_pos = search_pos + bracket_offset;
if input[bracket_pos..].starts_with("</")
&& let Some(close_end) = input[bracket_pos + 2..].find('>')
{
let close_start = bracket_pos + 2;
let tag_name = input[close_start..close_start + close_end].trim();
if XML_TAGS.contains(&tag_name) {
depth -= 1;
if depth == 0 {
*cursor = close_start + close_end + 1;
return Ok(bracket_pos);
}
}
search_pos = close_start + close_end + 1;
continue;
} else if let Some(close_bracket_offset) = input[bracket_pos..].find('>') {
let close_bracket_pos = bracket_pos + close_bracket_offset;
let tag_name = &input[bracket_pos + 1..close_bracket_pos].trim();
if XML_TAGS.contains(&tag_name) {
depth += 1;
}
}
search_pos = bracket_pos + 1;
}
anyhow::bail!("no closing tag found")
}
const REPLACEMENT_COST: u32 = 1;
const INSERTION_COST: u32 = 3;
const DELETION_COST: u32 = 10;
/// A fuzzy matcher that can process text chunks incrementally
/// and return the best match found so far at each step.
struct FuzzyMatcher<'a> {
snapshot: &'a BufferSnapshot,
query_lines: Vec<&'a str>,
matrix: SearchMatrix,
}
impl<'a> FuzzyMatcher<'a> {
fn new(snapshot: &'a BufferSnapshot, old_text: &'a str) -> Self {
let query_lines = old_text.lines().collect();
Self {
snapshot,
query_lines,
matrix: SearchMatrix::new(0),
}
}
fn match_range(&mut self, range: Range<usize>) -> Option<(u32, Range<usize>)> {
let point_range = range.to_point(&self.snapshot);
let buffer_line_count = (point_range.end.row - point_range.start.row + 1) as usize;
self.matrix
.reset(self.query_lines.len() + 1, buffer_line_count + 1);
let query_line_count = self.query_lines.len();
for row in 0..query_line_count {
let query_line = self.query_lines[row].trim();
let leading_deletion_cost = (row + 1) as u32 * DELETION_COST;
self.matrix.set(
row + 1,
0,
SearchState::new(leading_deletion_cost, SearchDirection::Up),
);
let mut buffer_lines = self.snapshot.text_for_range(range.clone()).lines();
let mut col = 0;
while let Some(buffer_line) = buffer_lines.next() {
let buffer_line = buffer_line.trim();
let up = SearchState::new(
self.matrix
.get(row, col + 1)
.cost
.saturating_add(DELETION_COST),
SearchDirection::Up,
);
let left = SearchState::new(
self.matrix
.get(row + 1, col)
.cost
.saturating_add(INSERTION_COST),
SearchDirection::Left,
);
let diagonal = SearchState::new(
if query_line == buffer_line {
self.matrix.get(row, col).cost
} else if fuzzy_eq(query_line, buffer_line) {
self.matrix.get(row, col).cost + REPLACEMENT_COST
} else {
self.matrix
.get(row, col)
.cost
.saturating_add(DELETION_COST + INSERTION_COST)
},
SearchDirection::Diagonal,
);
self.matrix
.set(row + 1, col + 1, up.min(left).min(diagonal));
col += 1;
}
}
// Find all matches with the best cost
let mut best_cost = u32::MAX;
let mut matches_with_best_cost = Vec::new();
for col in 1..=buffer_line_count {
let cost = self.matrix.get(query_line_count, col).cost;
if cost < best_cost {
best_cost = cost;
matches_with_best_cost.clear();
matches_with_best_cost.push(col as u32);
} else if cost == best_cost {
matches_with_best_cost.push(col as u32);
}
}
// Find ranges for the matches
for &match_end_col in &matches_with_best_cost {
let mut matched_lines = 0;
let mut query_row = query_line_count;
let mut match_start_col = match_end_col;
while query_row > 0 && match_start_col > 0 {
let current = self.matrix.get(query_row, match_start_col as usize);
match current.direction {
SearchDirection::Diagonal => {
query_row -= 1;
match_start_col -= 1;
matched_lines += 1;
}
SearchDirection::Up => {
query_row -= 1;
}
SearchDirection::Left => {
match_start_col -= 1;
}
}
}
let buffer_row_start = match_start_col + point_range.start.row;
let buffer_row_end = match_end_col + point_range.start.row;
let matched_buffer_row_count = buffer_row_end - buffer_row_start;
let matched_ratio = matched_lines as f32
/ (matched_buffer_row_count as f32).max(query_line_count as f32);
if matched_ratio >= 0.8 {
let buffer_start_ix = self
.snapshot
.point_to_offset(Point::new(buffer_row_start, 0));
let buffer_end_ix = self.snapshot.point_to_offset(Point::new(
buffer_row_end - 1,
self.snapshot.line_len(buffer_row_end - 1),
));
return Some((best_cost, buffer_start_ix..buffer_end_ix));
}
}
None
}
}
fn fuzzy_eq(left: &str, right: &str) -> bool {
const THRESHOLD: f64 = 0.8;
let min_levenshtein = left.len().abs_diff(right.len());
let min_normalized_levenshtein =
1. - (min_levenshtein as f64 / cmp::max(left.len(), right.len()) as f64);
if min_normalized_levenshtein < THRESHOLD {
return false;
}
strsim::normalized_levenshtein(left, right) >= THRESHOLD
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum SearchDirection {
Up,
Left,
Diagonal,
}
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct SearchState {
cost: u32,
direction: SearchDirection,
}
impl SearchState {
fn new(cost: u32, direction: SearchDirection) -> Self {
Self { cost, direction }
}
}
struct SearchMatrix {
cols: usize,
rows: usize,
data: Vec<SearchState>,
}
impl SearchMatrix {
fn new(cols: usize) -> Self {
SearchMatrix {
cols,
rows: 0,
data: Vec::new(),
}
}
fn reset(&mut self, rows: usize, cols: usize) {
self.rows = rows;
self.cols = cols;
self.data
.fill(SearchState::new(0, SearchDirection::Diagonal));
self.data.resize(
self.rows * self.cols,
SearchState::new(0, SearchDirection::Diagonal),
);
}
fn get(&self, row: usize, col: usize) -> SearchState {
debug_assert!(row < self.rows);
debug_assert!(col < self.cols);
self.data[row * self.cols + col]
}
fn set(&mut self, row: usize, col: usize, state: SearchState) {
debug_assert!(row < self.rows && col < self.cols);
self.data[row * self.cols + col] = state;
}
}
#[cfg(test)]
mod tests {
use super::*;
use gpui::TestAppContext;
use indoc::indoc;
use language::Point;
use project::{FakeFs, Project};
use serde_json::json;
use settings::SettingsStore;
use util::path;
#[test]
fn test_extract_xml_edits() {
let input = indoc! {r#"
<edits path="test.rs">
<old_text>
old content
</old_text>
<new_text>
new content
</new_text>
</edits>
"#};
let result = extract_xml_replacements(input).unwrap();
assert_eq!(result.file_path, "test.rs");
assert_eq!(result.replacements.len(), 1);
assert_eq!(result.replacements[0].0, "old content");
assert_eq!(result.replacements[0].1, "new content");
}
#[test]
fn test_extract_xml_edits_with_wrong_closing_tags() {
let input = indoc! {r#"
<edits path="test.rs">
<old_text>
old content
</new_text>
<new_text>
new content
</old_text>
</ edits >
"#};
let result = extract_xml_replacements(input).unwrap();
assert_eq!(result.file_path, "test.rs");
assert_eq!(result.replacements.len(), 1);
assert_eq!(result.replacements[0].0, "old content");
assert_eq!(result.replacements[0].1, "new content");
}
#[test]
fn test_extract_xml_edits_with_xml_like_content() {
let input = indoc! {r#"
<edits path="component.tsx">
<old_text>
<foo><bar></bar></foo>
</old_text>
<new_text>
<foo><bar><baz></baz></bar></foo>
</new_text>
</edits>
"#};
let result = extract_xml_replacements(input).unwrap();
assert_eq!(result.file_path, "component.tsx");
assert_eq!(result.replacements.len(), 1);
assert_eq!(result.replacements[0].0, "<foo><bar></bar></foo>");
assert_eq!(
result.replacements[0].1,
"<foo><bar><baz></baz></bar></foo>"
);
}
#[test]
fn test_extract_xml_edits_with_conflicting_content() {
let input = indoc! {r#"
<edits path="component.tsx">
<old_text>
<new_text></new_text>
</old_text>
<new_text>
<old_text></old_text>
</new_text>
</edits>
"#};
let result = extract_xml_replacements(input).unwrap();
assert_eq!(result.file_path, "component.tsx");
assert_eq!(result.replacements.len(), 1);
assert_eq!(result.replacements[0].0, "<new_text></new_text>");
assert_eq!(result.replacements[0].1, "<old_text></old_text>");
}
#[test]
fn test_extract_xml_edits_multiple_pairs() {
let input = indoc! {r#"
Some reasoning before edits. Lots of thinking going on here
<edits path="test.rs">
<old_text>
first old
</old_text>
<new_text>
first new
</new_text>
<old_text>
second old
</edits>
<new_text>
second new
</old_text>
</edits>
"#};
let result = extract_xml_replacements(input).unwrap();
assert_eq!(result.file_path, "test.rs");
assert_eq!(result.replacements.len(), 2);
assert_eq!(result.replacements[0].0, "first old");
assert_eq!(result.replacements[0].1, "first new");
assert_eq!(result.replacements[1].0, "second old");
assert_eq!(result.replacements[1].1, "second new");
}
#[test]
fn test_extract_xml_edits_unexpected_eof() {
let input = indoc! {r#"
<edits path="test.rs">
<old_text>
first old
</
"#};
extract_xml_replacements(input).expect_err("Unexpected end of file");
}
#[gpui::test]
async fn test_parse_xml_edits(cx: &mut TestAppContext) {
let fs = init_test(cx);
let buffer_1_text = indoc! {r#"
one two three four
five six seven eight
nine ten eleven twelve
thirteen fourteen fifteen
sixteen seventeen eighteen
"#};
fs.insert_tree(
path!("/root"),
json!({
"file1": buffer_1_text,
}),
)
.await;
let project = Project::test(fs, [path!("/root").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(path!("/root/file1"), cx)
})
.await
.unwrap();
let buffer_snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
let edits = indoc! {r#"
<edits path="root/file1">
<old_text>
nine ten eleven twelve
</old_text>
<new_text>
nine TEN eleven twelve!
</new_text>
</edits>
"#};
let included_ranges = [(buffer_snapshot.anchor_before(Point::new(1, 0))..Anchor::MAX)];
let (buffer, edits) = parse_xml_edits(edits, |_path| {
Some((&buffer_snapshot, included_ranges.as_slice()))
})
.await
.unwrap();
let edits = edits
.into_iter()
.map(|(range, text)| (range.to_point(&buffer), text))
.collect::<Vec<_>>();
assert_eq!(
edits,
&[
(Point::new(2, 5)..Point::new(2, 8), "TEN".into()),
(Point::new(2, 22)..Point::new(2, 22), "!".into())
]
);
}
fn init_test(cx: &mut TestAppContext) -> Arc<FakeFs> {
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
});
FakeFs::new(cx.background_executor.clone())
}
}
crates/zeta/src/zeta.rs
+3119 -2068
View File
File diff suppressed because it is too large Load Diff
crates/zeta/src/zeta1.rs
+500
View File
@@ -0,0 +1,500 @@
mod input_excerpt;
use std::{fmt::Write, ops::Range, path::Path, sync::Arc, time::Instant};
use crate::{
EditPredictionId, ZedUpdateRequiredError, Zeta,
prediction::{EditPrediction, EditPredictionInputs},
};
use anyhow::{Context as _, Result};
use cloud_llm_client::{
PredictEditsBody, PredictEditsGitInfo, PredictEditsResponse, predict_edits_v3::Event,
};
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, SharedString, Task};
use input_excerpt::excerpt_for_cursor_position;
use language::{
Anchor, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToPoint as _, text_diff,
};
use project::{Project, ProjectPath};
use release_channel::AppVersion;
use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification};
const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
pub(crate) const MAX_CONTEXT_TOKENS: usize = 150;
pub(crate) const MAX_REWRITE_TOKENS: usize = 350;
pub(crate) const MAX_EVENT_TOKENS: usize = 500;
pub(crate) fn request_prediction_with_zeta1(
zeta: &mut Zeta,
project: &Entity<Project>,
buffer: &Entity<Buffer>,
position: language::Anchor,
cx: &mut Context<Zeta>,
) -> Task<Result<Option<EditPrediction>>> {
let buffer = buffer.clone();
let buffer_snapshotted_at = Instant::now();
let snapshot = buffer.read(cx).snapshot();
let client = zeta.client.clone();
let llm_token = zeta.llm_token.clone();
let app_version = AppVersion::global(cx);
let zeta_project = zeta.get_or_init_zeta_project(project, cx);
let events = Arc::new(zeta_project.events(cx));
let (git_info, can_collect_file) = if let Some(file) = snapshot.file() {
let can_collect_file = zeta.can_collect_file(project, file, cx);
let git_info = if can_collect_file {
git_info_for_file(project, &ProjectPath::from_file(file.as_ref(), cx), cx)
} else {
None
};
(git_info, can_collect_file)
} else {
(None, false)
};
let full_path: Arc<Path> = snapshot
.file()
.map(|f| Arc::from(f.full_path(cx).as_path()))
.unwrap_or_else(|| Arc::from(Path::new("untitled")));
let full_path_str = full_path.to_string_lossy().into_owned();
let cursor_point = position.to_point(&snapshot);
let prompt_for_events = {
let events = events.clone();
move || prompt_for_events_impl(&events, MAX_EVENT_TOKENS)
};
let gather_task = gather_context(
full_path_str,
&snapshot,
cursor_point,
prompt_for_events,
cx,
);
cx.spawn(async move |this, cx| {
let GatherContextOutput {
mut body,
context_range,
editable_range,
included_events_count,
} = gather_task.await?;
let done_gathering_context_at = Instant::now();
let included_events = &events[events.len() - included_events_count..events.len()];
body.can_collect_data = can_collect_file
&& this
.read_with(cx, |this, _| this.can_collect_events(included_events))
.unwrap_or(false);
if body.can_collect_data {
body.git_info = git_info;
}
log::debug!(
"Events:\n{}\nExcerpt:\n{:?}",
body.input_events,
body.input_excerpt
);
let http_client = client.http_client();
let response = Zeta::send_api_request::<PredictEditsResponse>(
|request| {
let uri = if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
predict_edits_url
} else {
http_client
.build_zed_llm_url("/predict_edits/v2", &[])?
.as_str()
.into()
};
Ok(request
.uri(uri)
.body(serde_json::to_string(&body)?.into())?)
},
client,
llm_token,
app_version,
)
.await;
let inputs = EditPredictionInputs {
events: included_events.into(),
included_files: vec![cloud_llm_client::predict_edits_v3::IncludedFile {
path: full_path.clone(),
max_row: cloud_llm_client::predict_edits_v3::Line(snapshot.max_point().row),
excerpts: vec![cloud_llm_client::predict_edits_v3::Excerpt {
start_line: cloud_llm_client::predict_edits_v3::Line(context_range.start.row),
text: snapshot
.text_for_range(context_range)
.collect::<String>()
.into(),
}],
}],
cursor_point: cloud_llm_client::predict_edits_v3::Point {
column: cursor_point.column,
line: cloud_llm_client::predict_edits_v3::Line(cursor_point.row),
},
cursor_path: full_path,
};
// let response = perform_predict_edits(PerformPredictEditsParams {
// client,
// llm_token,
// app_version,
// body,
// })
// .await;
let (response, usage) = match response {
Ok(response) => response,
Err(err) => {
if err.is::<ZedUpdateRequiredError>() {
cx.update(|cx| {
this.update(cx, |zeta, _cx| {
zeta.update_required = true;
})
.ok();
let error_message: SharedString = err.to_string().into();
show_app_notification(
NotificationId::unique::<ZedUpdateRequiredError>(),
cx,
move |cx| {
cx.new(|cx| {
ErrorMessagePrompt::new(error_message.clone(), cx)
.with_link_button("Update Zed", "https://zed.dev/releases")
})
},
);
})
.ok();
}
return Err(err);
}
};
let received_response_at = Instant::now();
log::debug!("completion response: {}", &response.output_excerpt);
if let Some(usage) = usage {
this.update(cx, |this, cx| {
this.user_store.update(cx, |user_store, cx| {
user_store.update_edit_prediction_usage(usage, cx);
});
})
.ok();
}
let edit_prediction = process_completion_response(
response,
buffer,
&snapshot,
editable_range,
inputs,
buffer_snapshotted_at,
received_response_at,
cx,
)
.await;
let finished_at = Instant::now();
// record latency for ~1% of requests
if rand::random::<u8>() <= 2 {
telemetry::event!(
"Edit Prediction Request",
context_latency = done_gathering_context_at
.duration_since(buffer_snapshotted_at)
.as_millis(),
request_latency = received_response_at
.duration_since(done_gathering_context_at)
.as_millis(),
process_latency = finished_at.duration_since(received_response_at).as_millis()
);
}
edit_prediction
})
}
fn process_completion_response(
prediction_response: PredictEditsResponse,
buffer: Entity<Buffer>,
snapshot: &BufferSnapshot,
editable_range: Range<usize>,
inputs: EditPredictionInputs,
buffer_snapshotted_at: Instant,
received_response_at: Instant,
cx: &AsyncApp,
) -> Task<Result<Option<EditPrediction>>> {
let snapshot = snapshot.clone();
let request_id = prediction_response.request_id;
let output_excerpt = prediction_response.output_excerpt;
cx.spawn(async move |cx| {
let output_excerpt: Arc<str> = output_excerpt.into();
let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx
.background_spawn({
let output_excerpt = output_excerpt.clone();
let editable_range = editable_range.clone();
let snapshot = snapshot.clone();
async move { parse_edits(output_excerpt, editable_range, &snapshot) }
})
.await?
.into();
Ok(EditPrediction::new(
EditPredictionId(request_id.into()),
&buffer,
&snapshot,
edits,
buffer_snapshotted_at,
received_response_at,
inputs,
cx,
)
.await)
})
}
fn parse_edits(
output_excerpt: Arc<str>,
editable_range: Range<usize>,
snapshot: &BufferSnapshot,
) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
let content = output_excerpt.replace(CURSOR_MARKER, "");
let start_markers = content
.match_indices(EDITABLE_REGION_START_MARKER)
.collect::<Vec<_>>();
anyhow::ensure!(
start_markers.len() == 1,
"expected exactly one start marker, found {}",
start_markers.len()
);
let end_markers = content
.match_indices(EDITABLE_REGION_END_MARKER)
.collect::<Vec<_>>();
anyhow::ensure!(
end_markers.len() == 1,
"expected exactly one end marker, found {}",
end_markers.len()
);
let sof_markers = content
.match_indices(START_OF_FILE_MARKER)
.collect::<Vec<_>>();
anyhow::ensure!(
sof_markers.len() <= 1,
"expected at most one start-of-file marker, found {}",
sof_markers.len()
);
let codefence_start = start_markers[0].0;
let content = &content[codefence_start..];
let newline_ix = content.find('\n').context("could not find newline")?;
let content = &content[newline_ix + 1..];
let codefence_end = content
.rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
.context("could not find end marker")?;
let new_text = &content[..codefence_end];
let old_text = snapshot
.text_for_range(editable_range.clone())
.collect::<String>();
Ok(compute_edits(
old_text,
new_text,
editable_range.start,
snapshot,
))
}
pub fn compute_edits(
old_text: String,
new_text: &str,
offset: usize,
snapshot: &BufferSnapshot,
) -> Vec<(Range<Anchor>, Arc<str>)> {
text_diff(&old_text, new_text)
.into_iter()
.map(|(mut old_range, new_text)| {
old_range.start += offset;
old_range.end += offset;
let prefix_len = common_prefix(
snapshot.chars_for_range(old_range.clone()),
new_text.chars(),
);
old_range.start += prefix_len;
let suffix_len = common_prefix(
snapshot.reversed_chars_for_range(old_range.clone()),
new_text[prefix_len..].chars().rev(),
);
old_range.end = old_range.end.saturating_sub(suffix_len);
let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
let range = if old_range.is_empty() {
let anchor = snapshot.anchor_after(old_range.start);
anchor..anchor
} else {
snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
};
(range, new_text)
})
.collect()
}
fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
a.zip(b)
.take_while(|(a, b)| a == b)
.map(|(a, _)| a.len_utf8())
.sum()
}
fn git_info_for_file(
project: &Entity<Project>,
project_path: &ProjectPath,
cx: &App,
) -> Option<PredictEditsGitInfo> {
let git_store = project.read(cx).git_store().read(cx);
if let Some((repository, _repo_path)) =
git_store.repository_and_path_for_project_path(project_path, cx)
{
let repository = repository.read(cx);
let head_sha = repository
.head_commit
.as_ref()
.map(|head_commit| head_commit.sha.to_string());
let remote_origin_url = repository.remote_origin_url.clone();
let remote_upstream_url = repository.remote_upstream_url.clone();
if head_sha.is_none() && remote_origin_url.is_none() && remote_upstream_url.is_none() {
return None;
}
Some(PredictEditsGitInfo {
head_sha,
remote_origin_url,
remote_upstream_url,
})
} else {
None
}
}
pub struct GatherContextOutput {
pub body: PredictEditsBody,
pub context_range: Range<Point>,
pub editable_range: Range<usize>,
pub included_events_count: usize,
}
pub fn gather_context(
full_path_str: String,
snapshot: &BufferSnapshot,
cursor_point: language::Point,
prompt_for_events: impl FnOnce() -> (String, usize) + Send + 'static,
cx: &App,
) -> Task<Result<GatherContextOutput>> {
cx.background_spawn({
let snapshot = snapshot.clone();
async move {
let input_excerpt = excerpt_for_cursor_position(
cursor_point,
&full_path_str,
&snapshot,
MAX_REWRITE_TOKENS,
MAX_CONTEXT_TOKENS,
);
let (input_events, included_events_count) = prompt_for_events();
let editable_range = input_excerpt.editable_range.to_offset(&snapshot);
let body = PredictEditsBody {
input_events,
input_excerpt: input_excerpt.prompt,
can_collect_data: false,
diagnostic_groups: None,
git_info: None,
outline: None,
speculated_output: None,
};
Ok(GatherContextOutput {
body,
context_range: input_excerpt.context_range,
editable_range,
included_events_count,
})
}
})
}
fn prompt_for_events_impl(events: &[Arc<Event>], mut remaining_tokens: usize) -> (String, usize) {
let mut result = String::new();
for (ix, event) in events.iter().rev().enumerate() {
let event_string = format_event(event.as_ref());
let event_tokens = guess_token_count(event_string.len());
if event_tokens > remaining_tokens {
return (result, ix);
}
if !result.is_empty() {
result.insert_str(0, "\n\n");
}
result.insert_str(0, &event_string);
remaining_tokens -= event_tokens;
}
return (result, events.len());
}
pub fn format_event(event: &Event) -> String {
match event {
Event::BufferChange {
path,
old_path,
diff,
..
} => {
let mut prompt = String::new();
if old_path != path {
writeln!(
prompt,
"User renamed {} to {}\n",
old_path.display(),
path.display()
)
.unwrap();
}
if !diff.is_empty() {
write!(
prompt,
"User edited {}:\n```diff\n{}\n```",
path.display(),
diff
)
.unwrap();
}
prompt
}
}
}
/// Typical number of string bytes per token for the purposes of limiting model input. This is
/// intentionally low to err on the side of underestimating limits.
pub(crate) const BYTES_PER_TOKEN_GUESS: usize = 3;
fn guess_token_count(bytes: usize) -> usize {
bytes / BYTES_PER_TOKEN_GUESS
}
crates/zeta/src/input_excerpt.rs → crates/zeta/src/zeta1/input_excerpt.rs
+4 -2
View File
@@ -1,4 +1,4 @@
use crate::{
use super::{
CURSOR_MARKER, EDITABLE_REGION_END_MARKER, EDITABLE_REGION_START_MARKER, START_OF_FILE_MARKER,
guess_token_count,
};
@@ -7,6 +7,7 @@ use std::{fmt::Write, ops::Range};
#[derive(Debug)]
pub struct InputExcerpt {
pub context_range: Range<Point>,
pub editable_range: Range<Point>,
pub prompt: String,
}
@@ -63,6 +64,7 @@ pub fn excerpt_for_cursor_position(
write!(prompt, "\n```").unwrap();
InputExcerpt {
context_range,
editable_range,
prompt,
}
@@ -124,7 +126,7 @@ mod tests {
use super::*;
use gpui::{App, AppContext};
use indoc::indoc;
use language::{Buffer, Language, LanguageConfig, LanguageMatcher};
use language::{Buffer, Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
use std::sync::Arc;
#[gpui::test]
crates/zeta/src/zeta_tests.rs
+671
View File
@@ -0,0 +1,671 @@
use client::test::FakeServer;
use clock::{FakeSystemClock, ReplicaId};
use cloud_api_types::{CreateLlmTokenResponse, LlmToken};
use cloud_llm_client::{PredictEditsBody, PredictEditsResponse};
use gpui::TestAppContext;
use http_client::FakeHttpClient;
use indoc::indoc;
use language::Point;
use parking_lot::Mutex;
use serde_json::json;
use settings::SettingsStore;
use util::{path, rel_path::rel_path};
use crate::zeta1::MAX_EVENT_TOKENS;
use super::*;
const BSD_0_TXT: &str = include_str!("../license_examples/0bsd.txt");
#[gpui::test]
async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
});
let edit_preview = cx
.read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
.await;
let completion = EditPrediction {
edits,
edit_preview,
buffer: buffer.clone(),
snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
id: EditPredictionId("the-id".into()),
inputs: EditPredictionInputs {
events: Default::default(),
included_files: Default::default(),
cursor_point: cloud_llm_client::predict_edits_v3::Point {
line: Line(0),
column: 0,
},
cursor_path: Path::new("").into(),
},
buffer_snapshotted_at: Instant::now(),
response_received_at: Instant::now(),
};
cx.update(|cx| {
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..5, "REM".into()), (9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..2, "REM".into()), (6..8, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.undo(cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(2..5, "REM".into()), (9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(3..3, "EM".into()), (7..9, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into()), (8..10, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(9..11, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into()), (8..10, "".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
assert_eq!(
from_completion_edits(
&completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
&buffer,
cx
),
vec![(4..4, "M".into())]
);
buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
})
}
#[gpui::test]
async fn test_clean_up_diff(cx: &mut TestAppContext) {
init_test(cx);
assert_eq!(
apply_edit_prediction(
indoc! {"
fn main() {
let word_1 = \"lorem\";
let range = word.len()..word.len();
}
"},
indoc! {"
<|editable_region_start|>
fn main() {
let word_1 = \"lorem\";
let range = word_1.len()..word_1.len();
}
<|editable_region_end|>
"},
cx,
)
.await,
indoc! {"
fn main() {
let word_1 = \"lorem\";
let range = word_1.len()..word_1.len();
}
"},
);
assert_eq!(
apply_edit_prediction(
indoc! {"
fn main() {
let story = \"the quick\"
}
"},
indoc! {"
<|editable_region_start|>
fn main() {
let story = \"the quick brown fox jumps over the lazy dog\";
}
<|editable_region_end|>
"},
cx,
)
.await,
indoc! {"
fn main() {
let story = \"the quick brown fox jumps over the lazy dog\";
}
"},
);
}
#[gpui::test]
async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
init_test(cx);
let buffer_content = "lorem\n";
let completion_response = indoc! {"
```animals.js
<|start_of_file|>
<|editable_region_start|>
lorem
ipsum
<|editable_region_end|>
```"};
assert_eq!(
apply_edit_prediction(buffer_content, completion_response, cx).await,
"lorem\nipsum"
);
}
#[gpui::test]
async fn test_can_collect_data(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree(path!("/project"), json!({ "LICENSE": BSD_0_TXT }))
.await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(path!("/project/src/main.rs"), cx)
})
.await
.unwrap();
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
true
);
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Disabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_no_data_collection_for_remote_file(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
let project = Project::test(fs.clone(), [], cx).await;
let buffer = cx.new(|_cx| {
Buffer::remote(
language::BufferId::new(1).unwrap(),
ReplicaId::new(1),
language::Capability::ReadWrite,
"fn main() {\n println!(\"Hello\");\n}",
)
});
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_no_data_collection_for_private_file(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree(
path!("/project"),
json!({
"LICENSE": BSD_0_TXT,
".env": "SECRET_KEY=secret"
}),
)
.await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/project/.env", cx)
})
.await
.unwrap();
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_no_data_collection_for_untitled_buffer(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
let project = Project::test(fs.clone(), [], cx).await;
let buffer = cx.new(|cx| Buffer::local("", cx));
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_no_data_collection_when_closed_source(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree(path!("/project"), json!({ "main.rs": "fn main() {}" }))
.await;
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer("/project/main.rs", cx)
})
.await
.unwrap();
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_data_collection_status_changes_on_move(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree(
path!("/open_source_worktree"),
json!({ "LICENSE": BSD_0_TXT, "main.rs": "" }),
)
.await;
fs.insert_tree(path!("/closed_source_worktree"), json!({ "main.rs": "" }))
.await;
let project = Project::test(
fs.clone(),
[
path!("/open_source_worktree").as_ref(),
path!("/closed_source_worktree").as_ref(),
],
cx,
)
.await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(path!("/open_source_worktree/main.rs"), cx)
})
.await
.unwrap();
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
true
);
let closed_source_file = project
.update(cx, |project, cx| {
let worktree2 = project
.worktree_for_root_name("closed_source_worktree", cx)
.unwrap();
worktree2.update(cx, |worktree2, cx| {
worktree2.load_file(rel_path("main.rs"), cx)
})
})
.await
.unwrap()
.file;
buffer.update(cx, |buffer, cx| {
buffer.file_updated(closed_source_file, cx);
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
}
#[gpui::test]
async fn test_no_data_collection_for_events_in_uncollectable_buffers(cx: &mut TestAppContext) {
init_test(cx);
let fs = project::FakeFs::new(cx.executor());
fs.insert_tree(
path!("/worktree1"),
json!({ "LICENSE": BSD_0_TXT, "main.rs": "", "other.rs": "" }),
)
.await;
fs.insert_tree(path!("/worktree2"), json!({ "private.rs": "" }))
.await;
let project = Project::test(
fs.clone(),
[path!("/worktree1").as_ref(), path!("/worktree2").as_ref()],
cx,
)
.await;
let buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(path!("/worktree1/main.rs"), cx)
})
.await
.unwrap();
let private_buffer = project
.update(cx, |project, cx| {
project.open_local_buffer(path!("/worktree2/file.rs"), cx)
})
.await
.unwrap();
let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
zeta.update(cx, |zeta, _cx| {
zeta.data_collection_choice = DataCollectionChoice::Enabled
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
true
);
// this has a side effect of registering the buffer to watch for edits
run_edit_prediction(&private_buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
private_buffer.update(cx, |private_buffer, cx| {
private_buffer.edit([(0..0, "An edit for the history!")], None, cx);
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
false
);
// make an edit that uses too many bytes, causing private_buffer edit to not be able to be
// included
buffer.update(cx, |buffer, cx| {
buffer.edit(
[(
0..0,
" ".repeat(MAX_EVENT_TOKENS * zeta1::BYTES_PER_TOKEN_GUESS),
)],
None,
cx,
);
});
run_edit_prediction(&buffer, &project, &zeta, cx).await;
assert_eq!(
captured_request.lock().clone().unwrap().can_collect_data,
true
);
}
fn init_test(cx: &mut TestAppContext) {
cx.update(|cx| {
let settings_store = SettingsStore::test(cx);
cx.set_global(settings_store);
});
}
async fn apply_edit_prediction(
buffer_content: &str,
completion_response: &str,
cx: &mut TestAppContext,
) -> String {
let fs = project::FakeFs::new(cx.executor());
let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
let (zeta, _, response) = make_test_zeta(&project, cx).await;
*response.lock() = completion_response.to_string();
let edit_prediction = run_edit_prediction(&buffer, &project, &zeta, cx).await;
buffer.update(cx, |buffer, cx| {
buffer.edit(edit_prediction.edits.iter().cloned(), None, cx)
});
buffer.read_with(cx, |buffer, _| buffer.text())
}
async fn run_edit_prediction(
buffer: &Entity<Buffer>,
project: &Entity<Project>,
zeta: &Entity<Zeta>,
cx: &mut TestAppContext,
) -> EditPrediction {
let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
zeta.update(cx, |zeta, cx| zeta.register_buffer(buffer, &project, cx));
cx.background_executor.run_until_parked();
let prediction_task = zeta.update(cx, |zeta, cx| {
zeta.request_prediction(&project, buffer, cursor, cx)
});
prediction_task.await.unwrap().unwrap()
}
async fn make_test_zeta(
project: &Entity<Project>,
cx: &mut TestAppContext,
) -> (
Entity<Zeta>,
Arc<Mutex<Option<PredictEditsBody>>>,
Arc<Mutex<String>>,
) {
let default_response = indoc! {"
```main.rs
<|start_of_file|>
<|editable_region_start|>
hello world
<|editable_region_end|>
```"
};
let captured_request: Arc<Mutex<Option<PredictEditsBody>>> = Arc::new(Mutex::new(None));
let completion_response: Arc<Mutex<String>> =
Arc::new(Mutex::new(default_response.to_string()));
let http_client = FakeHttpClient::create({
let captured_request = captured_request.clone();
let completion_response = completion_response.clone();
let mut next_request_id = 0;
move |req| {
let captured_request = captured_request.clone();
let completion_response = completion_response.clone();
async move {
match (req.method(), req.uri().path()) {
(&Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
.status(200)
.body(
serde_json::to_string(&CreateLlmTokenResponse {
token: LlmToken("the-llm-token".to_string()),
})
.unwrap()
.into(),
)
.unwrap()),
(&Method::POST, "/predict_edits/v2") => {
let mut request_body = String::new();
req.into_body().read_to_string(&mut request_body).await?;
*captured_request.lock() =
Some(serde_json::from_str(&request_body).unwrap());
next_request_id += 1;
Ok(http_client::Response::builder()
.status(200)
.body(
serde_json::to_string(&PredictEditsResponse {
request_id: format!("request-{next_request_id}"),
output_excerpt: completion_response.lock().clone(),
})
.unwrap()
.into(),
)
.unwrap())
}
_ => Ok(http_client::Response::builder()
.status(404)
.body("Not Found".into())
.unwrap()),
}
}
}
});
let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
cx.update(|cx| {
RefreshLlmTokenListener::register(client.clone(), cx);
});
let _server = FakeServer::for_client(42, &client, cx).await;
let zeta = cx.new(|cx| {
let mut zeta = Zeta::new(client, project.read(cx).user_store(), cx);
zeta.set_edit_prediction_model(ZetaEditPredictionModel::Zeta1);
let worktrees = project.read(cx).worktrees(cx).collect::<Vec<_>>();
for worktree in worktrees {
let worktree_id = worktree.read(cx).id();
zeta.get_or_init_zeta_project(project, cx)
.license_detection_watchers
.entry(worktree_id)
.or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
}
zeta
});
(zeta, captured_request, completion_response)
}
fn to_completion_edits(
iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
buffer: &Entity<Buffer>,
cx: &App,
) -> Vec<(Range<Anchor>, Arc<str>)> {
let buffer = buffer.read(cx);
iterator
.into_iter()
.map(|(range, text)| {
(
buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
text,
)
})
.collect()
}
fn from_completion_edits(
editor_edits: &[(Range<Anchor>, Arc<str>)],
buffer: &Entity<Buffer>,
cx: &App,
) -> Vec<(Range<usize>, Arc<str>)> {
let buffer = buffer.read(cx);
editor_edits
.iter()
.map(|(range, text)| {
(
range.start.to_offset(buffer)..range.end.to_offset(buffer),
text.clone(),
)
})
.collect()
}
#[ctor::ctor]
fn init_logger() {
zlog::init_test();
}