zeta2: Build edit prediction prompt and process model output in client (#41870)
Release Notes: - N/A --------- Co-authored-by: Agus Zubiaga <agus@zed.dev> Co-authored-by: Ben Kunkle <ben@zed.dev> Co-authored-by: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
This commit is contained in:
Max Brunsfeld
co-authored by
Agus Zubiaga
Ben Kunkle
Piotr Osiewicz
Agus Zubiaga
Ben Kunkle
Piotr Osiewicz
parent
fb87972f44
commit
784fdcaee3
@@ -45,7 +45,6 @@ struct RetrievalRun {
|
||||
started_at: Instant,
|
||||
search_results_generated_at: Option<Instant>,
|
||||
search_results_executed_at: Option<Instant>,
|
||||
search_results_filtered_at: Option<Instant>,
|
||||
finished_at: Option<Instant>,
|
||||
}
|
||||
|
||||
@@ -117,17 +116,12 @@ impl Zeta2ContextView {
|
||||
self.handle_search_queries_executed(info, window, cx);
|
||||
}
|
||||
}
|
||||
ZetaDebugInfo::SearchResultsFiltered(info) => {
|
||||
if info.project == self.project {
|
||||
self.handle_search_results_filtered(info, window, cx);
|
||||
}
|
||||
}
|
||||
ZetaDebugInfo::ContextRetrievalFinished(info) => {
|
||||
if info.project == self.project {
|
||||
self.handle_context_retrieval_finished(info, window, cx);
|
||||
}
|
||||
}
|
||||
ZetaDebugInfo::EditPredicted(_) => {}
|
||||
ZetaDebugInfo::EditPredictionRequested(_) => {}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -159,7 +153,6 @@ impl Zeta2ContextView {
|
||||
started_at: info.timestamp,
|
||||
search_results_generated_at: None,
|
||||
search_results_executed_at: None,
|
||||
search_results_filtered_at: None,
|
||||
finished_at: None,
|
||||
});
|
||||
|
||||
@@ -218,18 +211,18 @@ impl Zeta2ContextView {
|
||||
|
||||
run.search_results_generated_at = Some(info.timestamp);
|
||||
run.search_queries = info
|
||||
.queries
|
||||
.regex_by_glob
|
||||
.into_iter()
|
||||
.map(|query| {
|
||||
.map(|(glob, regex)| {
|
||||
let mut regex_parser = regex_syntax::ast::parse::Parser::new();
|
||||
|
||||
GlobQueries {
|
||||
glob: query.glob,
|
||||
alternations: match regex_parser.parse(&query.regex) {
|
||||
glob,
|
||||
alternations: match regex_parser.parse(®ex) {
|
||||
Ok(regex_syntax::ast::Ast::Alternation(ref alt)) => {
|
||||
alt.asts.iter().map(|ast| ast.to_string()).collect()
|
||||
}
|
||||
_ => vec![query.regex],
|
||||
_ => vec![regex],
|
||||
},
|
||||
}
|
||||
})
|
||||
@@ -256,20 +249,6 @@ impl Zeta2ContextView {
|
||||
cx.notify();
|
||||
}
|
||||
|
||||
fn handle_search_results_filtered(
|
||||
&mut self,
|
||||
info: ZetaContextRetrievalDebugInfo,
|
||||
_window: &mut Window,
|
||||
cx: &mut Context<Self>,
|
||||
) {
|
||||
let Some(run) = self.runs.back_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
run.search_results_filtered_at = Some(info.timestamp);
|
||||
cx.notify();
|
||||
}
|
||||
|
||||
fn handle_go_back(
|
||||
&mut self,
|
||||
_: &Zeta2ContextGoBack,
|
||||
@@ -398,19 +377,10 @@ impl Zeta2ContextView {
|
||||
};
|
||||
div = div.child(format!("Ran search: {:>5} ms", (t2 - t1).as_millis()));
|
||||
|
||||
let Some(t3) = run.search_results_filtered_at else {
|
||||
return pending_message(div, "Filtering results...");
|
||||
};
|
||||
div =
|
||||
div.child(format!("Filtered results: {:>5} ms", (t3 - t2).as_millis()));
|
||||
|
||||
let Some(t4) = run.finished_at else {
|
||||
return pending_message(div, "Building excerpts");
|
||||
};
|
||||
div = div
|
||||
.child(format!("Build excerpts: {:>5} µs", (t4 - t3).as_micros()))
|
||||
.child(format!("Total: {:>5} ms", (t4 - t0).as_millis()));
|
||||
div
|
||||
div.child(format!(
|
||||
"Total: {:>5} ms",
|
||||
(run.finished_at.unwrap_or(t0) - t0).as_millis()
|
||||
))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use std::{cmp::Reverse, path::PathBuf, str::FromStr, sync::Arc, time::Duration};
|
||||
use chrono::TimeDelta;
|
||||
use client::{Client, UserStore};
|
||||
use cloud_llm_client::predict_edits_v3::{
|
||||
self, DeclarationScoreComponents, PredictEditsRequest, PredictEditsResponse, PromptFormat,
|
||||
DeclarationScoreComponents, PredictEditsRequest, PromptFormat,
|
||||
};
|
||||
use collections::HashMap;
|
||||
use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
|
||||
@@ -23,7 +23,7 @@ use ui_input::InputField;
|
||||
use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
|
||||
use workspace::{Item, SplitDirection, Workspace};
|
||||
use zeta2::{
|
||||
ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, LlmContextOptions, Zeta, Zeta2FeatureFlag,
|
||||
AgenticContextOptions, ContextMode, DEFAULT_SYNTAX_CONTEXT_OPTIONS, Zeta, Zeta2FeatureFlag,
|
||||
ZetaDebugInfo, ZetaEditPredictionDebugInfo, ZetaOptions,
|
||||
};
|
||||
|
||||
@@ -123,6 +123,7 @@ struct LastPrediction {
|
||||
context_editor: Entity<Editor>,
|
||||
prompt_editor: Entity<Editor>,
|
||||
retrieval_time: TimeDelta,
|
||||
request_time: Option<TimeDelta>,
|
||||
buffer: WeakEntity<Buffer>,
|
||||
position: language::Anchor,
|
||||
state: LastPredictionState,
|
||||
@@ -143,7 +144,7 @@ enum LastPredictionState {
|
||||
model_response_editor: Entity<Editor>,
|
||||
feedback_editor: Entity<Editor>,
|
||||
feedback: Option<Feedback>,
|
||||
response: predict_edits_v3::PredictEditsResponse,
|
||||
request_id: String,
|
||||
},
|
||||
Failed {
|
||||
message: String,
|
||||
@@ -217,7 +218,7 @@ impl Zeta2Inspector {
|
||||
});
|
||||
|
||||
match &options.context {
|
||||
ContextMode::Llm(_) => {
|
||||
ContextMode::Agentic(_) => {
|
||||
self.context_mode = ContextModeState::Llm;
|
||||
}
|
||||
ContextMode::Syntax(_) => {
|
||||
@@ -307,9 +308,11 @@ impl Zeta2Inspector {
|
||||
};
|
||||
|
||||
let context = match zeta_options.context {
|
||||
ContextMode::Llm(_context_options) => ContextMode::Llm(LlmContextOptions {
|
||||
excerpt: excerpt_options,
|
||||
}),
|
||||
ContextMode::Agentic(_context_options) => {
|
||||
ContextMode::Agentic(AgenticContextOptions {
|
||||
excerpt: excerpt_options,
|
||||
})
|
||||
}
|
||||
ContextMode::Syntax(context_options) => {
|
||||
let max_retrieved_declarations = match &this.context_mode {
|
||||
ContextModeState::Llm => {
|
||||
@@ -368,7 +371,7 @@ impl Zeta2Inspector {
|
||||
let language_registry = self.project.read(cx).languages().clone();
|
||||
async move |this, cx| {
|
||||
let mut languages = HashMap::default();
|
||||
let ZetaDebugInfo::EditPredicted(prediction) = prediction else {
|
||||
let ZetaDebugInfo::EditPredictionRequested(prediction) = prediction else {
|
||||
return;
|
||||
};
|
||||
for ext in prediction
|
||||
@@ -396,6 +399,8 @@ impl Zeta2Inspector {
|
||||
.await
|
||||
.log_err();
|
||||
|
||||
let json_language = language_registry.language_for_name("Json").await.log_err();
|
||||
|
||||
this.update_in(cx, |this, window, cx| {
|
||||
let context_editor = cx.new(|cx| {
|
||||
let mut excerpt_score_components = HashMap::default();
|
||||
@@ -492,25 +497,15 @@ impl Zeta2Inspector {
|
||||
|
||||
let task = cx.spawn_in(window, {
|
||||
let markdown_language = markdown_language.clone();
|
||||
let json_language = json_language.clone();
|
||||
async move |this, cx| {
|
||||
let response = response_rx.await;
|
||||
|
||||
this.update_in(cx, |this, window, cx| {
|
||||
if let Some(prediction) = this.last_prediction.as_mut() {
|
||||
prediction.state = match response {
|
||||
Ok(Ok(response)) => {
|
||||
if let Some(debug_info) = &response.debug_info {
|
||||
prediction.prompt_editor.update(
|
||||
cx,
|
||||
|prompt_editor, cx| {
|
||||
prompt_editor.set_text(
|
||||
debug_info.prompt.as_str(),
|
||||
window,
|
||||
cx,
|
||||
);
|
||||
},
|
||||
);
|
||||
}
|
||||
Ok((Ok(response), request_time)) => {
|
||||
prediction.request_time = Some(request_time);
|
||||
|
||||
let feedback_editor = cx.new(|cx| {
|
||||
let buffer = cx.new(|cx| {
|
||||
@@ -577,16 +572,11 @@ impl Zeta2Inspector {
|
||||
model_response_editor: cx.new(|cx| {
|
||||
let buffer = cx.new(|cx| {
|
||||
let mut buffer = Buffer::local(
|
||||
response
|
||||
.debug_info
|
||||
.as_ref()
|
||||
.map(|p| p.model_response.as_str())
|
||||
.unwrap_or(
|
||||
"(Debug info not available)",
|
||||
),
|
||||
serde_json::to_string_pretty(&response)
|
||||
.unwrap_or_default(),
|
||||
cx,
|
||||
);
|
||||
buffer.set_language(markdown_language, cx);
|
||||
buffer.set_language(json_language, cx);
|
||||
buffer
|
||||
});
|
||||
let buffer = cx.new(|cx| {
|
||||
@@ -607,10 +597,11 @@ impl Zeta2Inspector {
|
||||
}),
|
||||
feedback_editor,
|
||||
feedback: None,
|
||||
response,
|
||||
request_id: response.id.clone(),
|
||||
}
|
||||
}
|
||||
Ok(Err(err)) => {
|
||||
Ok((Err(err), request_time)) => {
|
||||
prediction.request_time = Some(request_time);
|
||||
LastPredictionState::Failed { message: err }
|
||||
}
|
||||
Err(oneshot::Canceled) => LastPredictionState::Failed {
|
||||
@@ -644,6 +635,7 @@ impl Zeta2Inspector {
|
||||
editor
|
||||
}),
|
||||
retrieval_time,
|
||||
request_time: None,
|
||||
buffer,
|
||||
position,
|
||||
state: LastPredictionState::Requested,
|
||||
@@ -700,7 +692,7 @@ impl Zeta2Inspector {
|
||||
feedback: feedback_state,
|
||||
feedback_editor,
|
||||
model_response_editor,
|
||||
response,
|
||||
request_id,
|
||||
..
|
||||
} = &mut last_prediction.state
|
||||
else {
|
||||
@@ -734,11 +726,10 @@ impl Zeta2Inspector {
|
||||
|
||||
telemetry::event!(
|
||||
"Zeta2 Prediction Rated",
|
||||
id = response.request_id,
|
||||
id = request_id,
|
||||
kind = kind,
|
||||
text = text,
|
||||
request = last_prediction.request,
|
||||
response = response,
|
||||
project_snapshot = project_snapshot,
|
||||
);
|
||||
})
|
||||
@@ -834,11 +825,11 @@ impl Zeta2Inspector {
|
||||
let current_options =
|
||||
this.zeta.read(cx).options().clone();
|
||||
match current_options.context.clone() {
|
||||
ContextMode::Llm(_) => {}
|
||||
ContextMode::Agentic(_) => {}
|
||||
ContextMode::Syntax(context_options) => {
|
||||
let options = ZetaOptions {
|
||||
context: ContextMode::Llm(
|
||||
LlmContextOptions {
|
||||
context: ContextMode::Agentic(
|
||||
AgenticContextOptions {
|
||||
excerpt: context_options.excerpt,
|
||||
},
|
||||
),
|
||||
@@ -865,7 +856,7 @@ impl Zeta2Inspector {
|
||||
let current_options =
|
||||
this.zeta.read(cx).options().clone();
|
||||
match current_options.context.clone() {
|
||||
ContextMode::Llm(context_options) => {
|
||||
ContextMode::Agentic(context_options) => {
|
||||
let options = ZetaOptions {
|
||||
context: ContextMode::Syntax(
|
||||
EditPredictionContextOptions {
|
||||
@@ -976,25 +967,6 @@ impl Zeta2Inspector {
|
||||
return None;
|
||||
};
|
||||
|
||||
let (prompt_planning_time, inference_time, parsing_time) =
|
||||
if let LastPredictionState::Success {
|
||||
response:
|
||||
PredictEditsResponse {
|
||||
debug_info: Some(debug_info),
|
||||
..
|
||||
},
|
||||
..
|
||||
} = &prediction.state
|
||||
{
|
||||
(
|
||||
Some(debug_info.prompt_planning_time),
|
||||
Some(debug_info.inference_time),
|
||||
Some(debug_info.parsing_time),
|
||||
)
|
||||
} else {
|
||||
(None, None, None)
|
||||
};
|
||||
|
||||
Some(
|
||||
v_flex()
|
||||
.p_4()
|
||||
@@ -1005,12 +977,7 @@ impl Zeta2Inspector {
|
||||
"Context retrieval",
|
||||
Some(prediction.retrieval_time),
|
||||
))
|
||||
.child(Self::render_duration(
|
||||
"Prompt planning",
|
||||
prompt_planning_time,
|
||||
))
|
||||
.child(Self::render_duration("Inference", inference_time))
|
||||
.child(Self::render_duration("Parsing", parsing_time)),
|
||||
.child(Self::render_duration("Request", prediction.request_time)),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user