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 + } +}