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