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