1use std::{collections::hash_map::Entry, path::PathBuf, str::FromStr, sync::Arc, time::Duration};
2
3use chrono::TimeDelta;
4use client::{Client, UserStore};
5use cloud_llm_client::predict_edits_v3::PromptFormat;
6use collections::HashMap;
7use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
8use futures::{StreamExt as _, channel::oneshot};
9use gpui::{
10 Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task, WeakEntity, actions,
11 prelude::*,
12};
13use language::{Buffer, DiskState};
14use project::{Project, WorktreeId};
15use ui::{ContextMenu, ContextMenuEntry, DropdownMenu, prelude::*};
16use ui_input::SingleLineInput;
17use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
18use workspace::{Item, SplitDirection, Workspace};
19use zeta2::{DEFAULT_CONTEXT_OPTIONS, PredictionDebugInfo, Zeta, ZetaOptions};
20
21use edit_prediction_context::{DeclarationStyle, EditPredictionExcerptOptions};
22
23actions!(
24 dev,
25 [
26 /// Opens the language server protocol logs viewer.
27 OpenZeta2Inspector
28 ]
29);
30
31pub fn init(cx: &mut App) {
32 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
33 workspace.register_action(move |workspace, _: &OpenZeta2Inspector, window, cx| {
34 let project = workspace.project();
35 workspace.split_item(
36 SplitDirection::Right,
37 Box::new(cx.new(|cx| {
38 Zeta2Inspector::new(
39 &project,
40 workspace.client(),
41 workspace.user_store(),
42 window,
43 cx,
44 )
45 })),
46 window,
47 cx,
48 );
49 });
50 })
51 .detach();
52}
53
54// TODO show included diagnostics, and events
55
56pub struct Zeta2Inspector {
57 focus_handle: FocusHandle,
58 project: Entity<Project>,
59 last_prediction: Option<LastPrediction>,
60 max_excerpt_bytes_input: Entity<SingleLineInput>,
61 min_excerpt_bytes_input: Entity<SingleLineInput>,
62 cursor_context_ratio_input: Entity<SingleLineInput>,
63 max_prompt_bytes_input: Entity<SingleLineInput>,
64 active_view: ActiveView,
65 zeta: Entity<Zeta>,
66 _active_editor_subscription: Option<Subscription>,
67 _update_state_task: Task<()>,
68 _receive_task: Task<()>,
69}
70
71#[derive(PartialEq)]
72enum ActiveView {
73 Context,
74 Inference,
75}
76
77struct LastPrediction {
78 context_editor: Entity<Editor>,
79 prompt_editor: Entity<Editor>,
80 retrieval_time: TimeDelta,
81 buffer: WeakEntity<Buffer>,
82 position: language::Anchor,
83 state: LastPredictionState,
84 _task: Option<Task<()>>,
85}
86
87enum LastPredictionState {
88 Requested,
89 Success {
90 inference_time: TimeDelta,
91 parsing_time: TimeDelta,
92 prompt_planning_time: TimeDelta,
93 model_response_editor: Entity<Editor>,
94 },
95 Failed {
96 message: String,
97 },
98}
99
100impl Zeta2Inspector {
101 pub fn new(
102 project: &Entity<Project>,
103 client: &Arc<Client>,
104 user_store: &Entity<UserStore>,
105 window: &mut Window,
106 cx: &mut Context<Self>,
107 ) -> Self {
108 let zeta = Zeta::global(client, user_store, cx);
109 let mut request_rx = zeta.update(cx, |zeta, _cx| zeta.debug_info());
110
111 let receive_task = cx.spawn_in(window, async move |this, cx| {
112 while let Some(prediction) = request_rx.next().await {
113 this.update_in(cx, |this, window, cx| {
114 this.update_last_prediction(prediction, window, cx)
115 })
116 .ok();
117 }
118 });
119
120 let mut this = Self {
121 focus_handle: cx.focus_handle(),
122 project: project.clone(),
123 last_prediction: None,
124 active_view: ActiveView::Context,
125 max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
126 min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
127 cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),
128 max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx),
129 zeta: zeta.clone(),
130 _active_editor_subscription: None,
131 _update_state_task: Task::ready(()),
132 _receive_task: receive_task,
133 };
134 this.set_input_options(&zeta.read(cx).options().clone(), window, cx);
135 this
136 }
137
138 fn set_input_options(
139 &mut self,
140 options: &ZetaOptions,
141 window: &mut Window,
142 cx: &mut Context<Self>,
143 ) {
144 self.max_excerpt_bytes_input.update(cx, |input, cx| {
145 input.set_text(options.context.excerpt.max_bytes.to_string(), window, cx);
146 });
147 self.min_excerpt_bytes_input.update(cx, |input, cx| {
148 input.set_text(options.context.excerpt.min_bytes.to_string(), window, cx);
149 });
150 self.cursor_context_ratio_input.update(cx, |input, cx| {
151 input.set_text(
152 format!(
153 "{:.2}",
154 options
155 .context
156 .excerpt
157 .target_before_cursor_over_total_bytes
158 ),
159 window,
160 cx,
161 );
162 });
163 self.max_prompt_bytes_input.update(cx, |input, cx| {
164 input.set_text(options.max_prompt_bytes.to_string(), window, cx);
165 });
166 cx.notify();
167 }
168
169 fn set_options(&mut self, options: ZetaOptions, cx: &mut Context<Self>) {
170 self.zeta.update(cx, |this, _cx| this.set_options(options));
171
172 const THROTTLE_TIME: Duration = Duration::from_millis(100);
173
174 if let Some(prediction) = self.last_prediction.as_mut() {
175 if let Some(buffer) = prediction.buffer.upgrade() {
176 let position = prediction.position;
177 let zeta = self.zeta.clone();
178 let project = self.project.clone();
179 prediction._task = Some(cx.spawn(async move |_this, cx| {
180 cx.background_executor().timer(THROTTLE_TIME).await;
181 if let Some(task) = zeta
182 .update(cx, |zeta, cx| {
183 zeta.refresh_prediction(&project, &buffer, position, cx)
184 })
185 .ok()
186 {
187 task.await.log_err();
188 }
189 }));
190 prediction.state = LastPredictionState::Requested;
191 } else {
192 self.last_prediction.take();
193 }
194 }
195
196 cx.notify();
197 }
198
199 fn number_input(
200 label: &'static str,
201 window: &mut Window,
202 cx: &mut Context<Self>,
203 ) -> Entity<SingleLineInput> {
204 let input = cx.new(|cx| {
205 SingleLineInput::new(window, cx, "")
206 .label(label)
207 .label_min_width(px(64.))
208 });
209
210 cx.subscribe_in(
211 &input.read(cx).editor().clone(),
212 window,
213 |this, _, event, _window, cx| {
214 let EditorEvent::BufferEdited = event else {
215 return;
216 };
217
218 fn number_input_value<T: FromStr + Default>(
219 input: &Entity<SingleLineInput>,
220 cx: &App,
221 ) -> T {
222 input
223 .read(cx)
224 .editor()
225 .read(cx)
226 .text(cx)
227 .parse::<T>()
228 .unwrap_or_default()
229 }
230
231 let mut context_options = DEFAULT_CONTEXT_OPTIONS.clone();
232 context_options.excerpt = EditPredictionExcerptOptions {
233 max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
234 min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
235 target_before_cursor_over_total_bytes: number_input_value(
236 &this.cursor_context_ratio_input,
237 cx,
238 ),
239 };
240
241 let zeta_options = this.zeta.read(cx).options();
242 this.set_options(
243 ZetaOptions {
244 context: context_options,
245 max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx),
246 max_diagnostic_bytes: zeta_options.max_diagnostic_bytes,
247 prompt_format: zeta_options.prompt_format,
248 file_indexing_parallelism: zeta_options.file_indexing_parallelism,
249 },
250 cx,
251 );
252 },
253 )
254 .detach();
255 input
256 }
257
258 fn update_last_prediction(
259 &mut self,
260 prediction: zeta2::PredictionDebugInfo,
261 window: &mut Window,
262 cx: &mut Context<Self>,
263 ) {
264 let project = self.project.read(cx);
265 let path_style = project.path_style(cx);
266 let Some(worktree_id) = project
267 .worktrees(cx)
268 .next()
269 .map(|worktree| worktree.read(cx).id())
270 else {
271 log::error!("Open a worktree to use edit prediction debug view");
272 self.last_prediction.take();
273 return;
274 };
275
276 self._update_state_task = cx.spawn_in(window, {
277 let language_registry = self.project.read(cx).languages().clone();
278 async move |this, cx| {
279 let mut languages = HashMap::default();
280 for lang_id in prediction
281 .context
282 .declarations
283 .iter()
284 .map(|snippet| snippet.declaration.identifier().language_id)
285 .chain(prediction.context.excerpt_text.language_id)
286 {
287 if let Entry::Vacant(entry) = languages.entry(lang_id) {
288 // Most snippets are gonna be the same language,
289 // so we think it's fine to do this sequentially for now
290 entry.insert(language_registry.language_for_id(lang_id).await.ok());
291 }
292 }
293
294 let markdown_language = language_registry
295 .language_for_name("Markdown")
296 .await
297 .log_err();
298
299 this.update_in(cx, |this, window, cx| {
300 let context_editor = cx.new(|cx| {
301 let multibuffer = cx.new(|cx| {
302 let mut multibuffer = MultiBuffer::new(language::Capability::ReadOnly);
303 let excerpt_file = Arc::new(ExcerptMetadataFile {
304 title: RelPath::unix("Cursor Excerpt").unwrap().into(),
305 path_style,
306 worktree_id,
307 });
308
309 let excerpt_buffer = cx.new(|cx| {
310 let mut buffer =
311 Buffer::local(prediction.context.excerpt_text.body, cx);
312 if let Some(language) = prediction
313 .context
314 .excerpt_text
315 .language_id
316 .as_ref()
317 .and_then(|id| languages.get(id))
318 {
319 buffer.set_language(language.clone(), cx);
320 }
321 buffer.file_updated(excerpt_file, cx);
322 buffer
323 });
324
325 multibuffer.push_excerpts(
326 excerpt_buffer,
327 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
328 cx,
329 );
330
331 for snippet in &prediction.context.declarations {
332 let path = this
333 .project
334 .read(cx)
335 .path_for_entry(snippet.declaration.project_entry_id(), cx);
336
337 let snippet_file = Arc::new(ExcerptMetadataFile {
338 title: RelPath::unix(&format!(
339 "{} (Score density: {})",
340 path.map(|p| p.path.display(path_style).to_string())
341 .unwrap_or_else(|| "".to_string()),
342 snippet.score_density(DeclarationStyle::Declaration)
343 ))
344 .unwrap()
345 .into(),
346 path_style,
347 worktree_id,
348 });
349
350 let excerpt_buffer = cx.new(|cx| {
351 let mut buffer =
352 Buffer::local(snippet.declaration.item_text().0, cx);
353 buffer.file_updated(snippet_file, cx);
354 if let Some(language) =
355 languages.get(&snippet.declaration.identifier().language_id)
356 {
357 buffer.set_language(language.clone(), cx);
358 }
359 buffer
360 });
361
362 multibuffer.push_excerpts(
363 excerpt_buffer,
364 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
365 cx,
366 );
367 }
368
369 multibuffer
370 });
371
372 Editor::new(EditorMode::full(), multibuffer, None, window, cx)
373 });
374
375 let PredictionDebugInfo {
376 response_rx,
377 position,
378 buffer,
379 retrieval_time,
380 local_prompt,
381 ..
382 } = prediction;
383
384 let task = cx.spawn_in(window, {
385 let markdown_language = markdown_language.clone();
386 async move |this, cx| {
387 let response = response_rx.await;
388
389 this.update_in(cx, |this, window, cx| {
390 if let Some(prediction) = this.last_prediction.as_mut() {
391 prediction.state = match response {
392 Ok(Ok(response)) => {
393 prediction.prompt_editor.update(
394 cx,
395 |prompt_editor, cx| {
396 prompt_editor.set_text(
397 response.prompt,
398 window,
399 cx,
400 );
401 },
402 );
403
404 LastPredictionState::Success {
405 prompt_planning_time: response.prompt_planning_time,
406 inference_time: response.inference_time,
407 parsing_time: response.parsing_time,
408 model_response_editor: cx.new(|cx| {
409 let buffer = cx.new(|cx| {
410 let mut buffer = Buffer::local(
411 response.model_response,
412 cx,
413 );
414 buffer.set_language(markdown_language, cx);
415 buffer
416 });
417 let buffer = cx.new(|cx| {
418 MultiBuffer::singleton(buffer, cx)
419 });
420 let mut editor = Editor::new(
421 EditorMode::full(),
422 buffer,
423 None,
424 window,
425 cx,
426 );
427 editor.set_read_only(true);
428 editor.set_show_line_numbers(false, cx);
429 editor.set_show_gutter(false, cx);
430 editor.set_show_scrollbars(false, cx);
431 editor
432 }),
433 }
434 }
435 Ok(Err(err)) => {
436 LastPredictionState::Failed { message: err }
437 }
438 Err(oneshot::Canceled) => LastPredictionState::Failed {
439 message: "Canceled".to_string(),
440 },
441 };
442 }
443 })
444 .ok();
445 }
446 });
447
448 this.last_prediction = Some(LastPrediction {
449 context_editor,
450 prompt_editor: cx.new(|cx| {
451 let buffer = cx.new(|cx| {
452 let mut buffer =
453 Buffer::local(local_prompt.unwrap_or_else(|err| err), cx);
454 buffer.set_language(markdown_language.clone(), cx);
455 buffer
456 });
457 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
458 let mut editor =
459 Editor::new(EditorMode::full(), buffer, None, window, cx);
460 editor.set_read_only(true);
461 editor.set_show_line_numbers(false, cx);
462 editor.set_show_gutter(false, cx);
463 editor.set_show_scrollbars(false, cx);
464 editor
465 }),
466 retrieval_time,
467 buffer,
468 position,
469 state: LastPredictionState::Requested,
470 _task: Some(task),
471 });
472 cx.notify();
473 })
474 .ok();
475 }
476 });
477 }
478
479 fn render_options(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
480 v_flex()
481 .gap_2()
482 .child(
483 h_flex()
484 .child(Headline::new("Options").size(HeadlineSize::Small))
485 .justify_between()
486 .child(
487 ui::Button::new("reset-options", "Reset")
488 .disabled(self.zeta.read(cx).options() == &zeta2::DEFAULT_OPTIONS)
489 .style(ButtonStyle::Outlined)
490 .size(ButtonSize::Large)
491 .on_click(cx.listener(|this, _, window, cx| {
492 this.set_input_options(&zeta2::DEFAULT_OPTIONS, window, cx);
493 })),
494 ),
495 )
496 .child(
497 v_flex()
498 .gap_2()
499 .child(
500 h_flex()
501 .gap_2()
502 .items_end()
503 .child(self.max_excerpt_bytes_input.clone())
504 .child(self.min_excerpt_bytes_input.clone())
505 .child(self.cursor_context_ratio_input.clone()),
506 )
507 .child(
508 h_flex()
509 .gap_2()
510 .items_end()
511 .child(self.max_prompt_bytes_input.clone())
512 .child(self.render_prompt_format_dropdown(window, cx)),
513 ),
514 )
515 }
516
517 fn render_prompt_format_dropdown(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
518 let active_format = self.zeta.read(cx).options().prompt_format;
519 let this = cx.weak_entity();
520
521 v_flex()
522 .gap_1p5()
523 .child(
524 Label::new("Prompt Format")
525 .size(LabelSize::Small)
526 .color(Color::Muted),
527 )
528 .child(
529 DropdownMenu::new(
530 "ep-prompt-format",
531 active_format.to_string(),
532 ContextMenu::build(window, cx, move |mut menu, _window, _cx| {
533 for prompt_format in PromptFormat::iter() {
534 menu = menu.item(
535 ContextMenuEntry::new(prompt_format.to_string())
536 .toggleable(IconPosition::End, active_format == prompt_format)
537 .handler({
538 let this = this.clone();
539 move |_window, cx| {
540 this.update(cx, |this, cx| {
541 let current_options =
542 this.zeta.read(cx).options().clone();
543 let options = ZetaOptions {
544 prompt_format,
545 ..current_options
546 };
547 this.set_options(options, cx);
548 })
549 .ok();
550 }
551 }),
552 )
553 }
554 menu
555 }),
556 )
557 .style(ui::DropdownStyle::Outlined),
558 )
559 }
560
561 fn render_tabs(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
562 if self.last_prediction.is_none() {
563 return None;
564 };
565
566 Some(
567 ui::ToggleButtonGroup::single_row(
568 "prediction",
569 [
570 ui::ToggleButtonSimple::new(
571 "Context",
572 cx.listener(|this, _, _, cx| {
573 this.active_view = ActiveView::Context;
574 cx.notify();
575 }),
576 ),
577 ui::ToggleButtonSimple::new(
578 "Inference",
579 cx.listener(|this, _, _, cx| {
580 this.active_view = ActiveView::Inference;
581 cx.notify();
582 }),
583 ),
584 ],
585 )
586 .style(ui::ToggleButtonGroupStyle::Outlined)
587 .selected_index(if self.active_view == ActiveView::Context {
588 0
589 } else {
590 1
591 })
592 .into_any_element(),
593 )
594 }
595
596 fn render_stats(&self) -> Option<Div> {
597 let Some(prediction) = self.last_prediction.as_ref() else {
598 return None;
599 };
600
601 let (prompt_planning_time, inference_time, parsing_time) = match &prediction.state {
602 LastPredictionState::Success {
603 inference_time,
604 parsing_time,
605 prompt_planning_time,
606 ..
607 } => (
608 Some(*prompt_planning_time),
609 Some(*inference_time),
610 Some(*parsing_time),
611 ),
612 LastPredictionState::Requested | LastPredictionState::Failed { .. } => {
613 (None, None, None)
614 }
615 };
616
617 Some(
618 v_flex()
619 .p_4()
620 .gap_2()
621 .min_w(px(160.))
622 .child(Headline::new("Stats").size(HeadlineSize::Small))
623 .child(Self::render_duration(
624 "Context retrieval",
625 Some(prediction.retrieval_time),
626 ))
627 .child(Self::render_duration(
628 "Prompt planning",
629 prompt_planning_time,
630 ))
631 .child(Self::render_duration("Inference", inference_time))
632 .child(Self::render_duration("Parsing", parsing_time)),
633 )
634 }
635
636 fn render_duration(name: &'static str, time: Option<chrono::TimeDelta>) -> Div {
637 h_flex()
638 .gap_1()
639 .child(Label::new(name).color(Color::Muted).size(LabelSize::Small))
640 .child(match time {
641 Some(time) => Label::new(if time.num_microseconds().unwrap_or(0) >= 1000 {
642 format!("{} ms", time.num_milliseconds())
643 } else {
644 format!("{} µs", time.num_microseconds().unwrap_or(0))
645 })
646 .size(LabelSize::Small),
647 None => Label::new("...").size(LabelSize::Small),
648 })
649 }
650
651 fn render_content(&self, cx: &mut Context<Self>) -> AnyElement {
652 match self.last_prediction.as_ref() {
653 None => v_flex()
654 .size_full()
655 .justify_center()
656 .items_center()
657 .child(Label::new("No prediction").size(LabelSize::Large))
658 .into_any(),
659 Some(prediction) => self.render_last_prediction(prediction, cx).into_any(),
660 }
661 }
662
663 fn render_last_prediction(&self, prediction: &LastPrediction, cx: &mut Context<Self>) -> Div {
664 match &self.active_view {
665 ActiveView::Context => div().size_full().child(prediction.context_editor.clone()),
666 ActiveView::Inference => h_flex()
667 .items_start()
668 .w_full()
669 .flex_1()
670 .border_t_1()
671 .border_color(cx.theme().colors().border)
672 .bg(cx.theme().colors().editor_background)
673 .child(
674 v_flex()
675 .flex_1()
676 .gap_2()
677 .p_4()
678 .h_full()
679 .child(
680 h_flex()
681 .justify_between()
682 .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall))
683 .child(match prediction.state {
684 LastPredictionState::Requested
685 | LastPredictionState::Failed { .. } => ui::Chip::new("Local")
686 .bg_color(cx.theme().status().warning_background)
687 .label_color(Color::Success),
688 LastPredictionState::Success { .. } => ui::Chip::new("Cloud")
689 .bg_color(cx.theme().status().success_background)
690 .label_color(Color::Success),
691 }),
692 )
693 .child(prediction.prompt_editor.clone()),
694 )
695 .child(ui::vertical_divider())
696 .child(
697 v_flex()
698 .flex_1()
699 .gap_2()
700 .h_full()
701 .p_4()
702 .child(ui::Headline::new("Model Response").size(ui::HeadlineSize::XSmall))
703 .child(match &prediction.state {
704 LastPredictionState::Success {
705 model_response_editor,
706 ..
707 } => model_response_editor.clone().into_any_element(),
708 LastPredictionState::Requested => v_flex()
709 .p_4()
710 .gap_2()
711 .child(Label::new("Loading...").buffer_font(cx))
712 .into_any(),
713 LastPredictionState::Failed { message } => v_flex()
714 .p_4()
715 .gap_2()
716 .child(Label::new(message.clone()).buffer_font(cx))
717 .into_any(),
718 }),
719 ),
720 }
721 }
722}
723
724impl Focusable for Zeta2Inspector {
725 fn focus_handle(&self, _cx: &App) -> FocusHandle {
726 self.focus_handle.clone()
727 }
728}
729
730impl Item for Zeta2Inspector {
731 type Event = ();
732
733 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
734 "Zeta2 Inspector".into()
735 }
736}
737
738impl EventEmitter<()> for Zeta2Inspector {}
739
740impl Render for Zeta2Inspector {
741 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
742 v_flex()
743 .size_full()
744 .bg(cx.theme().colors().editor_background)
745 .child(
746 h_flex()
747 .w_full()
748 .child(
749 v_flex()
750 .flex_1()
751 .p_4()
752 .h_full()
753 .justify_between()
754 .child(self.render_options(window, cx))
755 .gap_4()
756 .children(self.render_tabs(cx)),
757 )
758 .child(ui::vertical_divider())
759 .children(self.render_stats()),
760 )
761 .child(self.render_content(cx))
762 }
763}
764
765// Using same approach as commit view
766
767struct ExcerptMetadataFile {
768 title: Arc<RelPath>,
769 worktree_id: WorktreeId,
770 path_style: PathStyle,
771}
772
773impl language::File for ExcerptMetadataFile {
774 fn as_local(&self) -> Option<&dyn language::LocalFile> {
775 None
776 }
777
778 fn disk_state(&self) -> DiskState {
779 DiskState::New
780 }
781
782 fn path(&self) -> &Arc<RelPath> {
783 &self.title
784 }
785
786 fn full_path(&self, _: &App) -> PathBuf {
787 self.title.as_std_path().to_path_buf()
788 }
789
790 fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
791 self.title.file_name().unwrap()
792 }
793
794 fn path_style(&self, _: &App) -> PathStyle {
795 self.path_style
796 }
797
798 fn worktree_id(&self, _: &App) -> WorktreeId {
799 self.worktree_id
800 }
801
802 fn to_proto(&self, _: &App) -> language::proto::File {
803 unimplemented!()
804 }
805
806 fn is_private(&self) -> bool {
807 false
808 }
809}