1use std::{
2 cmp::Reverse, collections::hash_map::Entry, ops::Add as _, path::PathBuf, str::FromStr,
3 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 buffer
404 });
405
406 multibuffer.push_excerpts(
407 excerpt_buffer,
408 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
409 cx,
410 );
411
412 let mut declarations =
413 request.full_request.referenced_declarations.clone();
414 declarations.sort_unstable_by_key(|declaration| {
415 Reverse(OrderedFloat(declaration.declaration_score))
416 });
417
418 for snippet in &declarations {
419 let snippet_file = Arc::new(ExcerptMetadataFile {
420 title: RelPath::unix(&format!(
421 "{} (Score: {})",
422 snippet.path.display(),
423 snippet.declaration_score
424 ))
425 .unwrap()
426 .into(),
427 path_style,
428 worktree_id,
429 });
430
431 let excerpt_buffer = cx.new(|cx| {
432 let mut buffer = Buffer::local(snippet.text.clone(), cx);
433 buffer.file_updated(snippet_file, cx);
434 if let Some(ext) = snippet.path.extension()
435 && let Some(language) = languages.get(ext)
436 {
437 buffer.set_language(language.clone(), cx);
438 }
439 buffer
440 });
441
442 let excerpt_ids = multibuffer.push_excerpts(
443 excerpt_buffer,
444 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
445 cx,
446 );
447 let excerpt_id = excerpt_ids.first().unwrap();
448
449 excerpt_score_components
450 .insert(*excerpt_id, snippet.score_components.clone());
451 }
452
453 multibuffer
454 });
455
456 let mut editor =
457 Editor::new(EditorMode::full(), multibuffer, None, window, cx);
458 editor.register_addon(ZetaContextAddon {
459 excerpt_score_components,
460 });
461 editor
462 });
463
464 let ObservedPredictionRequest {
465 response,
466 position,
467 buffer,
468 local_prompt,
469 ..
470 } = request;
471
472 let response_task = response.clone();
473 let task = cx.spawn_in(window, {
474 let markdown_language = markdown_language.clone();
475 async move |this, cx| {
476 let response = response_task.await;
477 this.update_in(cx, |this, window, cx| {
478 if let Some(prediction) = this.current.as_mut() {
479 prediction.state = match response {
480 Ok(response) => {
481 if let Some(debug_info) = &response.debug_info {
482 prediction.prompt_editor.update(
483 cx,
484 |prompt_editor, cx| {
485 prompt_editor.set_text(
486 debug_info.prompt.as_str(),
487 window,
488 cx,
489 );
490 },
491 );
492 }
493
494 let feedback_editor = cx.new(|cx| {
495 let buffer = cx.new(|cx| {
496 let mut buffer = Buffer::local("", cx);
497 buffer.set_language(
498 markdown_language.clone(),
499 cx,
500 );
501 buffer
502 });
503 let buffer =
504 cx.new(|cx| MultiBuffer::singleton(buffer, cx));
505 let mut editor = Editor::new(
506 EditorMode::AutoHeight {
507 min_lines: 3,
508 max_lines: None,
509 },
510 buffer,
511 None,
512 window,
513 cx,
514 );
515 editor.set_placeholder_text(
516 "Write feedback here",
517 window,
518 cx,
519 );
520 editor.set_show_line_numbers(false, cx);
521 editor.set_show_gutter(false, cx);
522 editor.set_show_scrollbars(false, cx);
523 editor
524 });
525
526 cx.subscribe_in(
527 &feedback_editor,
528 window,
529 |this, editor, ev, window, cx| match ev {
530 EditorEvent::BufferEdited => {
531 if let Some(last_prediction) =
532 this.current.as_mut()
533 && let CurrentPredictionState::Success {
534 feedback: feedback_state,
535 ..
536 } = &mut last_prediction.state
537 {
538 if feedback_state.take().is_some() {
539 editor.update(cx, |editor, cx| {
540 editor.set_placeholder_text(
541 "Write feedback here",
542 window,
543 cx,
544 );
545 });
546 cx.notify();
547 }
548 }
549 }
550 _ => {}
551 },
552 )
553 .detach();
554
555 CurrentPredictionState::Success {
556 model_response_editor: cx.new(|cx| {
557 let buffer = cx.new(|cx| {
558 let mut buffer = Buffer::local(
559 response
560 .debug_info
561 .as_ref()
562 .map(|p| p.model_response.as_str())
563 .unwrap_or(
564 "(Debug info not available)",
565 ),
566 cx,
567 );
568 buffer.set_language(markdown_language, cx);
569 buffer
570 });
571 let buffer = cx.new(|cx| {
572 MultiBuffer::singleton(buffer, cx)
573 });
574 let mut editor = Editor::new(
575 EditorMode::full(),
576 buffer,
577 None,
578 window,
579 cx,
580 );
581 editor.set_read_only(true);
582 editor.set_show_line_numbers(false, cx);
583 editor.set_show_gutter(false, cx);
584 editor.set_show_scrollbars(false, cx);
585 editor
586 }),
587 feedback_editor,
588 feedback: None,
589 response,
590 }
591 }
592 Err(err) => CurrentPredictionState::Failed { message: err },
593 };
594 }
595 })
596 .ok();
597 }
598 });
599
600 this.current = Some(CurrentRequest {
601 index,
602 context_editor,
603 prompt_editor: cx.new(|cx| {
604 let buffer = cx.new(|cx| {
605 let mut buffer = Buffer::local(
606 local_prompt.as_ref().unwrap_or_else(|err| err),
607 cx,
608 );
609 buffer.set_language(markdown_language.clone(), cx);
610 buffer
611 });
612 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
613 let mut editor =
614 Editor::new(EditorMode::full(), buffer, None, window, cx);
615 editor.set_read_only(true);
616 editor.set_show_line_numbers(false, cx);
617 editor.set_show_gutter(false, cx);
618 editor.set_show_scrollbars(false, cx);
619 editor
620 }),
621 buffer: buffer.clone(),
622 position: *position,
623 state: CurrentPredictionState::Requested,
624 _task: Some(task),
625 });
626 cx.notify();
627 })
628 .ok();
629 }
630 });
631 }
632
633 fn handle_prev(
634 &mut self,
635 _action: &Zeta2InspectorPrevious,
636 window: &mut Window,
637 cx: &mut Context<Self>,
638 ) {
639 self.set_current_predition(
640 self.current
641 .as_ref()
642 .map(|c| c.index)
643 .unwrap_or_default()
644 .saturating_sub(1),
645 window,
646 cx,
647 );
648 }
649
650 fn handle_next(
651 &mut self,
652 _action: &Zeta2InspectorPrevious,
653 window: &mut Window,
654 cx: &mut Context<Self>,
655 ) {
656 self.set_current_predition(
657 self.current
658 .as_ref()
659 .map(|c| c.index)
660 .unwrap_or_default()
661 .add(1)
662 .min(self.requests.len() - 1),
663 window,
664 cx,
665 );
666 }
667
668 fn handle_rate_positive(
669 &mut self,
670 _action: &Zeta2RatePredictionPositive,
671 window: &mut Window,
672 cx: &mut Context<Self>,
673 ) {
674 self.handle_rate(Feedback::Positive, window, cx);
675 }
676
677 fn handle_rate_negative(
678 &mut self,
679 _action: &Zeta2RatePredictionNegative,
680 window: &mut Window,
681 cx: &mut Context<Self>,
682 ) {
683 self.handle_rate(Feedback::Negative, window, cx);
684 }
685
686 fn handle_rate(&mut self, kind: Feedback, window: &mut Window, cx: &mut Context<Self>) {
687 let Some(last_prediction) = self.current.as_mut() else {
688 return;
689 };
690 let request = &self.requests[last_prediction.index];
691 if !request.observed_request.full_request.can_collect_data {
692 return;
693 }
694
695 let project_snapshot_task = request.project_snapshot.clone();
696
697 cx.spawn_in(window, async move |this, cx| {
698 let project_snapshot = project_snapshot_task.await;
699 this.update_in(cx, |this, window, cx| {
700 let Some(last_prediction) = this.current.as_mut() else {
701 return;
702 };
703
704 // todo! move to Self::requests?
705 let CurrentPredictionState::Success {
706 feedback: feedback_state,
707 feedback_editor,
708 model_response_editor,
709 response,
710 ..
711 } = &mut last_prediction.state
712 else {
713 return;
714 };
715
716 *feedback_state = Some(kind);
717 let text = feedback_editor.update(cx, |feedback_editor, cx| {
718 feedback_editor.set_placeholder_text(
719 "Submitted. Edit or submit again to change.",
720 window,
721 cx,
722 );
723 feedback_editor.text(cx)
724 });
725 cx.notify();
726
727 cx.defer_in(window, {
728 let model_response_editor = model_response_editor.downgrade();
729 move |_, window, cx| {
730 if let Some(model_response_editor) = model_response_editor.upgrade() {
731 model_response_editor.focus_handle(cx).focus(window);
732 }
733 }
734 });
735
736 let kind = match kind {
737 Feedback::Positive => "positive",
738 Feedback::Negative => "negative",
739 };
740
741 let request = &this.requests[last_prediction.index];
742
743 telemetry::event!(
744 "Zeta2 Prediction Rated",
745 id = response.request_id,
746 kind = kind,
747 text = text,
748 request = request.observed_request.full_request,
749 response = response,
750 project_snapshot = project_snapshot,
751 );
752 })
753 .log_err();
754 })
755 .detach();
756 }
757
758 fn focus_feedback(&mut self, window: &mut Window, cx: &mut Context<Self>) {
759 if let Some(last_prediction) = self.current.as_mut() {
760 if let CurrentPredictionState::Success {
761 feedback_editor, ..
762 } = &mut last_prediction.state
763 {
764 feedback_editor.focus_handle(cx).focus(window);
765 }
766 };
767 }
768
769 fn render_options(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
770 v_flex()
771 .gap_2()
772 .child(
773 h_flex()
774 .child(Headline::new("Options").size(HeadlineSize::Small))
775 .justify_between()
776 .child(
777 ui::Button::new("reset-options", "Reset")
778 .disabled(self.zeta.read(cx).options() == &zeta2::DEFAULT_OPTIONS)
779 .style(ButtonStyle::Outlined)
780 .size(ButtonSize::Large)
781 .on_click(cx.listener(|this, _, window, cx| {
782 this.set_input_options(&zeta2::DEFAULT_OPTIONS, window, cx);
783 })),
784 ),
785 )
786 .child(
787 v_flex()
788 .gap_2()
789 .child(
790 h_flex()
791 .gap_2()
792 .items_end()
793 .child(self.max_excerpt_bytes_input.clone())
794 .child(self.min_excerpt_bytes_input.clone())
795 .child(self.cursor_context_ratio_input.clone()),
796 )
797 .child(
798 h_flex()
799 .gap_2()
800 .items_end()
801 .child(self.max_retrieved_declarations.clone())
802 .child(self.max_prompt_bytes_input.clone())
803 .child(self.render_prompt_format_dropdown(window, cx)),
804 ),
805 )
806 }
807
808 fn render_prompt_format_dropdown(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
809 let active_format = self.zeta.read(cx).options().prompt_format;
810 let this = cx.weak_entity();
811
812 v_flex()
813 .gap_1p5()
814 .child(
815 Label::new("Prompt Format")
816 .size(LabelSize::Small)
817 .color(Color::Muted),
818 )
819 .child(
820 DropdownMenu::new(
821 "ep-prompt-format",
822 active_format.to_string(),
823 ContextMenu::build(window, cx, move |mut menu, _window, _cx| {
824 for prompt_format in PromptFormat::iter() {
825 menu = menu.item(
826 ContextMenuEntry::new(prompt_format.to_string())
827 .toggleable(IconPosition::End, active_format == prompt_format)
828 .handler({
829 let this = this.clone();
830 move |_window, cx| {
831 this.update(cx, |this, cx| {
832 let current_options =
833 this.zeta.read(cx).options().clone();
834 let options = ZetaOptions {
835 prompt_format,
836 ..current_options
837 };
838 this.set_options(options, cx);
839 })
840 .ok();
841 }
842 }),
843 )
844 }
845 menu
846 }),
847 )
848 .style(ui::DropdownStyle::Outlined),
849 )
850 }
851
852 fn render_tabs(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
853 if self.current.is_none() {
854 return None;
855 };
856
857 Some(
858 ui::ToggleButtonGroup::single_row(
859 "prediction",
860 [
861 ui::ToggleButtonSimple::new(
862 "Context",
863 cx.listener(|this, _, _, cx| {
864 this.active_view = ActiveView::Context;
865 cx.notify();
866 }),
867 ),
868 ui::ToggleButtonSimple::new(
869 "Inference",
870 cx.listener(|this, _, window, cx| {
871 this.active_view = ActiveView::Inference;
872 this.focus_feedback(window, cx);
873 cx.notify();
874 }),
875 ),
876 ],
877 )
878 .style(ui::ToggleButtonGroupStyle::Outlined)
879 .selected_index(if self.active_view == ActiveView::Context {
880 0
881 } else {
882 1
883 })
884 .into_any_element(),
885 )
886 }
887
888 fn render_stats(&self) -> Option<Div> {
889 let Some(current) = self.current.as_ref() else {
890 return None;
891 };
892
893 let (prompt_planning_time, inference_time, parsing_time) =
894 if let CurrentPredictionState::Success {
895 response:
896 PredictEditsResponse {
897 debug_info: Some(debug_info),
898 ..
899 },
900 ..
901 } = ¤t.state
902 {
903 (
904 Some(debug_info.prompt_planning_time),
905 Some(debug_info.inference_time),
906 Some(debug_info.parsing_time),
907 )
908 } else {
909 (None, None, None)
910 };
911
912 Some(
913 v_flex()
914 .p_4()
915 .gap_2()
916 .min_w(px(160.))
917 .child(Headline::new("Stats").size(HeadlineSize::Small))
918 .child(Self::render_duration(
919 "Context retrieval",
920 Some(self.requests[current.index].observed_request.retrieval_time),
921 ))
922 .child(Self::render_duration(
923 "Prompt planning",
924 prompt_planning_time,
925 ))
926 .child(Self::render_duration("Inference", inference_time))
927 .child(Self::render_duration("Parsing", parsing_time)),
928 )
929 }
930
931 fn render_duration(name: &'static str, time: Option<chrono::TimeDelta>) -> Div {
932 h_flex()
933 .gap_1()
934 .child(Label::new(name).color(Color::Muted).size(LabelSize::Small))
935 .child(match time {
936 Some(time) => Label::new(if time.num_microseconds().unwrap_or(0) >= 1000 {
937 format!("{} ms", time.num_milliseconds())
938 } else {
939 format!("{} µs", time.num_microseconds().unwrap_or(0))
940 })
941 .size(LabelSize::Small),
942 None => Label::new("...").size(LabelSize::Small),
943 })
944 }
945
946 fn render_content(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
947 if !cx.has_flag::<Zeta2FeatureFlag>() {
948 return Self::render_message("`zeta2` feature flag is not enabled");
949 }
950
951 match self.current.as_ref() {
952 None => Self::render_message("No prediction"),
953 Some(prediction) => self
954 .render_current_prediction(prediction, window, cx)
955 .into_any(),
956 }
957 }
958
959 fn render_message(message: impl Into<SharedString>) -> AnyElement {
960 v_flex()
961 .size_full()
962 .justify_center()
963 .items_center()
964 .child(Label::new(message).size(LabelSize::Large))
965 .into_any()
966 }
967
968 fn render_current_prediction(
969 &self,
970 current: &CurrentRequest,
971 window: &mut Window,
972 cx: &mut Context<Self>,
973 ) -> Div {
974 match &self.active_view {
975 ActiveView::Context => div().size_full().child(current.context_editor.clone()),
976 ActiveView::Inference => h_flex()
977 .items_start()
978 .w_full()
979 .flex_1()
980 .border_t_1()
981 .border_color(cx.theme().colors().border)
982 .bg(cx.theme().colors().editor_background)
983 .child(
984 v_flex()
985 .flex_1()
986 .gap_2()
987 .p_4()
988 .h_full()
989 .child(
990 h_flex()
991 .justify_between()
992 .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall))
993 .child(match current.state {
994 CurrentPredictionState::Requested
995 | CurrentPredictionState::Failed { .. } => {
996 ui::Chip::new("Local")
997 .bg_color(cx.theme().status().warning_background)
998 .label_color(Color::Success)
999 }
1000 CurrentPredictionState::Success { .. } => {
1001 ui::Chip::new("Cloud")
1002 .bg_color(cx.theme().status().success_background)
1003 .label_color(Color::Success)
1004 }
1005 }),
1006 )
1007 .child(current.prompt_editor.clone()),
1008 )
1009 .child(ui::vertical_divider())
1010 .child(
1011 v_flex()
1012 .flex_1()
1013 .gap_2()
1014 .h_full()
1015 .child(
1016 v_flex()
1017 .flex_1()
1018 .gap_2()
1019 .p_4()
1020 .child(
1021 ui::Headline::new("Model Response")
1022 .size(ui::HeadlineSize::XSmall),
1023 )
1024 .child(match ¤t.state {
1025 CurrentPredictionState::Success {
1026 model_response_editor,
1027 ..
1028 } => model_response_editor.clone().into_any_element(),
1029 CurrentPredictionState::Requested => v_flex()
1030 .gap_2()
1031 .child(Label::new("Loading...").buffer_font(cx))
1032 .into_any_element(),
1033 CurrentPredictionState::Failed { message } => v_flex()
1034 .gap_2()
1035 .max_w_96()
1036 .child(Label::new(message.clone()).buffer_font(cx))
1037 .into_any_element(),
1038 }),
1039 )
1040 .child(ui::divider())
1041 .child(
1042 if self.requests[current.index]
1043 .observed_request
1044 .full_request
1045 .can_collect_data
1046 && let CurrentPredictionState::Success {
1047 feedback_editor,
1048 feedback: feedback_state,
1049 ..
1050 } = ¤t.state
1051 {
1052 v_flex()
1053 .key_context("Zeta2Feedback")
1054 .on_action(cx.listener(Self::handle_rate_positive))
1055 .on_action(cx.listener(Self::handle_rate_negative))
1056 .gap_2()
1057 .p_2()
1058 .child(feedback_editor.clone())
1059 .child(
1060 h_flex()
1061 .justify_end()
1062 .w_full()
1063 .child(
1064 ButtonLike::new("rate-positive")
1065 .when(
1066 *feedback_state == Some(Feedback::Positive),
1067 |this| this.style(ButtonStyle::Filled),
1068 )
1069 .children(
1070 KeyBinding::for_action(
1071 &Zeta2RatePredictionPositive,
1072 window,
1073 cx,
1074 )
1075 .map(|k| k.size(TextSize::Small.rems(cx))),
1076 )
1077 .child(ui::Icon::new(ui::IconName::ThumbsUp))
1078 .on_click(cx.listener(
1079 |this, _, window, cx| {
1080 this.handle_rate_positive(
1081 &Zeta2RatePredictionPositive,
1082 window,
1083 cx,
1084 );
1085 },
1086 )),
1087 )
1088 .child(
1089 ButtonLike::new("rate-negative")
1090 .when(
1091 *feedback_state == Some(Feedback::Negative),
1092 |this| this.style(ButtonStyle::Filled),
1093 )
1094 .children(
1095 KeyBinding::for_action(
1096 &Zeta2RatePredictionNegative,
1097 window,
1098 cx,
1099 )
1100 .map(|k| k.size(TextSize::Small.rems(cx))),
1101 )
1102 .child(ui::Icon::new(ui::IconName::ThumbsDown))
1103 .on_click(cx.listener(
1104 |this, _, window, cx| {
1105 this.handle_rate_negative(
1106 &Zeta2RatePredictionNegative,
1107 window,
1108 cx,
1109 );
1110 },
1111 )),
1112 ),
1113 )
1114 .into_any()
1115 } else {
1116 Empty.into_any_element()
1117 },
1118 ),
1119 ),
1120 }
1121 }
1122}
1123
1124impl Focusable for Zeta2Inspector {
1125 fn focus_handle(&self, _cx: &App) -> FocusHandle {
1126 self.focus_handle.clone()
1127 }
1128}
1129
1130impl Item for Zeta2Inspector {
1131 type Event = ();
1132
1133 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
1134 "Zeta2 Inspector".into()
1135 }
1136}
1137
1138impl EventEmitter<()> for Zeta2Inspector {}
1139
1140impl Render for Zeta2Inspector {
1141 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1142 v_flex()
1143 .track_focus(&self.focus_handle)
1144 .key_context("Zeta2Inspector")
1145 .size_full()
1146 .on_action(cx.listener(Self::handle_prev))
1147 .on_action(cx.listener(Self::handle_next))
1148 .bg(cx.theme().colors().editor_background)
1149 .child(
1150 h_flex()
1151 .w_full()
1152 .child(
1153 v_flex()
1154 .flex_1()
1155 .p_4()
1156 .h_full()
1157 .justify_between()
1158 .child(self.render_options(window, cx))
1159 .gap_4()
1160 .children(self.render_tabs(cx)),
1161 )
1162 .child(ui::vertical_divider())
1163 .children(self.render_stats()),
1164 )
1165 .child(self.render_content(window, cx))
1166 }
1167}
1168
1169// Using same approach as commit view
1170
1171struct ExcerptMetadataFile {
1172 title: Arc<RelPath>,
1173 worktree_id: WorktreeId,
1174 path_style: PathStyle,
1175}
1176
1177impl language::File for ExcerptMetadataFile {
1178 fn as_local(&self) -> Option<&dyn language::LocalFile> {
1179 None
1180 }
1181
1182 fn disk_state(&self) -> DiskState {
1183 DiskState::New
1184 }
1185
1186 fn path(&self) -> &Arc<RelPath> {
1187 &self.title
1188 }
1189
1190 fn full_path(&self, _: &App) -> PathBuf {
1191 self.title.as_std_path().to_path_buf()
1192 }
1193
1194 fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
1195 self.title.file_name().unwrap()
1196 }
1197
1198 fn path_style(&self, _: &App) -> PathStyle {
1199 self.path_style
1200 }
1201
1202 fn worktree_id(&self, _: &App) -> WorktreeId {
1203 self.worktree_id
1204 }
1205
1206 fn to_proto(&self, _: &App) -> language::proto::File {
1207 unimplemented!()
1208 }
1209
1210 fn is_private(&self) -> bool {
1211 false
1212 }
1213}
1214
1215struct ZetaContextAddon {
1216 excerpt_score_components: HashMap<editor::ExcerptId, DeclarationScoreComponents>,
1217}
1218
1219impl editor::Addon for ZetaContextAddon {
1220 fn to_any(&self) -> &dyn std::any::Any {
1221 self
1222 }
1223
1224 fn render_buffer_header_controls(
1225 &self,
1226 excerpt_info: &multi_buffer::ExcerptInfo,
1227 _window: &Window,
1228 _cx: &App,
1229 ) -> Option<AnyElement> {
1230 let score_components = self.excerpt_score_components.get(&excerpt_info.id)?.clone();
1231
1232 Some(
1233 div()
1234 .id(excerpt_info.id.to_proto() as usize)
1235 .child(ui::Icon::new(IconName::Info))
1236 .cursor(CursorStyle::PointingHand)
1237 .tooltip(move |_, cx| {
1238 cx.new(|_| ScoreComponentsTooltip::new(&score_components))
1239 .into()
1240 })
1241 .into_any(),
1242 )
1243 }
1244}
1245
1246struct ScoreComponentsTooltip {
1247 text: SharedString,
1248}
1249
1250impl ScoreComponentsTooltip {
1251 fn new(components: &DeclarationScoreComponents) -> Self {
1252 Self {
1253 text: format!("{:#?}", components).into(),
1254 }
1255 }
1256}
1257
1258impl Render for ScoreComponentsTooltip {
1259 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1260 div().pl_2().pt_2p5().child(
1261 div()
1262 .elevation_2(cx)
1263 .py_1()
1264 .px_2()
1265 .child(ui::Label::new(self.text.clone()).buffer_font(cx)),
1266 )
1267 }
1268}