1use std::{collections::hash_map::Entry, path::PathBuf, str::FromStr, sync::Arc, time::Duration};
2
3use chrono::TimeDelta;
4use client::{Client, UserStore};
5use cloud_llm_client::predict_edits_v3::PromptFormat;
6use collections::HashMap;
7use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
8use futures::StreamExt as _;
9use gpui::{
10 Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task, WeakEntity, actions,
11 prelude::*,
12};
13use language::{Buffer, DiskState};
14use project::{Project, WorktreeId};
15use ui::{ContextMenu, ContextMenuEntry, DropdownMenu, prelude::*};
16use ui_input::SingleLineInput;
17use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
18use workspace::{Item, SplitDirection, Workspace};
19use zeta2::{DEFAULT_CONTEXT_OPTIONS, Zeta, ZetaOptions};
20
21use edit_prediction_context::{DeclarationStyle, EditPredictionExcerptOptions};
22
23actions!(
24 dev,
25 [
26 /// Opens the language server protocol logs viewer.
27 OpenZeta2Inspector
28 ]
29);
30
31pub fn init(cx: &mut App) {
32 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
33 workspace.register_action(move |workspace, _: &OpenZeta2Inspector, window, cx| {
34 let project = workspace.project();
35 workspace.split_item(
36 SplitDirection::Right,
37 Box::new(cx.new(|cx| {
38 Zeta2Inspector::new(
39 &project,
40 workspace.client(),
41 workspace.user_store(),
42 window,
43 cx,
44 )
45 })),
46 window,
47 cx,
48 );
49 });
50 })
51 .detach();
52}
53
54// TODO show included diagnostics, and events
55
56pub struct Zeta2Inspector {
57 focus_handle: FocusHandle,
58 project: Entity<Project>,
59 last_prediction: Option<LastPredictionState>,
60 max_excerpt_bytes_input: Entity<SingleLineInput>,
61 min_excerpt_bytes_input: Entity<SingleLineInput>,
62 cursor_context_ratio_input: Entity<SingleLineInput>,
63 max_prompt_bytes_input: Entity<SingleLineInput>,
64 active_view: ActiveView,
65 zeta: Entity<Zeta>,
66 _active_editor_subscription: Option<Subscription>,
67 _update_state_task: Task<()>,
68 _receive_task: Task<()>,
69}
70
71#[derive(PartialEq)]
72enum ActiveView {
73 Context,
74 Inference,
75}
76
77enum LastPredictionState {
78 Failed(SharedString),
79 Success(LastPrediction),
80 Replaying {
81 prediction: LastPrediction,
82 _task: Task<()>,
83 },
84}
85
86struct LastPrediction {
87 context_editor: Entity<Editor>,
88 retrieval_time: TimeDelta,
89 prompt_planning_time: TimeDelta,
90 inference_time: TimeDelta,
91 parsing_time: TimeDelta,
92 prompt_editor: Entity<Editor>,
93 model_response_editor: Entity<Editor>,
94 buffer: WeakEntity<Buffer>,
95 position: language::Anchor,
96}
97
98impl Zeta2Inspector {
99 pub fn new(
100 project: &Entity<Project>,
101 client: &Arc<Client>,
102 user_store: &Entity<UserStore>,
103 window: &mut Window,
104 cx: &mut Context<Self>,
105 ) -> Self {
106 let zeta = Zeta::global(client, user_store, cx);
107 let mut request_rx = zeta.update(cx, |zeta, _cx| zeta.debug_info());
108
109 let receive_task = cx.spawn_in(window, async move |this, cx| {
110 while let Some(prediction_result) = request_rx.next().await {
111 this.update_in(cx, |this, window, cx| match prediction_result {
112 Ok(prediction) => {
113 this.update_last_prediction(prediction, window, cx);
114 }
115 Err(err) => {
116 this.last_prediction = Some(LastPredictionState::Failed(err.into()));
117 cx.notify();
118 }
119 })
120 .ok();
121 }
122 });
123
124 let mut this = Self {
125 focus_handle: cx.focus_handle(),
126 project: project.clone(),
127 last_prediction: None,
128 active_view: ActiveView::Context,
129 max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
130 min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
131 cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),
132 max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx),
133 zeta: zeta.clone(),
134 _active_editor_subscription: None,
135 _update_state_task: Task::ready(()),
136 _receive_task: receive_task,
137 };
138 this.set_input_options(&zeta.read(cx).options().clone(), window, cx);
139 this
140 }
141
142 fn set_input_options(
143 &mut self,
144 options: &ZetaOptions,
145 window: &mut Window,
146 cx: &mut Context<Self>,
147 ) {
148 self.max_excerpt_bytes_input.update(cx, |input, cx| {
149 input.set_text(options.context.excerpt.max_bytes.to_string(), window, cx);
150 });
151 self.min_excerpt_bytes_input.update(cx, |input, cx| {
152 input.set_text(options.context.excerpt.min_bytes.to_string(), window, cx);
153 });
154 self.cursor_context_ratio_input.update(cx, |input, cx| {
155 input.set_text(
156 format!(
157 "{:.2}",
158 options
159 .context
160 .excerpt
161 .target_before_cursor_over_total_bytes
162 ),
163 window,
164 cx,
165 );
166 });
167 self.max_prompt_bytes_input.update(cx, |input, cx| {
168 input.set_text(options.max_prompt_bytes.to_string(), window, cx);
169 });
170 cx.notify();
171 }
172
173 fn set_options(&mut self, options: ZetaOptions, cx: &mut Context<Self>) {
174 self.zeta.update(cx, |this, _cx| this.set_options(options));
175
176 const THROTTLE_TIME: Duration = Duration::from_millis(100);
177
178 if let Some(
179 LastPredictionState::Success(prediction)
180 | LastPredictionState::Replaying { prediction, .. },
181 ) = self.last_prediction.take()
182 {
183 if let Some(buffer) = prediction.buffer.upgrade() {
184 let position = prediction.position;
185 let zeta = self.zeta.clone();
186 let project = self.project.clone();
187 let task = cx.spawn(async move |_this, cx| {
188 cx.background_executor().timer(THROTTLE_TIME).await;
189 if let Some(task) = zeta
190 .update(cx, |zeta, cx| {
191 zeta.refresh_prediction(&project, &buffer, position, cx)
192 })
193 .ok()
194 {
195 task.await.log_err();
196 }
197 });
198 self.last_prediction = Some(LastPredictionState::Replaying {
199 prediction,
200 _task: task,
201 });
202 } else {
203 self.last_prediction = Some(LastPredictionState::Failed("Buffer dropped".into()));
204 }
205 }
206
207 cx.notify();
208 }
209
210 fn number_input(
211 label: &'static str,
212 window: &mut Window,
213 cx: &mut Context<Self>,
214 ) -> Entity<SingleLineInput> {
215 let input = cx.new(|cx| {
216 SingleLineInput::new(window, cx, "")
217 .label(label)
218 .label_min_width(px(64.))
219 });
220
221 cx.subscribe_in(
222 &input.read(cx).editor().clone(),
223 window,
224 |this, _, event, _window, cx| {
225 let EditorEvent::BufferEdited = event else {
226 return;
227 };
228
229 fn number_input_value<T: FromStr + Default>(
230 input: &Entity<SingleLineInput>,
231 cx: &App,
232 ) -> T {
233 input
234 .read(cx)
235 .editor()
236 .read(cx)
237 .text(cx)
238 .parse::<T>()
239 .unwrap_or_default()
240 }
241
242 let mut context_options = DEFAULT_CONTEXT_OPTIONS.clone();
243 context_options.excerpt = EditPredictionExcerptOptions {
244 max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
245 min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
246 target_before_cursor_over_total_bytes: number_input_value(
247 &this.cursor_context_ratio_input,
248 cx,
249 ),
250 };
251
252 let zeta_options = this.zeta.read(cx).options();
253 this.set_options(
254 ZetaOptions {
255 context: context_options,
256 max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx),
257 max_diagnostic_bytes: zeta_options.max_diagnostic_bytes,
258 prompt_format: zeta_options.prompt_format,
259 file_indexing_parallelism: zeta_options.file_indexing_parallelism,
260 },
261 cx,
262 );
263 },
264 )
265 .detach();
266 input
267 }
268
269 fn update_last_prediction(
270 &mut self,
271 prediction: zeta2::PredictionDebugInfo,
272 window: &mut Window,
273 cx: &mut Context<Self>,
274 ) {
275 let project = self.project.read(cx);
276 let path_style = project.path_style(cx);
277 let Some(worktree_id) = project
278 .worktrees(cx)
279 .next()
280 .map(|worktree| worktree.read(cx).id())
281 else {
282 log::error!("Open a worktree to use edit prediction debug view");
283 self.last_prediction.take();
284 return;
285 };
286
287 self._update_state_task = cx.spawn_in(window, {
288 let language_registry = self.project.read(cx).languages().clone();
289 async move |this, cx| {
290 let mut languages = HashMap::default();
291 for lang_id in prediction
292 .context
293 .declarations
294 .iter()
295 .map(|snippet| snippet.declaration.identifier().language_id)
296 .chain(prediction.context.excerpt_text.language_id)
297 {
298 if let Entry::Vacant(entry) = languages.entry(lang_id) {
299 // Most snippets are gonna be the same language,
300 // so we think it's fine to do this sequentially for now
301 entry.insert(language_registry.language_for_id(lang_id).await.ok());
302 }
303 }
304
305 let markdown_language = language_registry
306 .language_for_name("Markdown")
307 .await
308 .log_err();
309
310 this.update_in(cx, |this, window, cx| {
311 let context_editor = cx.new(|cx| {
312 let multibuffer = cx.new(|cx| {
313 let mut multibuffer = MultiBuffer::new(language::Capability::ReadOnly);
314 let excerpt_file = Arc::new(ExcerptMetadataFile {
315 title: RelPath::unix("Cursor Excerpt").unwrap().into(),
316 path_style,
317 worktree_id,
318 });
319
320 let excerpt_buffer = cx.new(|cx| {
321 let mut buffer =
322 Buffer::local(prediction.context.excerpt_text.body, cx);
323 if let Some(language) = prediction
324 .context
325 .excerpt_text
326 .language_id
327 .as_ref()
328 .and_then(|id| languages.get(id))
329 {
330 buffer.set_language(language.clone(), cx);
331 }
332 buffer.file_updated(excerpt_file, cx);
333 buffer
334 });
335
336 multibuffer.push_excerpts(
337 excerpt_buffer,
338 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
339 cx,
340 );
341
342 for snippet in &prediction.context.declarations {
343 let path = this
344 .project
345 .read(cx)
346 .path_for_entry(snippet.declaration.project_entry_id(), cx);
347
348 let snippet_file = Arc::new(ExcerptMetadataFile {
349 title: RelPath::unix(&format!(
350 "{} (Score density: {})",
351 path.map(|p| p.path.display(path_style).to_string())
352 .unwrap_or_else(|| "".to_string()),
353 snippet.score_density(DeclarationStyle::Declaration)
354 ))
355 .unwrap()
356 .into(),
357 path_style,
358 worktree_id,
359 });
360
361 let excerpt_buffer = cx.new(|cx| {
362 let mut buffer =
363 Buffer::local(snippet.declaration.item_text().0, cx);
364 buffer.file_updated(snippet_file, cx);
365 if let Some(language) =
366 languages.get(&snippet.declaration.identifier().language_id)
367 {
368 buffer.set_language(language.clone(), cx);
369 }
370 buffer
371 });
372
373 multibuffer.push_excerpts(
374 excerpt_buffer,
375 [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
376 cx,
377 );
378 }
379
380 multibuffer
381 });
382
383 Editor::new(EditorMode::full(), multibuffer, None, window, cx)
384 });
385
386 let last_prediction = LastPrediction {
387 context_editor,
388 prompt_editor: cx.new(|cx| {
389 let buffer = cx.new(|cx| {
390 let mut buffer = Buffer::local(prediction.request.prompt, cx);
391 buffer.set_language(markdown_language.clone(), cx);
392 buffer
393 });
394 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
395 let mut editor =
396 Editor::new(EditorMode::full(), buffer, None, window, cx);
397 editor.set_read_only(true);
398 editor.set_show_line_numbers(false, cx);
399 editor.set_show_gutter(false, cx);
400 editor.set_show_scrollbars(false, cx);
401 editor
402 }),
403 model_response_editor: cx.new(|cx| {
404 let buffer = cx.new(|cx| {
405 let mut buffer =
406 Buffer::local(prediction.request.model_response, cx);
407 buffer.set_language(markdown_language, cx);
408 buffer
409 });
410 let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
411 let mut editor =
412 Editor::new(EditorMode::full(), buffer, None, window, cx);
413 editor.set_read_only(true);
414 editor.set_show_line_numbers(false, cx);
415 editor.set_show_gutter(false, cx);
416 editor.set_show_scrollbars(false, cx);
417 editor
418 }),
419 retrieval_time: prediction.retrieval_time,
420 prompt_planning_time: prediction.request.prompt_planning_time,
421 inference_time: prediction.request.inference_time,
422 parsing_time: prediction.request.parsing_time,
423 buffer: prediction.buffer,
424 position: prediction.position,
425 };
426 this.last_prediction = Some(LastPredictionState::Success(last_prediction));
427 cx.notify();
428 })
429 .ok();
430 }
431 });
432 }
433
434 fn render_options(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
435 v_flex()
436 .gap_2()
437 .child(
438 h_flex()
439 .child(Headline::new("Options").size(HeadlineSize::Small))
440 .justify_between()
441 .child(
442 ui::Button::new("reset-options", "Reset")
443 .disabled(self.zeta.read(cx).options() == &zeta2::DEFAULT_OPTIONS)
444 .style(ButtonStyle::Outlined)
445 .size(ButtonSize::Large)
446 .on_click(cx.listener(|this, _, window, cx| {
447 this.set_input_options(&zeta2::DEFAULT_OPTIONS, window, cx);
448 })),
449 ),
450 )
451 .child(
452 v_flex()
453 .gap_2()
454 .child(
455 h_flex()
456 .gap_2()
457 .items_end()
458 .child(self.max_excerpt_bytes_input.clone())
459 .child(self.min_excerpt_bytes_input.clone())
460 .child(self.cursor_context_ratio_input.clone()),
461 )
462 .child(
463 h_flex()
464 .gap_2()
465 .items_end()
466 .child(self.max_prompt_bytes_input.clone())
467 .child(self.render_prompt_format_dropdown(window, cx)),
468 ),
469 )
470 }
471
472 fn render_prompt_format_dropdown(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
473 let active_format = self.zeta.read(cx).options().prompt_format;
474 let this = cx.weak_entity();
475
476 v_flex()
477 .gap_1p5()
478 .child(
479 Label::new("Prompt Format")
480 .size(LabelSize::Small)
481 .color(Color::Muted),
482 )
483 .child(
484 DropdownMenu::new(
485 "ep-prompt-format",
486 active_format.to_string(),
487 ContextMenu::build(window, cx, move |mut menu, _window, _cx| {
488 for prompt_format in PromptFormat::iter() {
489 menu = menu.item(
490 ContextMenuEntry::new(prompt_format.to_string())
491 .toggleable(IconPosition::End, active_format == prompt_format)
492 .handler({
493 let this = this.clone();
494 move |_window, cx| {
495 this.update(cx, |this, cx| {
496 let current_options =
497 this.zeta.read(cx).options().clone();
498 let options = ZetaOptions {
499 prompt_format,
500 ..current_options
501 };
502 this.set_options(options, cx);
503 })
504 .ok();
505 }
506 }),
507 )
508 }
509 menu
510 }),
511 )
512 .style(ui::DropdownStyle::Outlined),
513 )
514 }
515
516 fn render_tabs(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
517 let Some(LastPredictionState::Success { .. } | LastPredictionState::Replaying { .. }) =
518 self.last_prediction.as_ref()
519 else {
520 return None;
521 };
522
523 Some(
524 ui::ToggleButtonGroup::single_row(
525 "prediction",
526 [
527 ui::ToggleButtonSimple::new(
528 "Context",
529 cx.listener(|this, _, _, cx| {
530 this.active_view = ActiveView::Context;
531 cx.notify();
532 }),
533 ),
534 ui::ToggleButtonSimple::new(
535 "Inference",
536 cx.listener(|this, _, _, cx| {
537 this.active_view = ActiveView::Inference;
538 cx.notify();
539 }),
540 ),
541 ],
542 )
543 .style(ui::ToggleButtonGroupStyle::Outlined)
544 .selected_index(if self.active_view == ActiveView::Context {
545 0
546 } else {
547 1
548 })
549 .into_any_element(),
550 )
551 }
552
553 fn render_stats(&self) -> Option<Div> {
554 let Some(
555 LastPredictionState::Success(prediction)
556 | LastPredictionState::Replaying { prediction, .. },
557 ) = self.last_prediction.as_ref()
558 else {
559 return None;
560 };
561
562 Some(
563 v_flex()
564 .p_4()
565 .gap_2()
566 .min_w(px(160.))
567 .child(Headline::new("Stats").size(HeadlineSize::Small))
568 .child(Self::render_duration(
569 "Context retrieval",
570 prediction.retrieval_time,
571 ))
572 .child(Self::render_duration(
573 "Prompt planning",
574 prediction.prompt_planning_time,
575 ))
576 .child(Self::render_duration(
577 "Inference",
578 prediction.inference_time,
579 ))
580 .child(Self::render_duration("Parsing", prediction.parsing_time)),
581 )
582 }
583
584 fn render_duration(name: &'static str, time: chrono::TimeDelta) -> Div {
585 h_flex()
586 .gap_1()
587 .child(Label::new(name).color(Color::Muted).size(LabelSize::Small))
588 .child(
589 Label::new(if time.num_microseconds().unwrap_or(0) >= 1000 {
590 format!("{} ms", time.num_milliseconds())
591 } else {
592 format!("{} µs", time.num_microseconds().unwrap_or(0))
593 })
594 .size(LabelSize::Small),
595 )
596 }
597
598 fn render_content(&self, cx: &mut Context<Self>) -> AnyElement {
599 match self.last_prediction.as_ref() {
600 None => v_flex()
601 .size_full()
602 .justify_center()
603 .items_center()
604 .child(Label::new("No prediction").size(LabelSize::Large))
605 .into_any(),
606 Some(LastPredictionState::Success(prediction)) => {
607 self.render_last_prediction(prediction, cx).into_any()
608 }
609 Some(LastPredictionState::Replaying { prediction, _task }) => self
610 .render_last_prediction(prediction, cx)
611 .opacity(0.6)
612 .into_any(),
613 Some(LastPredictionState::Failed(err)) => v_flex()
614 .p_4()
615 .gap_2()
616 .child(Label::new(err.clone()).buffer_font(cx))
617 .into_any(),
618 }
619 }
620
621 fn render_last_prediction(&self, prediction: &LastPrediction, cx: &mut Context<Self>) -> Div {
622 match &self.active_view {
623 ActiveView::Context => div().size_full().child(prediction.context_editor.clone()),
624 ActiveView::Inference => h_flex()
625 .items_start()
626 .w_full()
627 .flex_1()
628 .border_t_1()
629 .border_color(cx.theme().colors().border)
630 .bg(cx.theme().colors().editor_background)
631 .child(
632 v_flex()
633 .flex_1()
634 .gap_2()
635 .p_4()
636 .h_full()
637 .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall))
638 .child(prediction.prompt_editor.clone()),
639 )
640 .child(ui::vertical_divider())
641 .child(
642 v_flex()
643 .flex_1()
644 .gap_2()
645 .h_full()
646 .p_4()
647 .child(ui::Headline::new("Model Response").size(ui::HeadlineSize::XSmall))
648 .child(prediction.model_response_editor.clone()),
649 ),
650 }
651 }
652}
653
654impl Focusable for Zeta2Inspector {
655 fn focus_handle(&self, _cx: &App) -> FocusHandle {
656 self.focus_handle.clone()
657 }
658}
659
660impl Item for Zeta2Inspector {
661 type Event = ();
662
663 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
664 "Zeta2 Inspector".into()
665 }
666}
667
668impl EventEmitter<()> for Zeta2Inspector {}
669
670impl Render for Zeta2Inspector {
671 fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
672 v_flex()
673 .size_full()
674 .bg(cx.theme().colors().editor_background)
675 .child(
676 h_flex()
677 .w_full()
678 .child(
679 v_flex()
680 .flex_1()
681 .p_4()
682 .h_full()
683 .justify_between()
684 .child(self.render_options(window, cx))
685 .gap_4()
686 .children(self.render_tabs(cx)),
687 )
688 .child(ui::vertical_divider())
689 .children(self.render_stats()),
690 )
691 .child(self.render_content(cx))
692 }
693}
694
695// Using same approach as commit view
696
697struct ExcerptMetadataFile {
698 title: Arc<RelPath>,
699 worktree_id: WorktreeId,
700 path_style: PathStyle,
701}
702
703impl language::File for ExcerptMetadataFile {
704 fn as_local(&self) -> Option<&dyn language::LocalFile> {
705 None
706 }
707
708 fn disk_state(&self) -> DiskState {
709 DiskState::New
710 }
711
712 fn path(&self) -> &Arc<RelPath> {
713 &self.title
714 }
715
716 fn full_path(&self, _: &App) -> PathBuf {
717 self.title.as_std_path().to_path_buf()
718 }
719
720 fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
721 self.title.file_name().unwrap()
722 }
723
724 fn path_style(&self, _: &App) -> PathStyle {
725 self.path_style
726 }
727
728 fn worktree_id(&self, _: &App) -> WorktreeId {
729 self.worktree_id
730 }
731
732 fn to_proto(&self, _: &App) -> language::proto::File {
733 unimplemented!()
734 }
735
736 fn is_private(&self) -> bool {
737 false
738 }
739}