From 0aad47493e184422b758350ea6c6330aca9b3633 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Tue, 23 Sep 2025 12:32:36 -0300 Subject: [PATCH] zeta2: Use global zeta in Inspector (#38718) The edit prediction debug tools has been renamed to zeta2 inspector because it's now zeta specific. It will now always display the last prediction request context, prompt, and model response. Release Notes: - N/A --------- Co-authored-by: Bennet Co-authored-by: Bennet Bo Fenner --- Cargo.lock | 60 +- Cargo.toml | 4 +- .../src/declaration_scoring.rs | 11 +- .../src/edit_prediction_context.rs | 2 +- crates/edit_prediction_context/src/excerpt.rs | 8 +- crates/zed/Cargo.toml | 2 +- crates/zed/src/main.rs | 2 +- crates/zeta2/Cargo.toml | 3 +- crates/zeta2/src/zeta2.rs | 116 ++- .../Cargo.toml | 16 +- .../LICENSE-GPL | 0 crates/zeta2_tools/src/zeta2_tools.rs | 663 ++++++++++++++++++ 12 files changed, 810 insertions(+), 77 deletions(-) rename crates/{edit_prediction_tools => zeta2_tools}/Cargo.toml (77%) rename crates/{edit_prediction_tools => zeta2_tools}/LICENSE-GPL (100%) create mode 100644 crates/zeta2_tools/src/zeta2_tools.rs diff --git a/Cargo.lock b/Cargo.lock index 39146342f2973691e00a9c5ab9a8d29e50539a47..c57d06ef884b78d22042324d02a78d23ec485608 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5214,33 +5214,6 @@ dependencies = [ "zlog", ] -[[package]] -name = "edit_prediction_tools" -version = "0.1.0" -dependencies = [ - "clap", - "collections", - "edit_prediction_context", - "editor", - "futures 0.3.31", - "gpui", - "indoc", - "language", - "log", - "pretty_assertions", - "project", - "serde", - "serde_json", - "settings", - "text", - "ui", - "ui_input", - "util", - "workspace", - "workspace-hack", - "zlog", -] - [[package]] name = "editor" version = "0.1.0" @@ -21274,7 +21247,6 @@ dependencies = [ "debugger_ui", "diagnostics", "edit_prediction_button", - "edit_prediction_tools", "editor", "env_logger 0.11.8", "extension", @@ -21386,6 +21358,7 @@ dependencies = [ "zed_env_vars", "zeta", "zeta2", + "zeta2_tools", "zlog", "zlog_settings", ] @@ -21669,6 +21642,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrayvec", + "chrono", "client", "cloud_llm_client", "edit_prediction", @@ -21689,6 +21663,36 @@ dependencies = [ "worktree", ] +[[package]] +name = "zeta2_tools" +version = "0.1.0" +dependencies = [ + "chrono", + "clap", + "client", + "collections", + "edit_prediction_context", + "editor", + "futures 0.3.31", + "gpui", + "indoc", + "language", + "log", + "pretty_assertions", + "project", + "serde", + "serde_json", + "settings", + "text", + "ui", + "ui_input", + "util", + "workspace", + "workspace-hack", + "zeta2", + "zlog", +] + [[package]] name = "zeta_cli" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 3c431a29eb53420dadde38a1c1ad30a1f61d44c1..8aabe6ad40c9d6b854d3453b08972dd8a4364e09 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -59,7 +59,7 @@ members = [ "crates/edit_prediction", "crates/edit_prediction_button", "crates/edit_prediction_context", - "crates/edit_prediction_tools", + "crates/zeta2_tools", "crates/editor", "crates/eval", "crates/explorer_command_injector", @@ -319,7 +319,7 @@ image_viewer = { path = "crates/image_viewer" } edit_prediction = { path = "crates/edit_prediction" } edit_prediction_button = { path = "crates/edit_prediction_button" } edit_prediction_context = { path = "crates/edit_prediction_context" } -edit_prediction_tools = { path = "crates/edit_prediction_tools" } +zeta2_tools = { path = "crates/zeta2_tools" } inspector_ui = { path = "crates/inspector_ui" } install_cli = { path = "crates/install_cli" } jj = { path = "crates/jj" } diff --git a/crates/edit_prediction_context/src/declaration_scoring.rs b/crates/edit_prediction_context/src/declaration_scoring.rs index fee7498a696c0608704dab6e8ab9f012c95660b5..f655387f680a53413161383f0678b21456c271f6 100644 --- a/crates/edit_prediction_context/src/declaration_scoring.rs +++ b/crates/edit_prediction_context/src/declaration_scoring.rs @@ -3,7 +3,7 @@ use itertools::Itertools as _; use language::BufferSnapshot; use ordered_float::OrderedFloat; use serde::Serialize; -use std::{collections::HashMap, ops::Range}; +use std::{cmp::Reverse, collections::HashMap, ops::Range}; use strum::EnumIter; use text::{Point, ToPoint}; @@ -159,11 +159,10 @@ pub fn scored_snippets( .collect::>(); snippets.sort_unstable_by_key(|snippet| { - OrderedFloat( - snippet - .score_density(SnippetStyle::Declaration) - .max(snippet.score_density(SnippetStyle::Signature)), - ) + let score_density = snippet + .score_density(SnippetStyle::Declaration) + .max(snippet.score_density(SnippetStyle::Signature)); + Reverse(OrderedFloat(score_density)) }); snippets diff --git a/crates/edit_prediction_context/src/edit_prediction_context.rs b/crates/edit_prediction_context/src/edit_prediction_context.rs index aeda74811296b70fc48198b1c3f72a50cfd7c31e..f3752abab991493660d197fe871838a71f6c8ad1 100644 --- a/crates/edit_prediction_context/src/edit_prediction_context.rs +++ b/crates/edit_prediction_context/src/edit_prediction_context.rs @@ -16,7 +16,7 @@ pub use excerpt::*; pub use reference::*; pub use syntax_index::*; -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct EditPredictionContext { pub excerpt: EditPredictionExcerpt, pub excerpt_text: EditPredictionExcerptText, diff --git a/crates/edit_prediction_context/src/excerpt.rs b/crates/edit_prediction_context/src/excerpt.rs index 764c9040f247561ba6058dfd2954f42085297a4c..58549d579dca2b589fb8da01e4963782845933e9 100644 --- a/crates/edit_prediction_context/src/excerpt.rs +++ b/crates/edit_prediction_context/src/excerpt.rs @@ -1,4 +1,4 @@ -use language::BufferSnapshot; +use language::{BufferSnapshot, LanguageId}; use std::ops::Range; use text::{Point, ToOffset as _, ToPoint as _}; use tree_sitter::{Node, TreeCursor}; @@ -20,7 +20,7 @@ use crate::{BufferDeclaration, declaration::DeclarationId, syntax_index::SyntaxI // // - Filter outer syntax layers that don't support edit prediction. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct EditPredictionExcerptOptions { /// Limit for the number of bytes in the window around the cursor. pub max_bytes: usize, @@ -31,6 +31,7 @@ pub struct EditPredictionExcerptOptions { pub target_before_cursor_over_total_bytes: f32, } +// TODO: consider merging these #[derive(Debug, Clone)] pub struct EditPredictionExcerpt { pub range: Range, @@ -42,6 +43,7 @@ pub struct EditPredictionExcerpt { pub struct EditPredictionExcerptText { pub body: String, pub parent_signatures: Vec, + pub language_id: Option, } impl EditPredictionExcerpt { @@ -54,9 +56,11 @@ impl EditPredictionExcerpt { .iter() .map(|(_, range)| buffer.text_for_range(range.clone()).collect::()) .collect(); + let language_id = buffer.language().map(|l| l.id()); EditPredictionExcerptText { body, parent_signatures, + language_id, } } diff --git a/crates/zed/Cargo.toml b/crates/zed/Cargo.toml index 3d7bbe8642c64d7cfff9bfb9169c3df17e18d854..c2d8733b8ae1f98faa7e1a5ba0d1851ba7c9ccc4 100644 --- a/crates/zed/Cargo.toml +++ b/crates/zed/Cargo.toml @@ -52,7 +52,7 @@ debugger_tools.workspace = true debugger_ui.workspace = true diagnostics.workspace = true editor.workspace = true -edit_prediction_tools.workspace = true +zeta2_tools.workspace = true env_logger.workspace = true extension.workspace = true extension_host.workspace = true diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs index 3dbbed0ce50ad84a1717f81afdf95c432b09259d..af882e76231458f160fac9354d2c57150560c00b 100644 --- a/crates/zed/src/main.rs +++ b/crates/zed/src/main.rs @@ -549,7 +549,7 @@ pub fn main() { language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx); agent_settings::init(cx); acp_tools::init(cx); - edit_prediction_tools::init(cx); + zeta2_tools::init(cx); web_search::init(cx); web_search_providers::init(app_state.client.clone(), cx); snippet_provider::init(cx); diff --git a/crates/zeta2/Cargo.toml b/crates/zeta2/Cargo.toml index 61c560baab543baf9d57b034807fe60cb566b24f..d362441cdd522e012d75bec9bde6daeff50e3e89 100644 --- a/crates/zeta2/Cargo.toml +++ b/crates/zeta2/Cargo.toml @@ -14,6 +14,7 @@ path = "src/zeta2.rs" [dependencies] anyhow.workspace = true arrayvec.workspace = true +chrono.workspace = true client.workspace = true cloud_llm_client.workspace = true edit_prediction.workspace = true @@ -29,8 +30,8 @@ serde_json.workspace = true thiserror.workspace = true util.workspace = true uuid.workspace = true -workspace.workspace = true workspace-hack.workspace = true +workspace.workspace = true worktree.workspace = true [dev-dependencies] diff --git a/crates/zeta2/src/zeta2.rs b/crates/zeta2/src/zeta2.rs index 240d44fae44d9a430e1ed64816e11428a5bdb3d0..4ecc0fa3acf9bc04935d3e7e443e0e610b907589 100644 --- a/crates/zeta2/src/zeta2.rs +++ b/crates/zeta2/src/zeta2.rs @@ -1,5 +1,6 @@ use anyhow::{Context as _, Result, anyhow}; use arrayvec::ArrayVec; +use chrono::TimeDelta; use client::{Client, EditPredictionUsage, UserStore}; use cloud_llm_client::predict_edits_v3::{self, Signature}; use cloud_llm_client::{ @@ -11,10 +12,11 @@ use edit_prediction_context::{ SyntaxIndexState, }; use futures::AsyncReadExt as _; +use futures::channel::mpsc; use gpui::http_client::Method; use gpui::{ - App, Entity, EntityId, Global, SemanticVersion, SharedString, Subscription, Task, http_client, - prelude::*, + App, Entity, EntityId, Global, SemanticVersion, SharedString, Subscription, Task, WeakEntity, + http_client, prelude::*, }; use language::{Anchor, Buffer, OffsetRangeExt as _, ToPoint}; use language::{BufferSnapshot, EditPreview}; @@ -28,7 +30,7 @@ use std::str::FromStr as _; use std::time::{Duration, Instant}; use std::{ops::Range, sync::Arc}; use thiserror::Error; -use util::ResultExt as _; +use util::{ResultExt as _, some_or_debug_panic}; use uuid::Uuid; use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification}; @@ -37,6 +39,12 @@ const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1); /// Maximum number of events to track. const MAX_EVENT_COUNT: usize = 16; +pub const DEFAULT_EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPredictionExcerptOptions { + max_bytes: 512, + min_bytes: 128, + target_before_cursor_over_total_bytes: 0.5, +}; + #[derive(Clone)] struct ZetaGlobal(Entity); @@ -50,8 +58,19 @@ pub struct Zeta { projects: HashMap, pub excerpt_options: EditPredictionExcerptOptions, update_required: bool, + debug_tx: Option>>, +} + +pub struct PredictionDebugInfo { + pub context: EditPredictionContext, + pub retrieval_time: TimeDelta, + pub request: RequestDebugInfo, + pub buffer: WeakEntity, + pub position: language::Anchor, } +pub type RequestDebugInfo = predict_edits_v3::DebugInfo; + struct ZetaProject { syntax_index: Entity, events: VecDeque, @@ -94,11 +113,7 @@ impl Zeta { projects: HashMap::new(), client, user_store, - excerpt_options: EditPredictionExcerptOptions { - max_bytes: 512, - min_bytes: 128, - target_before_cursor_over_total_bytes: 0.5, - }, + excerpt_options: DEFAULT_EXCERPT_OPTIONS, llm_token: LlmApiToken::default(), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, @@ -113,9 +128,24 @@ impl Zeta { }, ), update_required: false, + debug_tx: None, } } + pub fn debug_info(&mut self) -> mpsc::UnboundedReceiver> { + let (debug_watch_tx, debug_watch_rx) = mpsc::unbounded(); + self.debug_tx = Some(debug_watch_tx); + debug_watch_rx + } + + pub fn excerpt_options(&self) -> &EditPredictionExcerptOptions { + &self.excerpt_options + } + + pub fn set_excerpt_options(&mut self, options: EditPredictionExcerptOptions) { + self.excerpt_options = options; + } + pub fn usage(&self, cx: &App) -> Option { self.user_store.read(cx).edit_prediction_usage() } @@ -273,9 +303,11 @@ impl Zeta { .worktrees(cx) .map(|worktree| worktree.read(cx).snapshot()) .collect::>(); + let debug_tx = self.debug_tx.clone(); let request_task = cx.background_spawn({ let snapshot = snapshot.clone(); + let buffer = buffer.clone(); async move { let index_state = if let Some(index_state) = index_state { Some(index_state.lock_owned().await) @@ -285,35 +317,61 @@ impl Zeta { let cursor_point = position.to_point(&snapshot); - // TODO: make this only true if debug view is open - let debug_info = true; + let before_retrieval = chrono::Utc::now(); - let Some(request) = EditPredictionContext::gather_context( + let Some(context) = EditPredictionContext::gather_context( cursor_point, &snapshot, &excerpt_options, index_state.as_deref(), - ) - .map(|context| { - make_cloud_request( - excerpt_path.clone(), - context, - // TODO pass everything - Vec::new(), - false, - Vec::new(), - None, - debug_info, - &worktree_snapshots, - index_state.as_deref(), - ) - }) else { + ) else { return Ok(None); }; - anyhow::Ok(Some( - Self::perform_request(client, llm_token, app_version, request).await?, - )) + let debug_context = if let Some(debug_tx) = debug_tx { + Some((debug_tx, context.clone())) + } else { + None + }; + + let request = make_cloud_request( + excerpt_path.clone(), + context, + // TODO pass everything + Vec::new(), + false, + Vec::new(), + None, + debug_context.is_some(), + &worktree_snapshots, + index_state.as_deref(), + ); + + let retrieval_time = chrono::Utc::now() - before_retrieval; + let response = Self::perform_request(client, llm_token, app_version, request).await; + + if let Some((debug_tx, context)) = debug_context { + debug_tx + .unbounded_send(response.as_ref().map_err(|err| err.to_string()).and_then( + |response| { + let Some(request) = + some_or_debug_panic(response.0.debug_info.clone()) + else { + return Err("Missing debug info".to_string()); + }; + Ok(PredictionDebugInfo { + context, + request, + retrieval_time, + buffer: buffer.downgrade(), + position, + }) + }, + )) + .ok(); + } + + anyhow::Ok(Some(response?)) } }); diff --git a/crates/edit_prediction_tools/Cargo.toml b/crates/zeta2_tools/Cargo.toml similarity index 77% rename from crates/edit_prediction_tools/Cargo.toml rename to crates/zeta2_tools/Cargo.toml index ffd34abb2537006dd914d9bf9d30b735de91c5ba..fe651c14193db41f24af00d71835ffd4c6deb6eb 100644 --- a/crates/edit_prediction_tools/Cargo.toml +++ b/crates/zeta2_tools/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "edit_prediction_tools" +name = "zeta2_tools" version = "0.1.0" edition.workspace = true publish.workspace = true @@ -9,12 +9,15 @@ license = "GPL-3.0-or-later" workspace = true [lib] -path = "src/edit_prediction_tools.rs" +path = "src/zeta2_tools.rs" [dependencies] -edit_prediction_context.workspace = true +chrono.workspace = true +client.workspace = true collections.workspace = true +edit_prediction_context.workspace = true editor.workspace = true +futures.workspace = true gpui.workspace = true language.workspace = true log.workspace = true @@ -23,19 +26,20 @@ serde.workspace = true text.workspace = true ui.workspace = true ui_input.workspace = true +util.workspace = true workspace-hack.workspace = true workspace.workspace = true +zeta2.workspace = true [dev-dependencies] clap.workspace = true -futures.workspace = true gpui = { workspace = true, features = ["test-support"] } indoc.workspace = true language = { workspace = true, features = ["test-support"] } pretty_assertions.workspace = true -project = {workspace= true, features = ["test-support"]} +project = { workspace = true, features = ["test-support"] } serde_json.workspace = true -settings = {workspace= true, features = ["test-support"]} +settings = { workspace = true, features = ["test-support"] } text = { workspace = true, features = ["test-support"] } util = { workspace = true, features = ["test-support"] } zlog.workspace = true diff --git a/crates/edit_prediction_tools/LICENSE-GPL b/crates/zeta2_tools/LICENSE-GPL similarity index 100% rename from crates/edit_prediction_tools/LICENSE-GPL rename to crates/zeta2_tools/LICENSE-GPL diff --git a/crates/zeta2_tools/src/zeta2_tools.rs b/crates/zeta2_tools/src/zeta2_tools.rs new file mode 100644 index 0000000000000000000000000000000000000000..0d743e388095f9c4658d102c69faec2eb3ad0275 --- /dev/null +++ b/crates/zeta2_tools/src/zeta2_tools.rs @@ -0,0 +1,663 @@ +use std::{ + collections::hash_map::Entry, + ffi::OsStr, + path::{Path, PathBuf}, + str::FromStr, + sync::Arc, + time::Duration, +}; + +use chrono::TimeDelta; +use client::{Client, UserStore}; +use collections::HashMap; +use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer}; +use futures::StreamExt as _; +use gpui::{ + Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task, WeakEntity, actions, + prelude::*, +}; +use language::{Buffer, DiskState}; +use project::{Project, WorktreeId}; +use ui::prelude::*; +use ui_input::SingleLineInput; +use util::ResultExt; +use workspace::{Item, SplitDirection, Workspace}; +use zeta2::Zeta; + +use edit_prediction_context::{EditPredictionExcerptOptions, SnippetStyle}; + +actions!( + dev, + [ + /// Opens the language server protocol logs viewer. + OpenZeta2Inspector + ] +); + +pub fn init(cx: &mut App) { + cx.observe_new(move |workspace: &mut Workspace, _, _cx| { + workspace.register_action(move |workspace, _: &OpenZeta2Inspector, window, cx| { + let project = workspace.project(); + workspace.split_item( + SplitDirection::Right, + Box::new(cx.new(|cx| { + Zeta2Inspector::new( + &project, + workspace.client(), + workspace.user_store(), + window, + cx, + ) + })), + window, + cx, + ); + }); + }) + .detach(); +} + +pub struct Zeta2Inspector { + focus_handle: FocusHandle, + project: Entity, + last_prediction: Option, + max_bytes_input: Entity, + min_bytes_input: Entity, + cursor_context_ratio_input: Entity, + active_view: ActiveView, + zeta: Entity, + _active_editor_subscription: Option, + _update_state_task: Task<()>, + _receive_task: Task<()>, +} + +#[derive(PartialEq)] +enum ActiveView { + Context, + Inference, +} + +enum LastPredictionState { + Failed(SharedString), + Success(LastPrediction), + Replaying { + prediction: LastPrediction, + _task: Task<()>, + }, +} + +struct LastPrediction { + context_editor: Entity, + retrieval_time: TimeDelta, + prompt_planning_time: TimeDelta, + inference_time: TimeDelta, + parsing_time: TimeDelta, + prompt_editor: Entity, + model_response_editor: Entity, + buffer: WeakEntity, + position: language::Anchor, +} + +impl Zeta2Inspector { + pub fn new( + project: &Entity, + client: &Arc, + user_store: &Entity, + window: &mut Window, + cx: &mut Context, + ) -> Self { + let zeta = Zeta::global(client, user_store, cx); + let mut request_rx = zeta.update(cx, |zeta, _cx| zeta.debug_info()); + + let receive_task = cx.spawn_in(window, async move |this, cx| { + while let Some(prediction_result) = request_rx.next().await { + this.update_in(cx, |this, window, cx| match prediction_result { + Ok(prediction) => { + this.update_last_prediction(prediction, window, cx); + } + Err(err) => { + this.last_prediction = Some(LastPredictionState::Failed(err.into())); + cx.notify(); + } + }) + .ok(); + } + }); + + let mut this = Self { + focus_handle: cx.focus_handle(), + project: project.clone(), + last_prediction: None, + active_view: ActiveView::Context, + max_bytes_input: Self::number_input("Max Bytes", window, cx), + min_bytes_input: Self::number_input("Min Bytes", window, cx), + cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx), + zeta: zeta.clone(), + _active_editor_subscription: None, + _update_state_task: Task::ready(()), + _receive_task: receive_task, + }; + this.set_input_options(&zeta.read(cx).excerpt_options().clone(), window, cx); + this + } + + fn set_input_options( + &mut self, + options: &EditPredictionExcerptOptions, + window: &mut Window, + cx: &mut Context, + ) { + self.max_bytes_input.update(cx, |input, cx| { + input.set_text(options.max_bytes.to_string(), window, cx); + }); + self.min_bytes_input.update(cx, |input, cx| { + input.set_text(options.min_bytes.to_string(), window, cx); + }); + self.cursor_context_ratio_input.update(cx, |input, cx| { + input.set_text( + format!("{:.2}", options.target_before_cursor_over_total_bytes), + window, + cx, + ); + }); + cx.notify(); + } + + fn set_options(&mut self, options: EditPredictionExcerptOptions, cx: &mut Context) { + self.zeta + .update(cx, |this, _cx| this.set_excerpt_options(options)); + + const THROTTLE_TIME: Duration = Duration::from_millis(100); + + if let Some( + LastPredictionState::Success(prediction) + | LastPredictionState::Replaying { prediction, .. }, + ) = self.last_prediction.take() + { + if let Some(buffer) = prediction.buffer.upgrade() { + let position = prediction.position; + let zeta = self.zeta.clone(); + let project = self.project.clone(); + let task = cx.spawn(async move |_this, cx| { + cx.background_executor().timer(THROTTLE_TIME).await; + if let Some(task) = zeta + .update(cx, |zeta, cx| { + zeta.request_prediction(&project, &buffer, position, cx) + }) + .ok() + { + task.await.log_err(); + } + }); + self.last_prediction = Some(LastPredictionState::Replaying { + prediction, + _task: task, + }); + } else { + self.last_prediction = Some(LastPredictionState::Failed("Buffer dropped".into())); + } + } + + cx.notify(); + } + + fn number_input( + label: &'static str, + window: &mut Window, + cx: &mut Context, + ) -> Entity { + let input = cx.new(|cx| { + SingleLineInput::new(window, cx, "") + .label(label) + .label_min_width(px(64.)) + }); + + cx.subscribe_in( + &input.read(cx).editor().clone(), + window, + |this, _, event, _window, cx| { + let EditorEvent::BufferEdited = event else { + return; + }; + + fn number_input_value( + input: &Entity, + cx: &App, + ) -> T { + input + .read(cx) + .editor() + .read(cx) + .text(cx) + .parse::() + .unwrap_or_default() + } + + let options = EditPredictionExcerptOptions { + max_bytes: number_input_value(&this.max_bytes_input, cx), + min_bytes: number_input_value(&this.min_bytes_input, cx), + target_before_cursor_over_total_bytes: number_input_value( + &this.cursor_context_ratio_input, + cx, + ), + }; + + this.set_options(options, cx); + }, + ) + .detach(); + input + } + + fn update_last_prediction( + &mut self, + prediction: zeta2::PredictionDebugInfo, + window: &mut Window, + cx: &mut Context, + ) { + let Some(worktree_id) = self + .project + .read(cx) + .worktrees(cx) + .next() + .map(|worktree| worktree.read(cx).id()) + else { + log::error!("Open a worktree to use edit prediction debug view"); + self.last_prediction.take(); + return; + }; + + self._update_state_task = cx.spawn_in(window, { + let language_registry = self.project.read(cx).languages().clone(); + async move |this, cx| { + let mut languages = HashMap::default(); + for lang_id in prediction + .context + .snippets + .iter() + .map(|snippet| snippet.declaration.identifier().language_id) + .chain(prediction.context.excerpt_text.language_id) + { + if let Entry::Vacant(entry) = languages.entry(lang_id) { + // Most snippets are gonna be the same language, + // so we think it's fine to do this sequentially for now + entry.insert(language_registry.language_for_id(lang_id).await.ok()); + } + } + + let markdown_language = language_registry + .language_for_name("Markdown") + .await + .log_err(); + + this.update_in(cx, |this, window, cx| { + let context_editor = cx.new(|cx| { + let multibuffer = cx.new(|cx| { + let mut multibuffer = MultiBuffer::new(language::Capability::ReadOnly); + let excerpt_file = Arc::new(ExcerptMetadataFile { + title: PathBuf::from("Cursor Excerpt").into(), + worktree_id, + }); + + let excerpt_buffer = cx.new(|cx| { + let mut buffer = + Buffer::local(prediction.context.excerpt_text.body, cx); + if let Some(language) = prediction + .context + .excerpt_text + .language_id + .as_ref() + .and_then(|id| languages.get(id)) + { + buffer.set_language(language.clone(), cx); + } + buffer.file_updated(excerpt_file, cx); + buffer + }); + + multibuffer.push_excerpts( + excerpt_buffer, + [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)], + cx, + ); + + for snippet in &prediction.context.snippets { + let path = this + .project + .read(cx) + .path_for_entry(snippet.declaration.project_entry_id(), cx); + + let snippet_file = Arc::new(ExcerptMetadataFile { + title: PathBuf::from(format!( + "{} (Score density: {})", + path.map(|p| p.path.to_string_lossy().to_string()) + .unwrap_or_else(|| "".to_string()), + snippet.score_density(SnippetStyle::Declaration) + )) + .into(), + worktree_id, + }); + + let excerpt_buffer = cx.new(|cx| { + let mut buffer = + Buffer::local(snippet.declaration.item_text().0, cx); + buffer.file_updated(snippet_file, cx); + if let Some(language) = + languages.get(&snippet.declaration.identifier().language_id) + { + buffer.set_language(language.clone(), cx); + } + buffer + }); + + multibuffer.push_excerpts( + excerpt_buffer, + [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)], + cx, + ); + } + + multibuffer + }); + + Editor::new(EditorMode::full(), multibuffer, None, window, cx) + }); + + let last_prediction = LastPrediction { + context_editor, + prompt_editor: cx.new(|cx| { + let buffer = cx.new(|cx| { + let mut buffer = Buffer::local(prediction.request.prompt, cx); + buffer.set_language(markdown_language.clone(), cx); + buffer + }); + let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx)); + let mut editor = + Editor::new(EditorMode::full(), buffer, None, window, cx); + editor.set_read_only(true); + editor.set_show_line_numbers(false, cx); + editor.set_show_gutter(false, cx); + editor.set_show_scrollbars(false, cx); + editor + }), + model_response_editor: cx.new(|cx| { + let buffer = cx.new(|cx| { + let mut buffer = + Buffer::local(prediction.request.model_response, cx); + buffer.set_language(markdown_language, cx); + buffer + }); + let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx)); + let mut editor = + Editor::new(EditorMode::full(), buffer, None, window, cx); + editor.set_read_only(true); + editor.set_show_line_numbers(false, cx); + editor.set_show_gutter(false, cx); + editor.set_show_scrollbars(false, cx); + editor + }), + retrieval_time: prediction.retrieval_time, + prompt_planning_time: prediction.request.prompt_planning_time, + inference_time: prediction.request.inference_time, + parsing_time: prediction.request.parsing_time, + buffer: prediction.buffer, + position: prediction.position, + }; + this.last_prediction = Some(LastPredictionState::Success(last_prediction)); + cx.notify(); + }) + .ok(); + } + }); + } + + fn render_duration(name: &'static str, time: chrono::TimeDelta) -> Div { + h_flex() + .gap_1() + .child(Label::new(name).color(Color::Muted).size(LabelSize::Small)) + .child( + Label::new(if time.num_microseconds().unwrap_or(0) > 1000 { + format!("{} ms", time.num_milliseconds()) + } else { + format!("{} µs", time.num_microseconds().unwrap_or(0)) + }) + .size(LabelSize::Small), + ) + } + + fn render_last_prediction(&self, prediction: &LastPrediction, cx: &mut Context) -> Div { + match &self.active_view { + ActiveView::Context => div().size_full().child(prediction.context_editor.clone()), + ActiveView::Inference => h_flex() + .items_start() + .w_full() + .flex_1() + .border_t_1() + .border_color(cx.theme().colors().border) + .bg(cx.theme().colors().editor_background) + .child( + v_flex() + .flex_1() + .gap_2() + .p_4() + .h_full() + .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall)) + .child(prediction.prompt_editor.clone()), + ) + .child(ui::vertical_divider()) + .child( + v_flex() + .flex_1() + .gap_2() + .h_full() + .p_4() + .child(ui::Headline::new("Model Response").size(ui::HeadlineSize::XSmall)) + .child(prediction.model_response_editor.clone()), + ), + } + } +} + +impl Focusable for Zeta2Inspector { + fn focus_handle(&self, _cx: &App) -> FocusHandle { + self.focus_handle.clone() + } +} + +impl Item for Zeta2Inspector { + type Event = (); + + fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString { + "Zeta2 Inspector".into() + } +} + +impl EventEmitter<()> for Zeta2Inspector {} + +impl Render for Zeta2Inspector { + fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { + let content = match self.last_prediction.as_ref() { + None => v_flex() + .size_full() + .justify_center() + .items_center() + .child(Label::new("No prediction").size(LabelSize::Large)) + .into_any(), + Some(LastPredictionState::Success(prediction)) => { + self.render_last_prediction(prediction, cx).into_any() + } + Some(LastPredictionState::Replaying { prediction, _task }) => self + .render_last_prediction(prediction, cx) + .opacity(0.6) + .into_any(), + Some(LastPredictionState::Failed(err)) => v_flex() + .p_4() + .gap_2() + .child(Label::new(err.clone()).buffer_font(cx)) + .into_any(), + }; + + v_flex() + .size_full() + .bg(cx.theme().colors().editor_background) + .child( + h_flex() + .w_full() + .child( + v_flex() + .flex_1() + .p_4() + .h_full() + .justify_between() + .child( + v_flex() + .gap_2() + .child( + Headline::new("Excerpt Options").size(HeadlineSize::Small), + ) + .child( + h_flex() + .gap_2() + .items_end() + .child(self.max_bytes_input.clone()) + .child(self.min_bytes_input.clone()) + .child(self.cursor_context_ratio_input.clone()) + .child( + ui::Button::new("reset-options", "Reset") + .disabled( + self.zeta.read(cx).excerpt_options() + == &zeta2::DEFAULT_EXCERPT_OPTIONS, + ) + .style(ButtonStyle::Outlined) + .size(ButtonSize::Large) + .on_click(cx.listener( + |this, _, window, cx| { + this.set_input_options( + &zeta2::DEFAULT_EXCERPT_OPTIONS, + window, + cx, + ); + }, + )), + ), + ), + ) + .map(|this| { + if let Some( + LastPredictionState::Success { .. } + | LastPredictionState::Replaying { .. }, + ) = self.last_prediction.as_ref() + { + this.child( + ui::ToggleButtonGroup::single_row( + "prediction", + [ + ui::ToggleButtonSimple::new( + "Context", + cx.listener(|this, _, _, cx| { + this.active_view = ActiveView::Context; + cx.notify(); + }), + ), + ui::ToggleButtonSimple::new( + "Inference", + cx.listener(|this, _, _, cx| { + this.active_view = ActiveView::Inference; + cx.notify(); + }), + ), + ], + ) + .style(ui::ToggleButtonGroupStyle::Outlined) + .selected_index( + if self.active_view == ActiveView::Context { + 0 + } else { + 1 + }, + ), + ) + } else { + this + } + }), + ) + .child(ui::vertical_divider()) + .map(|this| { + if let Some( + LastPredictionState::Success(prediction) + | LastPredictionState::Replaying { prediction, .. }, + ) = self.last_prediction.as_ref() + { + this.child( + v_flex() + .p_4() + .gap_2() + .min_w(px(160.)) + .child(Headline::new("Stats").size(HeadlineSize::Small)) + .child(Self::render_duration( + "Context retrieval", + prediction.retrieval_time, + )) + .child(Self::render_duration( + "Prompt planning", + prediction.prompt_planning_time, + )) + .child(Self::render_duration( + "Inference", + prediction.inference_time, + )) + .child(Self::render_duration( + "Parsing", + prediction.parsing_time, + )), + ) + } else { + this + } + }), + ) + .child(content) + } +} + +// Using same approach as commit view + +struct ExcerptMetadataFile { + title: Arc, + worktree_id: WorktreeId, +} + +impl language::File for ExcerptMetadataFile { + fn as_local(&self) -> Option<&dyn language::LocalFile> { + None + } + + fn disk_state(&self) -> DiskState { + DiskState::New + } + + fn path(&self) -> &Arc { + &self.title + } + + fn full_path(&self, _: &App) -> PathBuf { + self.title.as_ref().into() + } + + fn file_name<'a>(&'a self, _: &'a App) -> &'a OsStr { + self.title.file_name().unwrap() + } + + fn worktree_id(&self, _: &App) -> WorktreeId { + self.worktree_id + } + + fn to_proto(&self, _: &App) -> language::proto::File { + unimplemented!() + } + + fn is_private(&self) -> bool { + false + } +}