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