From bf3c5705e750fee4f591c2a153a2ba552bab33b6 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Mon, 22 Sep 2025 18:47:03 -0300 Subject: [PATCH] Checkpoint: Displaying debug info Co-Authored-By: Bennet --- Cargo.lock | 62 +- Cargo.toml | 4 +- crates/edit_prediction_context/Cargo.toml | 1 + .../src/edit_prediction_context.rs | 2 +- crates/edit_prediction_context/src/excerpt.rs | 6 +- .../src/edit_prediction_tools.rs | 461 ------------- crates/zed/Cargo.toml | 2 +- crates/zed/src/main.rs | 2 +- crates/zeta2/Cargo.toml | 3 +- crates/zeta2/src/zeta2.rs | 160 +++-- .../Cargo.toml | 16 +- .../LICENSE-GPL | 0 crates/zeta2_tools/src/zeta2_tools.rs | 605 ++++++++++++++++++ 13 files changed, 763 insertions(+), 561 deletions(-) delete mode 100644 crates/edit_prediction_tools/src/edit_prediction_tools.rs 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 8e0ff733011d90d8672766508fed39cfae1598e5..52bb51ca6f138429e440327edaa07666c20b1e25 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5174,6 +5174,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrayvec", + "chrono", "clap", "cloud_llm_client", "collections", @@ -5199,33 +5200,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" @@ -21246,7 +21220,6 @@ dependencies = [ "debugger_ui", "diagnostics", "edit_prediction_button", - "edit_prediction_tools", "editor", "env_logger 0.11.8", "extension", @@ -21358,6 +21331,7 @@ dependencies = [ "zed_env_vars", "zeta", "zeta2", + "zeta2_tools", "zlog", "zlog_settings", ] @@ -21641,6 +21615,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrayvec", + "chrono", "client", "cloud_llm_client", "edit_prediction", @@ -21661,6 +21636,37 @@ 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", + "markdown", + "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 c03d45eab01b64518142cdc828e9dbf22ed071d6..0d34f0ba38396c6929f35a0473904a887f02c76b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,7 +58,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", @@ -316,7 +316,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/Cargo.toml b/crates/edit_prediction_context/Cargo.toml index 75880cad5f3e2807e525908656931853efa19a92..063ba1714387edf510cc1df93409ddf9371bde50 100644 --- a/crates/edit_prediction_context/Cargo.toml +++ b/crates/edit_prediction_context/Cargo.toml @@ -14,6 +14,7 @@ path = "src/edit_prediction_context.rs" [dependencies] anyhow.workspace = true arrayvec.workspace = true +chrono.workspace = true cloud_llm_client.workspace = true collections.workspace = true futures.workspace = true diff --git a/crates/edit_prediction_context/src/edit_prediction_context.rs b/crates/edit_prediction_context/src/edit_prediction_context.rs index ccd21365ba31a29e8a96405615ab92abdda699b2..fca064a13eeb72a08f2da123b6abaa5c205b4b72 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..780d35502aefd1cb5617b5d175749555a1ab61f9 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}; @@ -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/edit_prediction_tools/src/edit_prediction_tools.rs b/crates/edit_prediction_tools/src/edit_prediction_tools.rs deleted file mode 100644 index 381cb2704a42c66ac945b1db896d38ea41f04281..0000000000000000000000000000000000000000 --- a/crates/edit_prediction_tools/src/edit_prediction_tools.rs +++ /dev/null @@ -1,461 +0,0 @@ -use std::{ - collections::hash_map::Entry, - ffi::OsStr, - path::{Path, PathBuf}, - str::FromStr, - sync::Arc, - time::{Duration, Instant}, -}; - -use collections::HashMap; -use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer}; -use gpui::{ - Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task, WeakEntity, actions, - prelude::*, -}; -use language::{Buffer, DiskState}; -use project::{Project, WorktreeId}; -use text::ToPoint; -use ui::prelude::*; -use ui_input::SingleLineInput; -use workspace::{Item, SplitDirection, Workspace}; - -use edit_prediction_context::{ - EditPredictionContext, EditPredictionExcerptOptions, SnippetStyle, SyntaxIndex, -}; - -actions!( - dev, - [ - /// Opens the language server protocol logs viewer. - OpenEditPredictionContext - ] -); - -pub fn init(cx: &mut App) { - cx.observe_new(move |workspace: &mut Workspace, _, _cx| { - workspace.register_action( - move |workspace, _: &OpenEditPredictionContext, window, cx| { - let workspace_entity = cx.entity(); - let project = workspace.project(); - let active_editor = workspace.active_item_as::(cx); - workspace.split_item( - SplitDirection::Right, - Box::new(cx.new(|cx| { - EditPredictionTools::new( - &workspace_entity, - &project, - active_editor, - window, - cx, - ) - })), - window, - cx, - ); - }, - ); - }) - .detach(); -} - -pub struct EditPredictionTools { - focus_handle: FocusHandle, - project: Entity, - last_context: Option, - max_bytes_input: Entity, - min_bytes_input: Entity, - cursor_context_ratio_input: Entity, - // TODO move to project or provider? - syntax_index: Entity, - last_editor: WeakEntity, - _active_editor_subscription: Option, - _edit_prediction_context_task: Task<()>, -} - -struct ContextState { - context_editor: Entity, - retrieval_duration: Duration, -} - -impl EditPredictionTools { - pub fn new( - workspace: &Entity, - project: &Entity, - active_editor: Option>, - window: &mut Window, - cx: &mut Context, - ) -> Self { - cx.subscribe_in(workspace, window, |this, workspace, event, window, cx| { - if let workspace::Event::ActiveItemChanged = event { - if let Some(editor) = workspace.read(cx).active_item_as::(cx) { - this._active_editor_subscription = Some(cx.subscribe_in( - &editor, - window, - |this, editor, event, window, cx| { - if let EditorEvent::SelectionsChanged { .. } = event { - this.update_context(editor, window, cx); - } - }, - )); - this.update_context(&editor, window, cx); - } else { - this._active_editor_subscription = None; - } - } - }) - .detach(); - let syntax_index = cx.new(|cx| SyntaxIndex::new(project, cx)); - - let number_input = |label: &'static str, - value: &'static str, - window: &mut Window, - cx: &mut Context| - -> Entity { - let input = cx.new(|cx| { - let input = SingleLineInput::new(window, cx, "") - .label(label) - .label_min_width(px(64.)); - input.set_text(value, window, cx); - input - }); - cx.subscribe_in( - &input.read(cx).editor().clone(), - window, - |this, _, event, window, cx| { - if let EditorEvent::BufferEdited = event - && let Some(editor) = this.last_editor.upgrade() - { - this.update_context(&editor, window, cx); - } - }, - ) - .detach(); - input - }; - - let mut this = Self { - focus_handle: cx.focus_handle(), - project: project.clone(), - last_context: None, - max_bytes_input: number_input("Max Bytes", "512", window, cx), - min_bytes_input: number_input("Min Bytes", "128", window, cx), - cursor_context_ratio_input: number_input("Cursor Context Ratio", "0.5", window, cx), - syntax_index, - last_editor: WeakEntity::new_invalid(), - _active_editor_subscription: None, - _edit_prediction_context_task: Task::ready(()), - }; - - if let Some(editor) = active_editor { - this.update_context(&editor, window, cx); - } - - this - } - - fn update_context( - &mut self, - editor: &Entity, - window: &mut Window, - cx: &mut Context, - ) { - self.last_editor = editor.downgrade(); - - let editor = editor.read(cx); - let buffer = editor.buffer().clone(); - let cursor_position = editor.selections.newest_anchor().start; - - let Some(buffer) = buffer.read(cx).buffer_for_anchor(cursor_position, cx) else { - self.last_context.take(); - return; - }; - let current_buffer_snapshot = buffer.read(cx).snapshot(); - let cursor_position = cursor_position - .text_anchor - .to_point(¤t_buffer_snapshot); - - let language = current_buffer_snapshot.language().cloned(); - 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_context.take(); - return; - }; - - self._edit_prediction_context_task = cx.spawn_in(window, { - let language_registry = self.project.read(cx).languages().clone(); - async move |this, cx| { - cx.background_executor() - .timer(Duration::from_millis(50)) - .await; - - let mut start_time = None; - - let Ok(task) = this.update(cx, |this, cx| { - 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, - ), - }; - - start_time = Some(Instant::now()); - - // TODO! use global zeta instead - EditPredictionContext::gather_context_in_background( - cursor_position, - current_buffer_snapshot, - options, - Some(this.syntax_index.clone()), - cx, - ) - }) else { - this.update(cx, |this, _cx| { - this.last_context.take(); - }) - .ok(); - return; - }; - - let Some(context) = task.await else { - // TODO: Display message - this.update(cx, |this, _cx| { - this.last_context.take(); - }) - .ok(); - return; - }; - let retrieval_duration = start_time.unwrap().elapsed(); - - let mut languages = HashMap::default(); - for snippet in context.snippets.iter() { - let lang_id = snippet.declaration.identifier().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()); - } - } - - 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(context.excerpt_text.body, cx); - buffer.set_language(language, 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 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) - }); - - this.last_context = Some(ContextState { - context_editor, - retrieval_duration, - }); - cx.notify(); - }) - .ok(); - } - }); - } -} - -impl Focusable for EditPredictionTools { - fn focus_handle(&self, _cx: &App) -> FocusHandle { - self.focus_handle.clone() - } -} - -impl Item for EditPredictionTools { - type Event = (); - - fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString { - "Edit Prediction Context Debug View".into() - } - - fn tab_icon(&self, _window: &Window, _cx: &App) -> Option { - Some(Icon::new(IconName::ZedPredict)) - } -} - -impl EventEmitter<()> for EditPredictionTools {} - -impl Render for EditPredictionTools { - fn render(&mut self, _window: &mut Window, cx: &mut Context) -> impl IntoElement { - v_flex() - .size_full() - .bg(cx.theme().colors().editor_background) - .child( - h_flex() - .items_start() - .w_full() - .child( - v_flex() - .flex_1() - .p_4() - .gap_2() - .child(Headline::new("Excerpt Options").size(HeadlineSize::Small)) - .child( - h_flex() - .gap_2() - .child(self.max_bytes_input.clone()) - .child(self.min_bytes_input.clone()) - .child(self.cursor_context_ratio_input.clone()), - ), - ) - .child(ui::Divider::vertical()) - .when_some(self.last_context.as_ref(), |this, last_context| { - this.child( - v_flex() - .p_4() - .gap_2() - .min_w(px(160.)) - .child(Headline::new("Stats").size(HeadlineSize::Small)) - .child( - h_flex() - .gap_1() - .child( - Label::new("Time to retrieve") - .color(Color::Muted) - .size(LabelSize::Small), - ) - .child( - Label::new( - if last_context.retrieval_duration.as_micros() - > 1000 - { - format!( - "{} ms", - last_context.retrieval_duration.as_millis() - ) - } else { - format!( - "{} µs", - last_context.retrieval_duration.as_micros() - ) - }, - ) - .size(LabelSize::Small), - ), - ), - ) - }), - ) - .children(self.last_context.as_ref().map(|c| c.context_editor.clone())) - } -} - -// 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 - } -} 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 5bcb750f4c4a73eac489a22ef3ca8d2198589172..fdd468944fb8f52494fefa48c5e2b8a6a7081bc0 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,6 +12,7 @@ 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, @@ -23,13 +25,12 @@ use project::Project; use release_channel::AppVersion; use std::cmp; use std::collections::{HashMap, VecDeque, hash_map}; -use std::fmt::Write; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; 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}; @@ -51,8 +52,17 @@ pub struct Zeta { projects: HashMap, excerpt_options: EditPredictionExcerptOptions, update_required: bool, + debug_tx: Option>>, } +pub struct PredictionDebugInfo { + pub context: EditPredictionContext, + pub retrieval_time: TimeDelta, + pub request: RequestDebugInfo, +} + +pub type RequestDebugInfo = predict_edits_v3::DebugInfo; + struct ZetaProject { syntax_index: Entity, events: VecDeque, @@ -75,41 +85,41 @@ pub enum Event { impl Event { //TODO: Actually use the events this in the prompt - fn to_prompt(&self) -> String { - match self { - Event::BufferChange { - old_snapshot, - new_snapshot, - .. - } => { - let mut prompt = String::new(); - - let old_path = old_snapshot - .file() - .map(|f| f.path().as_ref()) - .unwrap_or(Path::new("untitled")); - let new_path = new_snapshot - .file() - .map(|f| f.path().as_ref()) - .unwrap_or(Path::new("untitled")); - if old_path != new_path { - writeln!(prompt, "User renamed {:?} to {:?}\n", old_path, new_path).unwrap(); - } - - let diff = language::unified_diff(&old_snapshot.text(), &new_snapshot.text()); - if !diff.is_empty() { - write!( - prompt, - "User edited {:?}:\n```diff\n{}\n```", - new_path, diff - ) - .unwrap(); - } - - prompt - } - } - } + // fn to_prompt(&self) -> String { + // match self { + // Event::BufferChange { + // old_snapshot, + // new_snapshot, + // .. + // } => { + // let mut prompt = String::new(); + + // let old_path = old_snapshot + // .file() + // .map(|f| f.path().as_ref()) + // .unwrap_or(Path::new("untitled")); + // let new_path = new_snapshot + // .file() + // .map(|f| f.path().as_ref()) + // .unwrap_or(Path::new("untitled")); + // if old_path != new_path { + // writeln!(prompt, "User renamed {:?} to {:?}\n", old_path, new_path).unwrap(); + // } + + // let diff = language::unified_diff(&old_snapshot.text(), &new_snapshot.text()); + // if !diff.is_empty() { + // write!( + // prompt, + // "User edited {:?}:\n```diff\n{}\n```", + // new_path, diff + // ) + // .unwrap(); + // } + + // prompt + // } + // } + // } } impl Zeta { @@ -153,9 +163,16 @@ 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 usage(&self, cx: &App) -> Option { self.user_store.read(cx).edit_prediction_usage() } @@ -313,6 +330,7 @@ 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(); @@ -325,35 +343,59 @@ 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, + }) + }, + )) + .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..5298cbb95207a45a1513ebee6b031c7d60201bec 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,15 +9,19 @@ 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 +markdown.workspace = true project.workspace = true serde.workspace = true text.workspace = true @@ -25,17 +29,17 @@ ui.workspace = true ui_input.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..72753aa9f652ff2b6e97e97feccdcfdce47a7173 --- /dev/null +++ b/crates/zeta2_tools/src/zeta2_tools.rs @@ -0,0 +1,605 @@ +use std::{ + collections::hash_map::Entry, + ffi::OsStr, + path::{Path, PathBuf}, + sync::Arc, +}; + +use chrono::TimeDelta; +use client::{Client, UserStore}; +use collections::HashMap; +use editor::{Editor, EditorMode, ExcerptRange, MultiBuffer}; +use futures::StreamExt as _; +use gpui::{ + BorderStyle, EdgesRefinement, Entity, EventEmitter, FocusHandle, Focusable, Length, + StyleRefinement, Subscription, Task, TextStyleRefinement, UnderlineStyle, actions, prelude::*, +}; +use language::{Buffer, DiskState}; +use markdown::{HeadingLevelStyles, Markdown, MarkdownElement, MarkdownStyle}; +use project::{Project, WorktreeId}; +use ui::prelude::*; +use ui_input::SingleLineInput; +use workspace::{Item, SplitDirection, Workspace}; +use zeta2::Zeta; + +use edit_prediction_context::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| { + EditPredictionTools::new( + &project, + workspace.client(), + workspace.user_store(), + window, + cx, + ) + })), + window, + cx, + ); + }); + }) + .detach(); +} + +pub struct EditPredictionTools { + focus_handle: FocusHandle, + project: Entity, + last_prediction: Option>, + max_bytes_input: Entity, + min_bytes_input: Entity, + cursor_context_ratio_input: Entity, + active_view: ActiveView, + _active_editor_subscription: Option, + _update_state_task: Task<()>, + _receive_task: Task<()>, +} + +#[derive(PartialEq)] +enum ActiveView { + Context, + Inference, +} + +struct LastPredictionState { + context_editor: Entity, + retrieval_time: TimeDelta, + prompt_planning_time: TimeDelta, + inference_time: TimeDelta, + parsing_time: TimeDelta, + prompt_md: Entity, + model_response_md: Entity, +} + +impl EditPredictionTools { + pub fn new( + project: &Entity, + client: &Arc, + user_store: &Entity, + window: &mut Window, + cx: &mut Context, + ) -> Self { + let number_input = |label: &'static str, + value: &'static str, + window: &mut Window, + cx: &mut Context| + -> Entity { + let input = cx.new(|cx| { + let input = SingleLineInput::new(window, cx, "") + .label(label) + .label_min_width(px(64.)); + input.set_text(value, window, cx); + input + }); + // todo! + // cx.subscribe_in( + // &input.read(cx).editor().clone(), + // window, + // |this, _, event, window, cx| { + // if let EditorEvent::BufferEdited = event + // && let Some(editor) = this.last_editor.upgrade() + // { + // this.update_context(&editor, window, cx); + // } + // }, + // ) + // .detach(); + input + }; + + 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(Err(err.into())); + cx.notify(); + } + }) + .ok(); + } + }); + + Self { + focus_handle: cx.focus_handle(), + project: project.clone(), + last_prediction: None, + active_view: ActiveView::Context, + max_bytes_input: number_input("Max Bytes", "512", window, cx), + min_bytes_input: number_input("Min Bytes", "128", window, cx), + cursor_context_ratio_input: number_input("Cursor Context Ratio", "0.5", window, cx), + _active_editor_subscription: None, + _update_state_task: Task::ready(()), + _receive_task: receive_task, + } + } + + 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| { + // 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, + // ), + // }; + + 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()); + } + } + + 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) + }); + + this.last_prediction = Some(Ok(LastPredictionState { + context_editor, + prompt_md: cx.new(|cx| { + Markdown::new(prediction.request.prompt.into(), None, None, cx) + }), + model_response_md: cx.new(|cx| { + Markdown::new(prediction.request.model_response.into(), None, None, cx) + }), + 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, + })); + 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), + ) + } +} + +impl Focusable for EditPredictionTools { + fn focus_handle(&self, _cx: &App) -> FocusHandle { + self.focus_handle.clone() + } +} + +impl Item for EditPredictionTools { + type Event = (); + + fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString { + "Zeta2 Inspector".into() + } +} + +impl EventEmitter<()> for EditPredictionTools {} + +impl Render for EditPredictionTools { + fn render(&mut self, window: &mut Window, cx: &mut Context) -> impl IntoElement { + v_flex() + .size_full() + .bg(cx.theme().colors().editor_background) + .child( + h_flex() + .items_start() + .w_full() + .child( + v_flex() + .flex_1() + .p_4() + .gap_2() + .child(Headline::new("Excerpt Options").size(HeadlineSize::Small)) + .child( + h_flex() + .gap_2() + .child(self.max_bytes_input.clone()) + .child(self.min_bytes_input.clone()) + .child(self.cursor_context_ratio_input.clone()), + ) + .child(div().flex_1()) + .when( + self.last_prediction.as_ref().is_some_and(|r| r.is_ok()), + |this| { + this.child( + ui::ToggleButtonGroup::single_row( + "prediction", + [ + ui::ToggleButtonSimple::new( + "Context", + cx.listener(|this, _, _, cx| { + this.active_view = ActiveView::Context; + cx.notify(); + }), + ) + .selected(self.active_view == ActiveView::Context), + ui::ToggleButtonSimple::new( + "Inference", + cx.listener(|this, _, _, cx| { + this.active_view = ActiveView::Inference; + cx.notify(); + }), + ) + .selected(self.active_view == ActiveView::Context), + ], + ) + .style(ui::ToggleButtonGroupStyle::Outlined), + ) + }, + ), + ) + .child(ui::vertical_divider()) + .when_some( + self.last_prediction.as_ref().and_then(|r| r.as_ref().ok()), + |this, last_prediction| { + 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", + last_prediction.retrieval_time, + )) + .child(Self::render_duration( + "Prompt planning", + last_prediction.prompt_planning_time, + )) + .child(Self::render_duration( + "Inference", + last_prediction.inference_time, + )) + .child(Self::render_duration( + "Parsing", + last_prediction.parsing_time, + )), + ) + }, + ), + ) + .children(self.last_prediction.as_ref().map(|result| { + match result { + Ok(state) => match &self.active_view { + ActiveView::Context => state.context_editor.clone().into_any_element(), + ActiveView::Inference => h_flex() + .items_start() + .w_full() + .gap_2() + .bg(cx.theme().colors().editor_background) + // todo! fix layout + .child( + v_flex() + .flex_1() + .p_4() + .gap_2() + .child( + ui::Headline::new("Prompt").size(ui::HeadlineSize::Small), + ) + .child(MarkdownElement::new( + state.prompt_md.clone(), + markdown_style(window, cx), + )), + ) + .child(ui::vertical_divider()) + .child( + v_flex() + .flex_1() + .p_4() + .gap_2() + .child( + ui::Headline::new("Model Response") + .size(ui::HeadlineSize::Small), + ) + .child(MarkdownElement::new( + state.model_response_md.clone(), + markdown_style(window, cx), + )), + ) + .into_any(), + }, + Err(err) => v_flex() + .p_4() + .gap_2() + .child(Label::new(err.clone()).buffer_font(cx)) + .into_any(), + } + })) + } +} + +// Mostly copied from agent-ui +fn markdown_style(window: &Window, cx: &App) -> MarkdownStyle { + let colors = cx.theme().colors(); + + let buffer_font_size = TextSize::Small.rems(cx); + let mut text_style = window.text_style(); + let line_height = buffer_font_size * 1.75; + + let font_size = TextSize::Small.rems(cx); + + let text_color = colors.text; + + text_style.refine(&TextStyleRefinement { + font_size: Some(font_size.into()), + line_height: Some(line_height.into()), + color: Some(text_color), + ..Default::default() + }); + + MarkdownStyle { + base_text_style: text_style.clone(), + syntax: cx.theme().syntax().clone(), + selection_background_color: colors.element_selection_background, + code_block_overflow_x_scroll: true, + table_overflow_x_scroll: true, + heading_level_styles: Some(HeadingLevelStyles { + h1: Some(TextStyleRefinement { + font_size: Some(rems(1.15).into()), + ..Default::default() + }), + h2: Some(TextStyleRefinement { + font_size: Some(rems(1.1).into()), + ..Default::default() + }), + h3: Some(TextStyleRefinement { + font_size: Some(rems(1.05).into()), + ..Default::default() + }), + h4: Some(TextStyleRefinement { + font_size: Some(rems(1.).into()), + ..Default::default() + }), + h5: Some(TextStyleRefinement { + font_size: Some(rems(0.95).into()), + ..Default::default() + }), + h6: Some(TextStyleRefinement { + font_size: Some(rems(0.875).into()), + ..Default::default() + }), + }), + code_block: StyleRefinement { + padding: EdgesRefinement { + top: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))), + left: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))), + right: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))), + bottom: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))), + }, + margin: EdgesRefinement { + top: Some(Length::Definite(Pixels(8.).into())), + left: Some(Length::Definite(Pixels(0.).into())), + right: Some(Length::Definite(Pixels(0.).into())), + bottom: Some(Length::Definite(Pixels(12.).into())), + }, + border_style: Some(BorderStyle::Solid), + border_widths: EdgesRefinement { + top: Some(AbsoluteLength::Pixels(Pixels(1.))), + left: Some(AbsoluteLength::Pixels(Pixels(1.))), + right: Some(AbsoluteLength::Pixels(Pixels(1.))), + bottom: Some(AbsoluteLength::Pixels(Pixels(1.))), + }, + border_color: Some(colors.border_variant), + background: Some(colors.editor_background.into()), + text: Some(TextStyleRefinement { + font_size: Some(buffer_font_size.into()), + ..Default::default() + }), + ..Default::default() + }, + inline_code: TextStyleRefinement { + font_size: Some(buffer_font_size.into()), + background_color: Some(colors.editor_foreground.opacity(0.08)), + ..Default::default() + }, + link: TextStyleRefinement { + background_color: Some(colors.editor_foreground.opacity(0.025)), + underline: Some(UnderlineStyle { + color: Some(colors.text_accent.opacity(0.5)), + thickness: px(1.), + ..Default::default() + }), + ..Default::default() + }, + ..Default::default() + } +} + +// 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 + } +}