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