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