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 39146342f2..c57d06ef88 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 3c431a29eb..8aabe6ad40 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 fee7498a69..f655387f68 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 aeda748112..f3752abab9 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 764c9040f2..58549d579d 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 3d7bbe8642..c2d8733b8a 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 3dbbed0ce5..af882e7623 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 61c560baab..d362441cdd 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 240d44fae4..4ecc0fa3ac 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 ffd34abb25..fe651c1419 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 0000000000..0d743e3880 --- /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 + } +}