1use std::{
2 cmp::Reverse, collections::hash_map::Entry, fmt::Debug, ops::Add as _, path::PathBuf,
3 str::FromStr, sync::Arc, time::Duration,
4};
5
6use client::{Client, UserStore};
7use cloud_llm_client::predict_edits_v3::{
8 self, DeclarationScoreComponents, PredictEditsResponse, PromptFormat,
9};
10use collections::{HashMap, HashSet};
11use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
12use feature_flags::FeatureFlagAppExt as _;
13use futures::{FutureExt, StreamExt as _, future::Shared};
14use gpui::{
15 CursorStyle, Empty, Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task,
16 WeakEntity, actions, prelude::*,
17};
18use language::{Buffer, DiskState};
19use ordered_float::OrderedFloat;
20use project::{Project, WorktreeId, telemetry_snapshot::TelemetrySnapshot};
21use ui::{ButtonLike, ContextMenu, ContextMenuEntry, DropdownMenu, KeyBinding, prelude::*};
22use ui_input::SingleLineInput;
23use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
24use workspace::{Item, SplitDirection, Workspace};
25use zeta2::{ObservedPredictionRequest, Zeta, Zeta2FeatureFlag, ZetaOptions};
26
27use edit_prediction_context::{EditPredictionContextOptions, EditPredictionExcerptOptions};
28
29actions!(
30 dev,
31 [
32 /// Opens the language server protocol logs viewer.
33 OpenZeta2Inspector,
34 /// Rate prediction as positive.
35 Zeta2RatePredictionPositive,
36 /// Rate prediction as negative.
37 Zeta2RatePredictionNegative,
38 /// Go to the previous request in the zeta2 inspector
39 Zeta2InspectorPrevious,
40 /// Go to the next request in the zeta2 inspector
41 Zeta2InspectorNext,
42 ]
43);
44
45pub fn init(cx: &mut App) {
46 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
47 workspace.register_action(move |workspace, _: &OpenZeta2Inspector, window, cx| {
48 let project = workspace.project();
49 workspace.split_item(
50 SplitDirection::Right,
51 Box::new(cx.new(|cx| {
52 Zeta2Inspector::new(
53 &project,
54 workspace.client(),
55 workspace.user_store(),
56 window,
57 cx,
58 )
59 })),
60 window,
61 cx,
62 );
63 });
64 })
65 .detach();
66}
67
68// TODO show included diagnostics, and events
69
70pub struct Zeta2Inspector {
71 focus_handle: FocusHandle,
72 project: Entity<Project>,
73 requests: Vec<InspectorObservedPredictionRequest>,
74 current: Option<CurrentRequest>,
75 max_excerpt_bytes_input: Entity<SingleLineInput>,
76 min_excerpt_bytes_input: Entity<SingleLineInput>,
77 cursor_context_ratio_input: Entity<SingleLineInput>,
78 max_prompt_bytes_input: Entity<SingleLineInput>,
79 max_retrieved_declarations: Entity<SingleLineInput>,
80 active_view: ActiveView,
81 zeta: Entity<Zeta>,
82 _active_editor_subscription: Option<Subscription>,
83 _update_state_task: Task<()>,
84 _receive_task: Task<()>,
85}
86
87struct InspectorObservedPredictionRequest {
88 observed_request: ObservedPredictionRequest,
89 project_snapshot: Shared<Task<Arc<TelemetrySnapshot>>>,
90}
91
92#[derive(PartialEq)]
93enum ActiveView {
94 Context,
95 Inference,
96}
97
98struct CurrentRequest {
99 index: usize,
100 context_editor: Entity<Editor>,
101 prompt_editor: Entity<Editor>,
102 buffer: WeakEntity<Buffer>,
103 position: language::Anchor,
104 state: CurrentPredictionState,
105 _task: Option<Task<()>>,
106}
107
108#[derive(Clone, Copy, PartialEq)]
109enum Feedback {
110 Positive,
111 Negative,
112}
113
114enum CurrentPredictionState {
115 Requested,
116 Success {
117 model_response_editor: Entity<Editor>,
118 feedback_editor: Entity<Editor>,
119 feedback: Option<Feedback>,
120 response: predict_edits_v3::PredictEditsResponse,
121 },
122 Failed {
123 message: String,
124 },
125}
126
127impl Zeta2Inspector {
128 pub fn new(
129 project: &Entity<Project>,
130 client: &Arc<Client>,
131 user_store: &Entity<UserStore>,
132 window: &mut Window,
133 cx: &mut Context<Self>,
134 ) -> Self {
135 let zeta = Zeta::global(client, user_store, cx);
136 let mut request_rx = zeta.update(cx, |zeta, _cx| zeta.observe_requests());
137
138 let receive_task = cx.spawn_in(window, async move |this, cx| {
139 while let Some(request) = request_rx.next().await {
140 this.update_in(cx, |this, window, cx| {
141 let project_snapshot_task = TelemetrySnapshot::new(&this.project, cx);
142 this.requests.push(InspectorObservedPredictionRequest {
143 observed_request: request,
144 project_snapshot: cx
145 .foreground_executor()
146 .spawn(async move { Arc::new(project_snapshot_task.await) })
147 .shared(),
148 });
149
150 this.set_current_predition(this.requests.len() - 1, window, cx)
151 })
152 .ok();
153 }
154 });
155
156 let mut this = Self {
157 focus_handle: cx.focus_handle(),
158 project: project.clone(),
159 requests: Vec::new(),
160 current: None,
161 active_view: ActiveView::Inference,
162 max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
163 min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
164 cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),
165 max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx),
166 max_retrieved_declarations: Self::number_input("Max Retrieved Definitions", window, cx),
167 zeta: zeta.clone(),
168 _active_editor_subscription: None,
169 _update_state_task: Task::ready(()),
170 _receive_task: receive_task,
171 };
172 this.set_input_options(&zeta.read(cx).options().clone(), window, cx);
173 this
174 }
175
176 fn set_input_options(
177 &mut self,
178 options: &ZetaOptions,
179 window: &mut Window,
180 cx: &mut Context<Self>,
181 ) {
182 self.max_excerpt_bytes_input.update(cx, |input, cx| {
183 input.set_text(options.context.excerpt.max_bytes.to_string(), window, cx);
184 });
185 self.min_excerpt_bytes_input.update(cx, |input, cx| {
186 input.set_text(options.context.excerpt.min_bytes.to_string(), window, cx);
187 });
188 self.cursor_context_ratio_input.update(cx, |input, cx| {
189 input.set_text(
190 format!(
191 "{:.2}",
192 options
193 .context
194 .excerpt
195 .target_before_cursor_over_total_bytes
196 ),
197 window,
198 cx,
199 );
200 });
201 self.max_prompt_bytes_input.update(cx, |input, cx| {
202 input.set_text(options.max_prompt_bytes.to_string(), window, cx);
203 });
204 self.max_retrieved_declarations.update(cx, |input, cx| {
205 input.set_text(
206 options.context.max_retrieved_declarations.to_string(),
207 window,
208 cx,
209 );
210 });
211 cx.notify();
212 }
213
214 fn set_options(&mut self, options: ZetaOptions, cx: &mut Context<Self>) {
215 self.zeta.update(cx, |this, _cx| this.set_options(options));
216
217 const THROTTLE_TIME: Duration = Duration::from_millis(100);
218
219 if let Some(prediction) = self.current.as_mut() {
220 if let Some(buffer) = prediction.buffer.upgrade() {
221 let position = prediction.position;
222 let zeta = self.zeta.clone();
223 let project = self.project.clone();
224 prediction._task = Some(cx.spawn(async move |_this, cx| {
225 cx.background_executor().timer(THROTTLE_TIME).await;
226 if let Some(task) = zeta
227 .update(cx, |zeta, cx| {
228 zeta.refresh_prediction(&project, &buffer, position, cx)
229 })
230 .ok()
231 {
232 task.await.log_err();
233 }
234 }));
235 prediction.state = CurrentPredictionState::Requested;
236 } else {
237 self.current.take();
238 }
239 }
240
241 cx.notify();
242 }
243
244 fn number_input(
245 label: &'static str,
246 window: &mut Window,
247 cx: &mut Context<Self>,
248 ) -> Entity<SingleLineInput> {
249 let input = cx.new(|cx| {
250 SingleLineInput::new(window, cx, "")
251 .label(label)
252 .label_min_width(px(64.))
253 });
254
255 cx.subscribe_in(
256 &input.read(cx).editor().clone(),
257 window,
258 |this, _, event, _window, cx| {
259 let EditorEvent::BufferEdited = event else {
260 return;
261 };
262
263 fn number_input_value<T: FromStr + Default>(
264 input: &Entity<SingleLineInput>,
265 cx: &App,
266 ) -> T {
267 input
268 .read(cx)
269 .editor()
270 .read(cx)
271 .text(cx)
272 .parse::<T>()
273 .unwrap_or_default()
274 }
275
276 let zeta_options = this.zeta.read(cx).options().clone();
277
278 let context_options = EditPredictionContextOptions {
279 excerpt: EditPredictionExcerptOptions {
280 max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
281 min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
282 target_before_cursor_over_total_bytes: number_input_value(
283 &this.cursor_context_ratio_input,
284 cx,
285 ),
286 },
287 max_retrieved_declarations: number_input_value(
288 &this.max_retrieved_declarations,
289 cx,
290 ),
291 ..zeta_options.context
292 };
293
294 this.set_options(
295 ZetaOptions {
296 context: context_options,
297 max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx),
298 max_diagnostic_bytes: zeta_options.max_diagnostic_bytes,
299 prompt_format: zeta_options.prompt_format,
300 file_indexing_parallelism: zeta_options.file_indexing_parallelism,
301 },
302 cx,
303 );
304 },
305 )
306 .detach();
307 input
308 }
309
310 fn set_current_predition(&mut self, index: usize, window: &mut Window, cx: &mut Context<Self>) {
311 let project = self.project.read(cx);
312 let path_style = project.path_style(cx);
313 let Some(worktree_id) = project
314 .worktrees(cx)
315 .next()
316 .map(|worktree| worktree.read(cx).id())
317 else {
318 log::error!("Open a worktree to use edit prediction debug view");
319 self.current.take();
320 return;
321 };
322
323 self._update_state_task = cx.spawn_in(window, {
324 let language_registry = self.project.read(cx).languages().clone();
325 async move |this, cx| {
326 let mut languages = HashMap::default();
327
328 let file_extensions = this
329 .read_with(cx, |this, _| {
330 let mut file_extensions = HashSet::default();
331 let request = &this.requests[index];
332
333 for ext in request
334 .observed_request
335 .full_request
336 .referenced_declarations
337 .iter()
338 .filter_map(|snippet| snippet.path.extension())
339 .chain(
340 request
341 .observed_request
342 .full_request
343 .excerpt_path
344 .extension(),
345 )
346 {
347 if !file_extensions.contains(ext) {
348 file_extensions.insert(ext.to_owned());
349 }
350 }
351 file_extensions
352 })
353 .unwrap();
354
355 for ext in file_extensions {
356 if let Entry::Vacant(entry) = languages.entry(ext) {
357 // Most snippets are gonna be the same language,
358 // so we think it's fine to do this sequentially for now
359 let ext_str = entry.key().to_string_lossy().to_string();
360 entry.insert(
361 language_registry
362 .language_for_name_or_extension(&ext_str)
363 .await
364 .ok(),
365 );
366 }
367 }
368
369 let markdown_language = language_registry
370 .language_for_name("Markdown")
371 .await
372 .log_err();
373
374 this.update_in(cx, |this, window, cx| {
375 let InspectorObservedPredictionRequest {
376 observed_request: request,
377 ..
378 } = &this.requests[index];
379
380 let context_editor = cx.new(|cx| {
381 let mut excerpt_score_components = HashMap::default();
382
383 let multibuffer = cx.new(|cx| {
384 let mut multibuffer = MultiBuffer::new(language::Capability::ReadOnly);
385 let excerpt_file = Arc::new(ExcerptMetadataFile {
386 title: RelPath::unix("Cursor Excerpt").unwrap().into(),
387 path_style,
388 worktree_id,
389 });
390
391 let excerpt_buffer = cx.new(|cx| {
392 let mut buffer =
393 Buffer::local(request.full_request.excerpt.clone(), cx);
394 if let Some(language) = request
395 .full_request
396 .excerpt_path
397 .extension()
398 .and_then(|ext| languages.get(ext))
399 {
400 buffer.set_language(language.clone(), cx);
401 }
402 buffer.file_updated(excerpt_file, cx);
403 #[cfg(debug_assertions)]
404 buffer.debug(
405 &language::Point::new(
406 request.full_request.cursor_point.line.0,
407 request.full_request.cursor_point.column,
408 ),
409 CursorMarker,
410 );
411 buffer
412 });
413
414 multibuffer.push_excerpts(
415 excerpt_buffer,
416 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
417 cx,
418 );
419
420 let mut declarations =
421 request.full_request.referenced_declarations.clone();
422 declarations.sort_unstable_by_key(|declaration| {
423 Reverse(OrderedFloat(declaration.declaration_score))
424 });
425
426 for snippet in &declarations {
427 let snippet_file = Arc::new(ExcerptMetadataFile {
428 title: RelPath::unix(&format!(
429 "{} (Score: {})",
430 snippet.path.display(),
431 snippet.declaration_score
432 ))
433 .unwrap()
434 .into(),
435 path_style,
436 worktree_id,
437 });
438
439 let excerpt_buffer = cx.new(|cx| {
440 let mut buffer = Buffer::local(snippet.text.clone(), cx);
441 buffer.file_updated(snippet_file, cx);
442 if let Some(ext) = snippet.path.extension()
443 && let Some(language) = languages.get(ext)
444 {
445 buffer.set_language(language.clone(), cx);
446 }
447 buffer
448 });
449
450 let excerpt_ids = multibuffer.push_excerpts(
451 excerpt_buffer,
452 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
453 cx,
454 );
455 let excerpt_id = excerpt_ids.first().unwrap();
456
457 excerpt_score_components
458 .insert(*excerpt_id, snippet.score_components.clone());
459 }
460
461 multibuffer
462 });
463
464 let mut editor =
465 Editor::new(EditorMode::full(), multibuffer, None, window, cx);
466 editor.register_addon(ZetaContextAddon {
467 excerpt_score_components,
468 });
469 editor
470 });
471
472 let ObservedPredictionRequest {
473 response,
474 position,
475 buffer,
476 local_prompt,
477 ..
478 } = request;
479
480 let response_task = response.clone();
481 let task = cx.spawn_in(window, {
482 let markdown_language = markdown_language.clone();
483 async move |this, cx| {
484 let response = response_task.await;
485 this.update_in(cx, |this, window, cx| {
486 if let Some(prediction) = this.current.as_mut() {
487 prediction.state = match response {
488 Ok(response) => {
489 if let Some(debug_info) = &response.debug_info {
490 prediction.prompt_editor.update(
491 cx,
492 |prompt_editor, cx| {
493 prompt_editor.set_text(
494 debug_info.prompt.as_str(),
495 window,
496 cx,
497 );
498 },
499 );
500 }
501
502 let feedback_editor = cx.new(|cx| {
503 let buffer = cx.new(|cx| {
504 let mut buffer = Buffer::local("", cx);
505 buffer.set_language(
506 markdown_language.clone(),
507 cx,
508 );
509 buffer
510 });
511 let buffer =
512 cx.new(|cx| MultiBuffer::singleton(buffer, cx));
513 let mut editor = Editor::new(
514 EditorMode::AutoHeight {
515 min_lines: 3,
516 max_lines: None,
517 },
518 buffer,
519 None,
520 window,
521 cx,
522 );
523 editor.set_placeholder_text(
524 "Write feedback here",
525 window,
526 cx,
527 );
528 editor.set_show_line_numbers(false, cx);
529 editor.set_show_gutter(false, cx);
530 editor.set_show_scrollbars(false, cx);
531 editor
532 });
533
534 cx.subscribe_in(
535 &feedback_editor,
536 window,
537 |this, editor, ev, window, cx| match ev {
538 EditorEvent::BufferEdited => {
539 if let Some(last_prediction) =
540 this.current.as_mut()
541 && let CurrentPredictionState::Success {
542 feedback: feedback_state,
543 ..
544 } = &mut last_prediction.state
545 {
546 if feedback_state.take().is_some() {
547 editor.update(cx, |editor, cx| {
548 editor.set_placeholder_text(
549 "Write feedback here",
550 window,
551 cx,
552 );
553 });
554 cx.notify();
555 }
556 }
557 }
558 _ => {}
559 },
560 )
561 .detach();
562
563 CurrentPredictionState::Success {
564 model_response_editor: cx.new(|cx| {
565 let buffer = cx.new(|cx| {
566 let mut buffer = Buffer::local(
567 response
568 .debug_info
569 .as_ref()
570 .map(|p| p.model_response.as_str())
571 .unwrap_or(
572 "(Debug info not available)",
573 ),
574 cx,
575 );
576 buffer.set_language(markdown_language, cx);
577 buffer
578 });
579 let buffer = cx.new(|cx| {
580 MultiBuffer::singleton(buffer, cx)
581 });
582 let mut editor = Editor::new(
583 EditorMode::full(),
584 buffer,
585 None,
586 window,
587 cx,
588 );
589 editor.set_read_only(true);
590 editor.set_show_line_numbers(false, cx);
591 editor.set_show_gutter(false, cx);
592 editor.set_show_scrollbars(false, cx);
593 editor
594 }),
595 feedback_editor,
596 feedback: None,
597 response,
598 }
599 }
600 Err(err) => CurrentPredictionState::Failed { message: err },
601 };
602 }
603 })
604 .ok();
605 }
606 });
607
608 this.current = Some(CurrentRequest {
609 index,
610 context_editor,
611 prompt_editor: cx.new(|cx| {
612 let buffer = cx.new(|cx| {
613 let mut buffer = Buffer::local(
614 local_prompt.as_ref().unwrap_or_else(|err| err),
615 cx,
616 );
617 buffer.set_language(markdown_language.clone(), cx);
618 buffer
619 });
620 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
621 let mut editor =
622 Editor::new(EditorMode::full(), buffer, None, window, cx);
623 editor.set_read_only(true);
624 editor.set_show_line_numbers(false, cx);
625 editor.set_show_gutter(false, cx);
626 editor.set_show_scrollbars(false, cx);
627 editor
628 }),
629 buffer: buffer.clone(),
630 position: *position,
631 state: CurrentPredictionState::Requested,
632 _task: Some(task),
633 });
634 cx.notify();
635 })
636 .ok();
637 }
638 });
639 }
640
641 fn handle_prev(
642 &mut self,
643 _action: &Zeta2InspectorPrevious,
644 window: &mut Window,
645 cx: &mut Context<Self>,
646 ) {
647 self.set_current_predition(
648 self.current
649 .as_ref()
650 .map(|c| c.index)
651 .unwrap_or_default()
652 .saturating_sub(1),
653 window,
654 cx,
655 );
656 }
657
658 fn handle_next(
659 &mut self,
660 _action: &Zeta2InspectorPrevious,
661 window: &mut Window,
662 cx: &mut Context<Self>,
663 ) {
664 self.set_current_predition(
665 self.current
666 .as_ref()
667 .map(|c| c.index)
668 .unwrap_or_default()
669 .add(1)
670 .min(self.requests.len() - 1),
671 window,
672 cx,
673 );
674 }
675
676 fn handle_rate_positive(
677 &mut self,
678 _action: &Zeta2RatePredictionPositive,
679 window: &mut Window,
680 cx: &mut Context<Self>,
681 ) {
682 self.handle_rate(Feedback::Positive, window, cx);
683 }
684
685 fn handle_rate_negative(
686 &mut self,
687 _action: &Zeta2RatePredictionNegative,
688 window: &mut Window,
689 cx: &mut Context<Self>,
690 ) {
691 self.handle_rate(Feedback::Negative, window, cx);
692 }
693
694 fn handle_rate(&mut self, kind: Feedback, window: &mut Window, cx: &mut Context<Self>) {
695 let Some(last_prediction) = self.current.as_mut() else {
696 return;
697 };
698 let request = &self.requests[last_prediction.index];
699 if !request.observed_request.full_request.can_collect_data {
700 return;
701 }
702
703 let project_snapshot_task = request.project_snapshot.clone();
704
705 cx.spawn_in(window, async move |this, cx| {
706 let project_snapshot = project_snapshot_task.await;
707 this.update_in(cx, |this, window, cx| {
708 let Some(last_prediction) = this.current.as_mut() else {
709 return;
710 };
711
712 // todo! move to Self::requests?
713 let CurrentPredictionState::Success {
714 feedback: feedback_state,
715 feedback_editor,
716 model_response_editor,
717 response,
718 ..
719 } = &mut last_prediction.state
720 else {
721 return;
722 };
723
724 *feedback_state = Some(kind);
725 let text = feedback_editor.update(cx, |feedback_editor, cx| {
726 feedback_editor.set_placeholder_text(
727 "Submitted. Edit or submit again to change.",
728 window,
729 cx,
730 );
731 feedback_editor.text(cx)
732 });
733 cx.notify();
734
735 cx.defer_in(window, {
736 let model_response_editor = model_response_editor.downgrade();
737 move |_, window, cx| {
738 if let Some(model_response_editor) = model_response_editor.upgrade() {
739 model_response_editor.focus_handle(cx).focus(window);
740 }
741 }
742 });
743
744 let kind = match kind {
745 Feedback::Positive => "positive",
746 Feedback::Negative => "negative",
747 };
748
749 let request = &this.requests[last_prediction.index];
750
751 telemetry::event!(
752 "Zeta2 Prediction Rated",
753 id = response.request_id,
754 kind = kind,
755 text = text,
756 request = request.observed_request.full_request,
757 response = response,
758 project_snapshot = project_snapshot,
759 );
760 })
761 .log_err();
762 })
763 .detach();
764 }
765
766 fn focus_feedback(&mut self, window: &mut Window, cx: &mut Context<Self>) {
767 if let Some(last_prediction) = self.current.as_mut() {
768 if let CurrentPredictionState::Success {
769 feedback_editor, ..
770 } = &mut last_prediction.state
771 {
772 feedback_editor.focus_handle(cx).focus(window);
773 }
774 };
775 }
776
777 fn render_options(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
778 v_flex()
779 .gap_2()
780 .child(
781 h_flex()
782 .child(Headline::new("Options").size(HeadlineSize::Small))
783 .justify_between()
784 .child(
785 ui::Button::new("reset-options", "Reset")
786 .disabled(self.zeta.read(cx).options() == &zeta2::DEFAULT_OPTIONS)
787 .style(ButtonStyle::Outlined)
788 .size(ButtonSize::Large)
789 .on_click(cx.listener(|this, _, window, cx| {
790 this.set_input_options(&zeta2::DEFAULT_OPTIONS, window, cx);
791 })),
792 ),
793 )
794 .child(
795 v_flex()
796 .gap_2()
797 .child(
798 h_flex()
799 .gap_2()
800 .items_end()
801 .child(self.max_excerpt_bytes_input.clone())
802 .child(self.min_excerpt_bytes_input.clone())
803 .child(self.cursor_context_ratio_input.clone()),
804 )
805 .child(
806 h_flex()
807 .gap_2()
808 .items_end()
809 .child(self.max_retrieved_declarations.clone())
810 .child(self.max_prompt_bytes_input.clone())
811 .child(self.render_prompt_format_dropdown(window, cx)),
812 ),
813 )
814 }
815
816 fn render_prompt_format_dropdown(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
817 let active_format = self.zeta.read(cx).options().prompt_format;
818 let this = cx.weak_entity();
819
820 v_flex()
821 .gap_1p5()
822 .child(
823 Label::new("Prompt Format")
824 .size(LabelSize::Small)
825 .color(Color::Muted),
826 )
827 .child(
828 DropdownMenu::new(
829 "ep-prompt-format",
830 active_format.to_string(),
831 ContextMenu::build(window, cx, move |mut menu, _window, _cx| {
832 for prompt_format in PromptFormat::iter() {
833 menu = menu.item(
834 ContextMenuEntry::new(prompt_format.to_string())
835 .toggleable(IconPosition::End, active_format == prompt_format)
836 .handler({
837 let this = this.clone();
838 move |_window, cx| {
839 this.update(cx, |this, cx| {
840 let current_options =
841 this.zeta.read(cx).options().clone();
842 let options = ZetaOptions {
843 prompt_format,
844 ..current_options
845 };
846 this.set_options(options, cx);
847 })
848 .ok();
849 }
850 }),
851 )
852 }
853 menu
854 }),
855 )
856 .style(ui::DropdownStyle::Outlined),
857 )
858 }
859
860 fn render_tabs(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
861 if self.current.is_none() {
862 return None;
863 };
864
865 Some(
866 ui::ToggleButtonGroup::single_row(
867 "prediction",
868 [
869 ui::ToggleButtonSimple::new(
870 "Context",
871 cx.listener(|this, _, _, cx| {
872 this.active_view = ActiveView::Context;
873 cx.notify();
874 }),
875 ),
876 ui::ToggleButtonSimple::new(
877 "Inference",
878 cx.listener(|this, _, window, cx| {
879 this.active_view = ActiveView::Inference;
880 this.focus_feedback(window, cx);
881 cx.notify();
882 }),
883 ),
884 ],
885 )
886 .style(ui::ToggleButtonGroupStyle::Outlined)
887 .selected_index(if self.active_view == ActiveView::Context {
888 0
889 } else {
890 1
891 })
892 .into_any_element(),
893 )
894 }
895
896 fn render_stats(&self) -> Option<Div> {
897 let Some(current) = self.current.as_ref() else {
898 return None;
899 };
900
901 let (prompt_planning_time, inference_time, parsing_time) =
902 if let CurrentPredictionState::Success {
903 response:
904 PredictEditsResponse {
905 debug_info: Some(debug_info),
906 ..
907 },
908 ..
909 } = ¤t.state
910 {
911 (
912 Some(debug_info.prompt_planning_time),
913 Some(debug_info.inference_time),
914 Some(debug_info.parsing_time),
915 )
916 } else {
917 (None, None, None)
918 };
919
920 Some(
921 v_flex()
922 .p_4()
923 .gap_2()
924 .min_w(px(160.))
925 .child(Headline::new("Stats").size(HeadlineSize::Small))
926 .child(Self::render_duration(
927 "Context retrieval",
928 Some(self.requests[current.index].observed_request.retrieval_time),
929 ))
930 .child(Self::render_duration(
931 "Prompt planning",
932 prompt_planning_time,
933 ))
934 .child(Self::render_duration("Inference", inference_time))
935 .child(Self::render_duration("Parsing", parsing_time)),
936 )
937 }
938
939 fn render_duration(name: &'static str, time: Option<chrono::TimeDelta>) -> Div {
940 h_flex()
941 .gap_1()
942 .child(Label::new(name).color(Color::Muted).size(LabelSize::Small))
943 .child(match time {
944 Some(time) => Label::new(if time.num_microseconds().unwrap_or(0) >= 1000 {
945 format!("{} ms", time.num_milliseconds())
946 } else {
947 format!("{} µs", time.num_microseconds().unwrap_or(0))
948 })
949 .size(LabelSize::Small),
950 None => Label::new("...").size(LabelSize::Small),
951 })
952 }
953
954 fn render_content(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
955 if !cx.has_flag::<Zeta2FeatureFlag>() {
956 return Self::render_message("`zeta2` feature flag is not enabled");
957 }
958
959 match self.current.as_ref() {
960 None => Self::render_message("No prediction"),
961 Some(prediction) => self
962 .render_current_prediction(prediction, window, cx)
963 .into_any(),
964 }
965 }
966
967 fn render_message(message: impl Into<SharedString>) -> AnyElement {
968 v_flex()
969 .size_full()
970 .justify_center()
971 .items_center()
972 .child(Label::new(message).size(LabelSize::Large))
973 .into_any()
974 }
975
976 fn render_current_prediction(
977 &self,
978 current: &CurrentRequest,
979 window: &mut Window,
980 cx: &mut Context<Self>,
981 ) -> Div {
982 match &self.active_view {
983 ActiveView::Context => div().size_full().child(current.context_editor.clone()),
984 ActiveView::Inference => h_flex()
985 .items_start()
986 .w_full()
987 .flex_1()
988 .border_t_1()
989 .border_color(cx.theme().colors().border)
990 .bg(cx.theme().colors().editor_background)
991 .child(
992 v_flex()
993 .flex_1()
994 .gap_2()
995 .p_4()
996 .h_full()
997 .child(
998 h_flex()
999 .justify_between()
1000 .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall))
1001 .child(match current.state {
1002 CurrentPredictionState::Requested
1003 | CurrentPredictionState::Failed { .. } => {
1004 ui::Chip::new("Local")
1005 .bg_color(cx.theme().status().warning_background)
1006 .label_color(Color::Success)
1007 }
1008 CurrentPredictionState::Success { .. } => {
1009 ui::Chip::new("Cloud")
1010 .bg_color(cx.theme().status().success_background)
1011 .label_color(Color::Success)
1012 }
1013 }),
1014 )
1015 .child(current.prompt_editor.clone()),
1016 )
1017 .child(ui::vertical_divider())
1018 .child(
1019 v_flex()
1020 .flex_1()
1021 .gap_2()
1022 .h_full()
1023 .child(
1024 v_flex()
1025 .flex_1()
1026 .gap_2()
1027 .p_4()
1028 .child(
1029 ui::Headline::new("Model Response")
1030 .size(ui::HeadlineSize::XSmall),
1031 )
1032 .child(match ¤t.state {
1033 CurrentPredictionState::Success {
1034 model_response_editor,
1035 ..
1036 } => model_response_editor.clone().into_any_element(),
1037 CurrentPredictionState::Requested => v_flex()
1038 .gap_2()
1039 .child(Label::new("Loading...").buffer_font(cx))
1040 .into_any_element(),
1041 CurrentPredictionState::Failed { message } => v_flex()
1042 .gap_2()
1043 .max_w_96()
1044 .child(Label::new(message.clone()).buffer_font(cx))
1045 .into_any_element(),
1046 }),
1047 )
1048 .child(ui::divider())
1049 .child(
1050 if self.requests[current.index]
1051 .observed_request
1052 .full_request
1053 .can_collect_data
1054 && let CurrentPredictionState::Success {
1055 feedback_editor,
1056 feedback: feedback_state,
1057 ..
1058 } = ¤t.state
1059 {
1060 v_flex()
1061 .key_context("Zeta2Feedback")
1062 .on_action(cx.listener(Self::handle_rate_positive))
1063 .on_action(cx.listener(Self::handle_rate_negative))
1064 .gap_2()
1065 .p_2()
1066 .child(feedback_editor.clone())
1067 .child(
1068 h_flex()
1069 .justify_end()
1070 .w_full()
1071 .child(
1072 ButtonLike::new("rate-positive")
1073 .when(
1074 *feedback_state == Some(Feedback::Positive),
1075 |this| this.style(ButtonStyle::Filled),
1076 )
1077 .children(
1078 KeyBinding::for_action(
1079 &Zeta2RatePredictionPositive,
1080 window,
1081 cx,
1082 )
1083 .map(|k| k.size(TextSize::Small.rems(cx))),
1084 )
1085 .child(ui::Icon::new(ui::IconName::ThumbsUp))
1086 .on_click(cx.listener(
1087 |this, _, window, cx| {
1088 this.handle_rate_positive(
1089 &Zeta2RatePredictionPositive,
1090 window,
1091 cx,
1092 );
1093 },
1094 )),
1095 )
1096 .child(
1097 ButtonLike::new("rate-negative")
1098 .when(
1099 *feedback_state == Some(Feedback::Negative),
1100 |this| this.style(ButtonStyle::Filled),
1101 )
1102 .children(
1103 KeyBinding::for_action(
1104 &Zeta2RatePredictionNegative,
1105 window,
1106 cx,
1107 )
1108 .map(|k| k.size(TextSize::Small.rems(cx))),
1109 )
1110 .child(ui::Icon::new(ui::IconName::ThumbsDown))
1111 .on_click(cx.listener(
1112 |this, _, window, cx| {
1113 this.handle_rate_negative(
1114 &Zeta2RatePredictionNegative,
1115 window,
1116 cx,
1117 );
1118 },
1119 )),
1120 ),
1121 )
1122 .into_any()
1123 } else {
1124 Empty.into_any_element()
1125 },
1126 ),
1127 ),
1128 }
1129 }
1130}
1131
1132impl Focusable for Zeta2Inspector {
1133 fn focus_handle(&self, _cx: &App) -> FocusHandle {
1134 self.focus_handle.clone()
1135 }
1136}
1137
1138impl Item for Zeta2Inspector {
1139 type Event = ();
1140
1141 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
1142 "Zeta2 Inspector".into()
1143 }
1144}
1145
1146impl EventEmitter<()> for Zeta2Inspector {}
1147
1148impl Render for Zeta2Inspector {
1149 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1150 v_flex()
1151 .track_focus(&self.focus_handle)
1152 .key_context("Zeta2Inspector")
1153 .size_full()
1154 .on_action(cx.listener(Self::handle_prev))
1155 .on_action(cx.listener(Self::handle_next))
1156 .bg(cx.theme().colors().editor_background)
1157 .child(
1158 h_flex()
1159 .w_full()
1160 .child(
1161 v_flex()
1162 .flex_1()
1163 .p_4()
1164 .h_full()
1165 .justify_between()
1166 .child(self.render_options(window, cx))
1167 .gap_4()
1168 .children(self.render_tabs(cx)),
1169 )
1170 .child(ui::vertical_divider())
1171 .children(self.render_stats()),
1172 )
1173 .child(self.render_content(window, cx))
1174 }
1175}
1176
1177// Using same approach as commit view
1178
1179struct ExcerptMetadataFile {
1180 title: Arc<RelPath>,
1181 worktree_id: WorktreeId,
1182 path_style: PathStyle,
1183}
1184
1185impl language::File for ExcerptMetadataFile {
1186 fn as_local(&self) -> Option<&dyn language::LocalFile> {
1187 None
1188 }
1189
1190 fn disk_state(&self) -> DiskState {
1191 DiskState::New
1192 }
1193
1194 fn path(&self) -> &Arc<RelPath> {
1195 &self.title
1196 }
1197
1198 fn full_path(&self, _: &App) -> PathBuf {
1199 self.title.as_std_path().to_path_buf()
1200 }
1201
1202 fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
1203 self.title.file_name().unwrap()
1204 }
1205
1206 fn path_style(&self, _: &App) -> PathStyle {
1207 self.path_style
1208 }
1209
1210 fn worktree_id(&self, _: &App) -> WorktreeId {
1211 self.worktree_id
1212 }
1213
1214 fn to_proto(&self, _: &App) -> language::proto::File {
1215 unimplemented!()
1216 }
1217
1218 fn is_private(&self) -> bool {
1219 false
1220 }
1221}
1222
1223#[cfg(debug_assertions)]
1224struct CursorMarker;
1225
1226#[cfg(debug_assertions)]
1227impl Debug for CursorMarker {
1228 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1229 write!(f, "Cursor")
1230 }
1231}
1232
1233struct ZetaContextAddon {
1234 excerpt_score_components: HashMap<editor::ExcerptId, DeclarationScoreComponents>,
1235}
1236
1237impl editor::Addon for ZetaContextAddon {
1238 fn to_any(&self) -> &dyn std::any::Any {
1239 self
1240 }
1241
1242 fn render_buffer_header_controls(
1243 &self,
1244 excerpt_info: &multi_buffer::ExcerptInfo,
1245 _window: &Window,
1246 _cx: &App,
1247 ) -> Option<AnyElement> {
1248 let score_components = self.excerpt_score_components.get(&excerpt_info.id)?.clone();
1249
1250 Some(
1251 div()
1252 .id(excerpt_info.id.to_proto() as usize)
1253 .child(ui::Icon::new(IconName::Info))
1254 .cursor(CursorStyle::PointingHand)
1255 .tooltip(move |_, cx| {
1256 cx.new(|_| ScoreComponentsTooltip::new(&score_components))
1257 .into()
1258 })
1259 .into_any(),
1260 )
1261 }
1262}
1263
1264struct ScoreComponentsTooltip {
1265 text: SharedString,
1266}
1267
1268impl ScoreComponentsTooltip {
1269 fn new(components: &DeclarationScoreComponents) -> Self {
1270 Self {
1271 text: format!("{:#?}", components).into(),
1272 }
1273 }
1274}
1275
1276impl Render for ScoreComponentsTooltip {
1277 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1278 div().pl_2().pt_2p5().child(
1279 div()
1280 .elevation_2(cx)
1281 .py_1()
1282 .px_2()
1283 .child(ui::Label::new(self.text.clone()).buffer_font(cx)),
1284 )
1285 }
1286}