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