zeta2_tools.rs

   1use std::{cmp::Reverse, path::PathBuf, str::FromStr, sync::Arc, time::Duration};
   2
   3use chrono::TimeDelta;
   4use client::{Client, UserStore};
   5use cloud_llm_client::predict_edits_v3::{
   6    self, DeclarationScoreComponents, PredictEditsRequest, PredictEditsResponse, PromptFormat,
   7};
   8use collections::HashMap;
   9use editor::{Editor, EditorEvent, EditorMode, ExcerptRange, MultiBuffer};
  10use feature_flags::FeatureFlagAppExt as _;
  11use futures::{FutureExt, StreamExt as _, channel::oneshot, future::Shared};
  12use gpui::{
  13    CursorStyle, Empty, Entity, EventEmitter, FocusHandle, Focusable, Subscription, Task,
  14    WeakEntity, actions, prelude::*,
  15};
  16use language::{Buffer, DiskState};
  17use ordered_float::OrderedFloat;
  18use project::{Project, WorktreeId, telemetry_snapshot::TelemetrySnapshot};
  19use ui::{ButtonLike, ContextMenu, ContextMenuEntry, DropdownMenu, KeyBinding, prelude::*};
  20use ui_input::InputField;
  21use util::{ResultExt, paths::PathStyle, rel_path::RelPath};
  22use workspace::{Item, SplitDirection, Workspace};
  23use zeta2::{PredictionDebugInfo, Zeta, Zeta2FeatureFlag, ZetaOptions};
  24
  25use edit_prediction_context::{EditPredictionContextOptions, EditPredictionExcerptOptions};
  26
  27actions!(
  28    dev,
  29    [
  30        /// Opens the language server protocol logs viewer.
  31        OpenZeta2Inspector,
  32        /// Rate prediction as positive.
  33        Zeta2RatePredictionPositive,
  34        /// Rate prediction as negative.
  35        Zeta2RatePredictionNegative,
  36    ]
  37);
  38
  39pub fn init(cx: &mut App) {
  40    cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
  41        workspace.register_action(move |workspace, _: &OpenZeta2Inspector, window, cx| {
  42            let project = workspace.project();
  43            workspace.split_item(
  44                SplitDirection::Right,
  45                Box::new(cx.new(|cx| {
  46                    Zeta2Inspector::new(
  47                        &project,
  48                        workspace.client(),
  49                        workspace.user_store(),
  50                        window,
  51                        cx,
  52                    )
  53                })),
  54                window,
  55                cx,
  56            );
  57        });
  58    })
  59    .detach();
  60}
  61
  62// TODO show included diagnostics, and events
  63
  64pub struct Zeta2Inspector {
  65    focus_handle: FocusHandle,
  66    project: Entity<Project>,
  67    last_prediction: Option<LastPrediction>,
  68    max_excerpt_bytes_input: Entity<InputField>,
  69    min_excerpt_bytes_input: Entity<InputField>,
  70    cursor_context_ratio_input: Entity<InputField>,
  71    max_prompt_bytes_input: Entity<InputField>,
  72    max_retrieved_declarations: Entity<InputField>,
  73    active_view: ActiveView,
  74    zeta: Entity<Zeta>,
  75    _active_editor_subscription: Option<Subscription>,
  76    _update_state_task: Task<()>,
  77    _receive_task: Task<()>,
  78}
  79
  80#[derive(PartialEq)]
  81enum ActiveView {
  82    Context,
  83    Inference,
  84}
  85
  86struct LastPrediction {
  87    context_editor: Entity<Editor>,
  88    prompt_editor: Entity<Editor>,
  89    retrieval_time: TimeDelta,
  90    buffer: WeakEntity<Buffer>,
  91    position: language::Anchor,
  92    state: LastPredictionState,
  93    request: PredictEditsRequest,
  94    project_snapshot: Shared<Task<Arc<TelemetrySnapshot>>>,
  95    _task: Option<Task<()>>,
  96}
  97
  98#[derive(Clone, Copy, PartialEq)]
  99enum Feedback {
 100    Positive,
 101    Negative,
 102}
 103
 104enum LastPredictionState {
 105    Requested,
 106    Success {
 107        model_response_editor: Entity<Editor>,
 108        feedback_editor: Entity<Editor>,
 109        feedback: Option<Feedback>,
 110        response: predict_edits_v3::PredictEditsResponse,
 111    },
 112    Failed {
 113        message: String,
 114    },
 115}
 116
 117impl Zeta2Inspector {
 118    pub fn new(
 119        project: &Entity<Project>,
 120        client: &Arc<Client>,
 121        user_store: &Entity<UserStore>,
 122        window: &mut Window,
 123        cx: &mut Context<Self>,
 124    ) -> Self {
 125        let zeta = Zeta::global(client, user_store, cx);
 126        let mut request_rx = zeta.update(cx, |zeta, _cx| zeta.debug_info());
 127
 128        let receive_task = cx.spawn_in(window, async move |this, cx| {
 129            while let Some(prediction) = request_rx.next().await {
 130                this.update_in(cx, |this, window, cx| {
 131                    this.update_last_prediction(prediction, window, cx)
 132                })
 133                .ok();
 134            }
 135        });
 136
 137        let mut this = Self {
 138            focus_handle: cx.focus_handle(),
 139            project: project.clone(),
 140            last_prediction: None,
 141            active_view: ActiveView::Inference,
 142            max_excerpt_bytes_input: Self::number_input("Max Excerpt Bytes", window, cx),
 143            min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx),
 144            cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx),
 145            max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx),
 146            max_retrieved_declarations: Self::number_input("Max Retrieved Definitions", window, cx),
 147            zeta: zeta.clone(),
 148            _active_editor_subscription: None,
 149            _update_state_task: Task::ready(()),
 150            _receive_task: receive_task,
 151        };
 152        this.set_input_options(&zeta.read(cx).options().clone(), window, cx);
 153        this
 154    }
 155
 156    fn set_input_options(
 157        &mut self,
 158        options: &ZetaOptions,
 159        window: &mut Window,
 160        cx: &mut Context<Self>,
 161    ) {
 162        self.max_excerpt_bytes_input.update(cx, |input, cx| {
 163            input.set_text(options.context.excerpt.max_bytes.to_string(), window, cx);
 164        });
 165        self.min_excerpt_bytes_input.update(cx, |input, cx| {
 166            input.set_text(options.context.excerpt.min_bytes.to_string(), window, cx);
 167        });
 168        self.cursor_context_ratio_input.update(cx, |input, cx| {
 169            input.set_text(
 170                format!(
 171                    "{:.2}",
 172                    options
 173                        .context
 174                        .excerpt
 175                        .target_before_cursor_over_total_bytes
 176                ),
 177                window,
 178                cx,
 179            );
 180        });
 181        self.max_prompt_bytes_input.update(cx, |input, cx| {
 182            input.set_text(options.max_prompt_bytes.to_string(), window, cx);
 183        });
 184        self.max_retrieved_declarations.update(cx, |input, cx| {
 185            input.set_text(
 186                options.context.max_retrieved_declarations.to_string(),
 187                window,
 188                cx,
 189            );
 190        });
 191        cx.notify();
 192    }
 193
 194    fn set_options(&mut self, options: ZetaOptions, cx: &mut Context<Self>) {
 195        self.zeta.update(cx, |this, _cx| this.set_options(options));
 196
 197        const THROTTLE_TIME: Duration = Duration::from_millis(100);
 198
 199        if let Some(prediction) = self.last_prediction.as_mut() {
 200            if let Some(buffer) = prediction.buffer.upgrade() {
 201                let position = prediction.position;
 202                let zeta = self.zeta.clone();
 203                let project = self.project.clone();
 204                prediction._task = Some(cx.spawn(async move |_this, cx| {
 205                    cx.background_executor().timer(THROTTLE_TIME).await;
 206                    if let Some(task) = zeta
 207                        .update(cx, |zeta, cx| {
 208                            zeta.refresh_prediction(&project, &buffer, position, cx)
 209                        })
 210                        .ok()
 211                    {
 212                        task.await.log_err();
 213                    }
 214                }));
 215                prediction.state = LastPredictionState::Requested;
 216            } else {
 217                self.last_prediction.take();
 218            }
 219        }
 220
 221        cx.notify();
 222    }
 223
 224    fn number_input(
 225        label: &'static str,
 226        window: &mut Window,
 227        cx: &mut Context<Self>,
 228    ) -> Entity<InputField> {
 229        let input = cx.new(|cx| {
 230            InputField::new(window, cx, "")
 231                .label(label)
 232                .label_min_width(px(64.))
 233        });
 234
 235        cx.subscribe_in(
 236            &input.read(cx).editor().clone(),
 237            window,
 238            |this, _, event, _window, cx| {
 239                let EditorEvent::BufferEdited = event else {
 240                    return;
 241                };
 242
 243                fn number_input_value<T: FromStr + Default>(
 244                    input: &Entity<InputField>,
 245                    cx: &App,
 246                ) -> T {
 247                    input
 248                        .read(cx)
 249                        .editor()
 250                        .read(cx)
 251                        .text(cx)
 252                        .parse::<T>()
 253                        .unwrap_or_default()
 254                }
 255
 256                let zeta_options = this.zeta.read(cx).options().clone();
 257
 258                let context_options = EditPredictionContextOptions {
 259                    excerpt: EditPredictionExcerptOptions {
 260                        max_bytes: number_input_value(&this.max_excerpt_bytes_input, cx),
 261                        min_bytes: number_input_value(&this.min_excerpt_bytes_input, cx),
 262                        target_before_cursor_over_total_bytes: number_input_value(
 263                            &this.cursor_context_ratio_input,
 264                            cx,
 265                        ),
 266                    },
 267                    max_retrieved_declarations: number_input_value(
 268                        &this.max_retrieved_declarations,
 269                        cx,
 270                    ),
 271                    ..zeta_options.context
 272                };
 273
 274                this.set_options(
 275                    ZetaOptions {
 276                        context: context_options,
 277                        max_prompt_bytes: number_input_value(&this.max_prompt_bytes_input, cx),
 278                        max_diagnostic_bytes: zeta_options.max_diagnostic_bytes,
 279                        prompt_format: zeta_options.prompt_format,
 280                        file_indexing_parallelism: zeta_options.file_indexing_parallelism,
 281                    },
 282                    cx,
 283                );
 284            },
 285        )
 286        .detach();
 287        input
 288    }
 289
 290    fn update_last_prediction(
 291        &mut self,
 292        prediction: zeta2::PredictionDebugInfo,
 293        window: &mut Window,
 294        cx: &mut Context<Self>,
 295    ) {
 296        let project = self.project.read(cx);
 297        let path_style = project.path_style(cx);
 298        let Some(worktree_id) = project
 299            .worktrees(cx)
 300            .next()
 301            .map(|worktree| worktree.read(cx).id())
 302        else {
 303            log::error!("Open a worktree to use edit prediction debug view");
 304            self.last_prediction.take();
 305            return;
 306        };
 307
 308        self._update_state_task = cx.spawn_in(window, {
 309            let language_registry = self.project.read(cx).languages().clone();
 310            async move |this, cx| {
 311                let mut languages = HashMap::default();
 312                for ext in prediction
 313                    .request
 314                    .referenced_declarations
 315                    .iter()
 316                    .filter_map(|snippet| snippet.path.extension())
 317                    .chain(prediction.request.excerpt_path.extension())
 318                {
 319                    if !languages.contains_key(ext) {
 320                        // Most snippets are gonna be the same language,
 321                        // so we think it's fine to do this sequentially for now
 322                        languages.insert(
 323                            ext.to_owned(),
 324                            language_registry
 325                                .language_for_name_or_extension(&ext.to_string_lossy())
 326                                .await
 327                                .ok(),
 328                        );
 329                    }
 330                }
 331
 332                let markdown_language = language_registry
 333                    .language_for_name("Markdown")
 334                    .await
 335                    .log_err();
 336
 337                this.update_in(cx, |this, window, cx| {
 338                    let context_editor = cx.new(|cx| {
 339                        let mut excerpt_score_components = HashMap::default();
 340
 341                        let multibuffer = cx.new(|cx| {
 342                            let mut multibuffer = MultiBuffer::new(language::Capability::ReadOnly);
 343                            let excerpt_file = Arc::new(ExcerptMetadataFile {
 344                                title: RelPath::unix("Cursor Excerpt").unwrap().into(),
 345                                path_style,
 346                                worktree_id,
 347                            });
 348
 349                            let excerpt_buffer = cx.new(|cx| {
 350                                let mut buffer =
 351                                    Buffer::local(prediction.request.excerpt.clone(), cx);
 352                                if let Some(language) = prediction
 353                                    .request
 354                                    .excerpt_path
 355                                    .extension()
 356                                    .and_then(|ext| languages.get(ext))
 357                                {
 358                                    buffer.set_language(language.clone(), cx);
 359                                }
 360                                buffer.file_updated(excerpt_file, cx);
 361                                buffer
 362                            });
 363
 364                            multibuffer.push_excerpts(
 365                                excerpt_buffer,
 366                                [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
 367                                cx,
 368                            );
 369
 370                            let mut declarations =
 371                                prediction.request.referenced_declarations.clone();
 372                            declarations.sort_unstable_by_key(|declaration| {
 373                                Reverse(OrderedFloat(declaration.declaration_score))
 374                            });
 375
 376                            for snippet in &declarations {
 377                                let snippet_file = Arc::new(ExcerptMetadataFile {
 378                                    title: RelPath::unix(&format!(
 379                                        "{} (Score: {})",
 380                                        snippet.path.display(),
 381                                        snippet.declaration_score
 382                                    ))
 383                                    .unwrap()
 384                                    .into(),
 385                                    path_style,
 386                                    worktree_id,
 387                                });
 388
 389                                let excerpt_buffer = cx.new(|cx| {
 390                                    let mut buffer = Buffer::local(snippet.text.clone(), cx);
 391                                    buffer.file_updated(snippet_file, cx);
 392                                    if let Some(ext) = snippet.path.extension()
 393                                        && let Some(language) = languages.get(ext)
 394                                    {
 395                                        buffer.set_language(language.clone(), cx);
 396                                    }
 397                                    buffer
 398                                });
 399
 400                                let excerpt_ids = multibuffer.push_excerpts(
 401                                    excerpt_buffer,
 402                                    [ExcerptRange::new(text::Anchor::MIN..text::Anchor::MAX)],
 403                                    cx,
 404                                );
 405                                let excerpt_id = excerpt_ids.first().unwrap();
 406
 407                                excerpt_score_components
 408                                    .insert(*excerpt_id, snippet.score_components.clone());
 409                            }
 410
 411                            multibuffer
 412                        });
 413
 414                        let mut editor =
 415                            Editor::new(EditorMode::full(), multibuffer, None, window, cx);
 416                        editor.register_addon(ZetaContextAddon {
 417                            excerpt_score_components,
 418                        });
 419                        editor
 420                    });
 421
 422                    let PredictionDebugInfo {
 423                        response_rx,
 424                        position,
 425                        buffer,
 426                        retrieval_time,
 427                        local_prompt,
 428                        ..
 429                    } = prediction;
 430
 431                    let task = cx.spawn_in(window, {
 432                        let markdown_language = markdown_language.clone();
 433                        async move |this, cx| {
 434                            let response = response_rx.await;
 435
 436                            this.update_in(cx, |this, window, cx| {
 437                                if let Some(prediction) = this.last_prediction.as_mut() {
 438                                    prediction.state = match response {
 439                                        Ok(Ok(response)) => {
 440                                            if let Some(debug_info) = &response.debug_info {
 441                                                prediction.prompt_editor.update(
 442                                                    cx,
 443                                                    |prompt_editor, cx| {
 444                                                        prompt_editor.set_text(
 445                                                            debug_info.prompt.as_str(),
 446                                                            window,
 447                                                            cx,
 448                                                        );
 449                                                    },
 450                                                );
 451                                            }
 452
 453                                            let feedback_editor = cx.new(|cx| {
 454                                                let buffer = cx.new(|cx| {
 455                                                    let mut buffer = Buffer::local("", cx);
 456                                                    buffer.set_language(
 457                                                        markdown_language.clone(),
 458                                                        cx,
 459                                                    );
 460                                                    buffer
 461                                                });
 462                                                let buffer =
 463                                                    cx.new(|cx| MultiBuffer::singleton(buffer, cx));
 464                                                let mut editor = Editor::new(
 465                                                    EditorMode::AutoHeight {
 466                                                        min_lines: 3,
 467                                                        max_lines: None,
 468                                                    },
 469                                                    buffer,
 470                                                    None,
 471                                                    window,
 472                                                    cx,
 473                                                );
 474                                                editor.set_placeholder_text(
 475                                                    "Write feedback here",
 476                                                    window,
 477                                                    cx,
 478                                                );
 479                                                editor.set_show_line_numbers(false, cx);
 480                                                editor.set_show_gutter(false, cx);
 481                                                editor.set_show_scrollbars(false, cx);
 482                                                editor
 483                                            });
 484
 485                                            cx.subscribe_in(
 486                                                &feedback_editor,
 487                                                window,
 488                                                |this, editor, ev, window, cx| match ev {
 489                                                    EditorEvent::BufferEdited => {
 490                                                        if let Some(last_prediction) =
 491                                                            this.last_prediction.as_mut()
 492                                                            && let LastPredictionState::Success {
 493                                                                feedback: feedback_state,
 494                                                                ..
 495                                                            } = &mut last_prediction.state
 496                                                        {
 497                                                            if feedback_state.take().is_some() {
 498                                                                editor.update(cx, |editor, cx| {
 499                                                                    editor.set_placeholder_text(
 500                                                                        "Write feedback here",
 501                                                                        window,
 502                                                                        cx,
 503                                                                    );
 504                                                                });
 505                                                                cx.notify();
 506                                                            }
 507                                                        }
 508                                                    }
 509                                                    _ => {}
 510                                                },
 511                                            )
 512                                            .detach();
 513
 514                                            LastPredictionState::Success {
 515                                                model_response_editor: cx.new(|cx| {
 516                                                    let buffer = cx.new(|cx| {
 517                                                        let mut buffer = Buffer::local(
 518                                                            response
 519                                                                .debug_info
 520                                                                .as_ref()
 521                                                                .map(|p| p.model_response.as_str())
 522                                                                .unwrap_or(
 523                                                                    "(Debug info not available)",
 524                                                                ),
 525                                                            cx,
 526                                                        );
 527                                                        buffer.set_language(markdown_language, cx);
 528                                                        buffer
 529                                                    });
 530                                                    let buffer = cx.new(|cx| {
 531                                                        MultiBuffer::singleton(buffer, cx)
 532                                                    });
 533                                                    let mut editor = Editor::new(
 534                                                        EditorMode::full(),
 535                                                        buffer,
 536                                                        None,
 537                                                        window,
 538                                                        cx,
 539                                                    );
 540                                                    editor.set_read_only(true);
 541                                                    editor.set_show_line_numbers(false, cx);
 542                                                    editor.set_show_gutter(false, cx);
 543                                                    editor.set_show_scrollbars(false, cx);
 544                                                    editor
 545                                                }),
 546                                                feedback_editor,
 547                                                feedback: None,
 548                                                response,
 549                                            }
 550                                        }
 551                                        Ok(Err(err)) => {
 552                                            LastPredictionState::Failed { message: err }
 553                                        }
 554                                        Err(oneshot::Canceled) => LastPredictionState::Failed {
 555                                            message: "Canceled".to_string(),
 556                                        },
 557                                    };
 558                                }
 559                            })
 560                            .ok();
 561                        }
 562                    });
 563
 564                    let project_snapshot_task = TelemetrySnapshot::new(&this.project, cx);
 565
 566                    this.last_prediction = Some(LastPrediction {
 567                        context_editor,
 568                        prompt_editor: cx.new(|cx| {
 569                            let buffer = cx.new(|cx| {
 570                                let mut buffer =
 571                                    Buffer::local(local_prompt.unwrap_or_else(|err| err), cx);
 572                                buffer.set_language(markdown_language.clone(), cx);
 573                                buffer
 574                            });
 575                            let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
 576                            let mut editor =
 577                                Editor::new(EditorMode::full(), buffer, None, window, cx);
 578                            editor.set_read_only(true);
 579                            editor.set_show_line_numbers(false, cx);
 580                            editor.set_show_gutter(false, cx);
 581                            editor.set_show_scrollbars(false, cx);
 582                            editor
 583                        }),
 584                        retrieval_time,
 585                        buffer,
 586                        position,
 587                        state: LastPredictionState::Requested,
 588                        project_snapshot: cx
 589                            .foreground_executor()
 590                            .spawn(async move { Arc::new(project_snapshot_task.await) })
 591                            .shared(),
 592                        request: prediction.request,
 593                        _task: Some(task),
 594                    });
 595                    cx.notify();
 596                })
 597                .ok();
 598            }
 599        });
 600    }
 601
 602    fn handle_rate_positive(
 603        &mut self,
 604        _action: &Zeta2RatePredictionPositive,
 605        window: &mut Window,
 606        cx: &mut Context<Self>,
 607    ) {
 608        self.handle_rate(Feedback::Positive, window, cx);
 609    }
 610
 611    fn handle_rate_negative(
 612        &mut self,
 613        _action: &Zeta2RatePredictionNegative,
 614        window: &mut Window,
 615        cx: &mut Context<Self>,
 616    ) {
 617        self.handle_rate(Feedback::Negative, window, cx);
 618    }
 619
 620    fn handle_rate(&mut self, kind: Feedback, window: &mut Window, cx: &mut Context<Self>) {
 621        let Some(last_prediction) = self.last_prediction.as_mut() else {
 622            return;
 623        };
 624        if !last_prediction.request.can_collect_data {
 625            return;
 626        }
 627
 628        let project_snapshot_task = last_prediction.project_snapshot.clone();
 629
 630        cx.spawn_in(window, async move |this, cx| {
 631            let project_snapshot = project_snapshot_task.await;
 632            this.update_in(cx, |this, window, cx| {
 633                let Some(last_prediction) = this.last_prediction.as_mut() else {
 634                    return;
 635                };
 636
 637                let LastPredictionState::Success {
 638                    feedback: feedback_state,
 639                    feedback_editor,
 640                    model_response_editor,
 641                    response,
 642                    ..
 643                } = &mut last_prediction.state
 644                else {
 645                    return;
 646                };
 647
 648                *feedback_state = Some(kind);
 649                let text = feedback_editor.update(cx, |feedback_editor, cx| {
 650                    feedback_editor.set_placeholder_text(
 651                        "Submitted. Edit or submit again to change.",
 652                        window,
 653                        cx,
 654                    );
 655                    feedback_editor.text(cx)
 656                });
 657                cx.notify();
 658
 659                cx.defer_in(window, {
 660                    let model_response_editor = model_response_editor.downgrade();
 661                    move |_, window, cx| {
 662                        if let Some(model_response_editor) = model_response_editor.upgrade() {
 663                            model_response_editor.focus_handle(cx).focus(window);
 664                        }
 665                    }
 666                });
 667
 668                let kind = match kind {
 669                    Feedback::Positive => "positive",
 670                    Feedback::Negative => "negative",
 671                };
 672
 673                telemetry::event!(
 674                    "Zeta2 Prediction Rated",
 675                    id = response.request_id,
 676                    kind = kind,
 677                    text = text,
 678                    request = last_prediction.request,
 679                    response = response,
 680                    project_snapshot = project_snapshot,
 681                );
 682            })
 683            .log_err();
 684        })
 685        .detach();
 686    }
 687
 688    fn focus_feedback(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 689        if let Some(last_prediction) = self.last_prediction.as_mut() {
 690            if let LastPredictionState::Success {
 691                feedback_editor, ..
 692            } = &mut last_prediction.state
 693            {
 694                feedback_editor.focus_handle(cx).focus(window);
 695            }
 696        };
 697    }
 698
 699    fn render_options(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
 700        v_flex()
 701            .gap_2()
 702            .child(
 703                h_flex()
 704                    .child(Headline::new("Options").size(HeadlineSize::Small))
 705                    .justify_between()
 706                    .child(
 707                        ui::Button::new("reset-options", "Reset")
 708                            .disabled(self.zeta.read(cx).options() == &zeta2::DEFAULT_OPTIONS)
 709                            .style(ButtonStyle::Outlined)
 710                            .size(ButtonSize::Large)
 711                            .on_click(cx.listener(|this, _, window, cx| {
 712                                this.set_input_options(&zeta2::DEFAULT_OPTIONS, window, cx);
 713                            })),
 714                    ),
 715            )
 716            .child(
 717                v_flex()
 718                    .gap_2()
 719                    .child(
 720                        h_flex()
 721                            .gap_2()
 722                            .items_end()
 723                            .child(self.max_excerpt_bytes_input.clone())
 724                            .child(self.min_excerpt_bytes_input.clone())
 725                            .child(self.cursor_context_ratio_input.clone()),
 726                    )
 727                    .child(
 728                        h_flex()
 729                            .gap_2()
 730                            .items_end()
 731                            .child(self.max_retrieved_declarations.clone())
 732                            .child(self.max_prompt_bytes_input.clone())
 733                            .child(self.render_prompt_format_dropdown(window, cx)),
 734                    ),
 735            )
 736    }
 737
 738    fn render_prompt_format_dropdown(&self, window: &mut Window, cx: &mut Context<Self>) -> Div {
 739        let active_format = self.zeta.read(cx).options().prompt_format;
 740        let this = cx.weak_entity();
 741
 742        v_flex()
 743            .gap_1p5()
 744            .child(
 745                Label::new("Prompt Format")
 746                    .size(LabelSize::Small)
 747                    .color(Color::Muted),
 748            )
 749            .child(
 750                DropdownMenu::new(
 751                    "ep-prompt-format",
 752                    active_format.to_string(),
 753                    ContextMenu::build(window, cx, move |mut menu, _window, _cx| {
 754                        for prompt_format in PromptFormat::iter() {
 755                            menu = menu.item(
 756                                ContextMenuEntry::new(prompt_format.to_string())
 757                                    .toggleable(IconPosition::End, active_format == prompt_format)
 758                                    .handler({
 759                                        let this = this.clone();
 760                                        move |_window, cx| {
 761                                            this.update(cx, |this, cx| {
 762                                                let current_options =
 763                                                    this.zeta.read(cx).options().clone();
 764                                                let options = ZetaOptions {
 765                                                    prompt_format,
 766                                                    ..current_options
 767                                                };
 768                                                this.set_options(options, cx);
 769                                            })
 770                                            .ok();
 771                                        }
 772                                    }),
 773                            )
 774                        }
 775                        menu
 776                    }),
 777                )
 778                .style(ui::DropdownStyle::Outlined),
 779            )
 780    }
 781
 782    fn render_tabs(&self, cx: &mut Context<Self>) -> Option<AnyElement> {
 783        if self.last_prediction.is_none() {
 784            return None;
 785        };
 786
 787        Some(
 788            ui::ToggleButtonGroup::single_row(
 789                "prediction",
 790                [
 791                    ui::ToggleButtonSimple::new(
 792                        "Context",
 793                        cx.listener(|this, _, _, cx| {
 794                            this.active_view = ActiveView::Context;
 795                            cx.notify();
 796                        }),
 797                    ),
 798                    ui::ToggleButtonSimple::new(
 799                        "Inference",
 800                        cx.listener(|this, _, window, cx| {
 801                            this.active_view = ActiveView::Inference;
 802                            this.focus_feedback(window, cx);
 803                            cx.notify();
 804                        }),
 805                    ),
 806                ],
 807            )
 808            .style(ui::ToggleButtonGroupStyle::Outlined)
 809            .selected_index(if self.active_view == ActiveView::Context {
 810                0
 811            } else {
 812                1
 813            })
 814            .into_any_element(),
 815        )
 816    }
 817
 818    fn render_stats(&self) -> Option<Div> {
 819        let Some(prediction) = self.last_prediction.as_ref() else {
 820            return None;
 821        };
 822
 823        let (prompt_planning_time, inference_time, parsing_time) =
 824            if let LastPredictionState::Success {
 825                response:
 826                    PredictEditsResponse {
 827                        debug_info: Some(debug_info),
 828                        ..
 829                    },
 830                ..
 831            } = &prediction.state
 832            {
 833                (
 834                    Some(debug_info.prompt_planning_time),
 835                    Some(debug_info.inference_time),
 836                    Some(debug_info.parsing_time),
 837                )
 838            } else {
 839                (None, None, None)
 840            };
 841
 842        Some(
 843            v_flex()
 844                .p_4()
 845                .gap_2()
 846                .min_w(px(160.))
 847                .child(Headline::new("Stats").size(HeadlineSize::Small))
 848                .child(Self::render_duration(
 849                    "Context retrieval",
 850                    Some(prediction.retrieval_time),
 851                ))
 852                .child(Self::render_duration(
 853                    "Prompt planning",
 854                    prompt_planning_time,
 855                ))
 856                .child(Self::render_duration("Inference", inference_time))
 857                .child(Self::render_duration("Parsing", parsing_time)),
 858        )
 859    }
 860
 861    fn render_duration(name: &'static str, time: Option<chrono::TimeDelta>) -> Div {
 862        h_flex()
 863            .gap_1()
 864            .child(Label::new(name).color(Color::Muted).size(LabelSize::Small))
 865            .child(match time {
 866                Some(time) => Label::new(if time.num_microseconds().unwrap_or(0) >= 1000 {
 867                    format!("{} ms", time.num_milliseconds())
 868                } else {
 869                    format!("{} µs", time.num_microseconds().unwrap_or(0))
 870                })
 871                .size(LabelSize::Small),
 872                None => Label::new("...").size(LabelSize::Small),
 873            })
 874    }
 875
 876    fn render_content(&self, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {
 877        if !cx.has_flag::<Zeta2FeatureFlag>() {
 878            return Self::render_message("`zeta2` feature flag is not enabled");
 879        }
 880
 881        match self.last_prediction.as_ref() {
 882            None => Self::render_message("No prediction"),
 883            Some(prediction) => self
 884                .render_last_prediction(prediction, window, cx)
 885                .into_any(),
 886        }
 887    }
 888
 889    fn render_message(message: impl Into<SharedString>) -> AnyElement {
 890        v_flex()
 891            .size_full()
 892            .justify_center()
 893            .items_center()
 894            .child(Label::new(message).size(LabelSize::Large))
 895            .into_any()
 896    }
 897
 898    fn render_last_prediction(
 899        &self,
 900        prediction: &LastPrediction,
 901        window: &mut Window,
 902        cx: &mut Context<Self>,
 903    ) -> Div {
 904        match &self.active_view {
 905            ActiveView::Context => div().size_full().child(prediction.context_editor.clone()),
 906            ActiveView::Inference => h_flex()
 907                .items_start()
 908                .w_full()
 909                .flex_1()
 910                .border_t_1()
 911                .border_color(cx.theme().colors().border)
 912                .bg(cx.theme().colors().editor_background)
 913                .child(
 914                    v_flex()
 915                        .flex_1()
 916                        .gap_2()
 917                        .p_4()
 918                        .h_full()
 919                        .child(
 920                            h_flex()
 921                                .justify_between()
 922                                .child(ui::Headline::new("Prompt").size(ui::HeadlineSize::XSmall))
 923                                .child(match prediction.state {
 924                                    LastPredictionState::Requested
 925                                    | LastPredictionState::Failed { .. } => ui::Chip::new("Local")
 926                                        .bg_color(cx.theme().status().warning_background)
 927                                        .label_color(Color::Success),
 928                                    LastPredictionState::Success { .. } => ui::Chip::new("Cloud")
 929                                        .bg_color(cx.theme().status().success_background)
 930                                        .label_color(Color::Success),
 931                                }),
 932                        )
 933                        .child(prediction.prompt_editor.clone()),
 934                )
 935                .child(ui::vertical_divider())
 936                .child(
 937                    v_flex()
 938                        .flex_1()
 939                        .gap_2()
 940                        .h_full()
 941                        .child(
 942                            v_flex()
 943                                .flex_1()
 944                                .gap_2()
 945                                .p_4()
 946                                .child(
 947                                    ui::Headline::new("Model Response")
 948                                        .size(ui::HeadlineSize::XSmall),
 949                                )
 950                                .child(match &prediction.state {
 951                                    LastPredictionState::Success {
 952                                        model_response_editor,
 953                                        ..
 954                                    } => model_response_editor.clone().into_any_element(),
 955                                    LastPredictionState::Requested => v_flex()
 956                                        .gap_2()
 957                                        .child(Label::new("Loading...").buffer_font(cx))
 958                                        .into_any_element(),
 959                                    LastPredictionState::Failed { message } => v_flex()
 960                                        .gap_2()
 961                                        .max_w_96()
 962                                        .child(Label::new(message.clone()).buffer_font(cx))
 963                                        .into_any_element(),
 964                                }),
 965                        )
 966                        .child(ui::divider())
 967                        .child(
 968                            if prediction.request.can_collect_data
 969                                && let LastPredictionState::Success {
 970                                    feedback_editor,
 971                                    feedback: feedback_state,
 972                                    ..
 973                                } = &prediction.state
 974                            {
 975                                v_flex()
 976                                    .key_context("Zeta2Feedback")
 977                                    .on_action(cx.listener(Self::handle_rate_positive))
 978                                    .on_action(cx.listener(Self::handle_rate_negative))
 979                                    .gap_2()
 980                                    .p_2()
 981                                    .child(feedback_editor.clone())
 982                                    .child(
 983                                        h_flex()
 984                                            .justify_end()
 985                                            .w_full()
 986                                            .child(
 987                                                ButtonLike::new("rate-positive")
 988                                                    .when(
 989                                                        *feedback_state == Some(Feedback::Positive),
 990                                                        |this| this.style(ButtonStyle::Filled),
 991                                                    )
 992                                                    .children(
 993                                                        KeyBinding::for_action(
 994                                                            &Zeta2RatePredictionPositive,
 995                                                            window,
 996                                                            cx,
 997                                                        )
 998                                                        .map(|k| k.size(TextSize::Small.rems(cx))),
 999                                                    )
1000                                                    .child(ui::Icon::new(ui::IconName::ThumbsUp))
1001                                                    .on_click(cx.listener(
1002                                                        |this, _, window, cx| {
1003                                                            this.handle_rate_positive(
1004                                                                &Zeta2RatePredictionPositive,
1005                                                                window,
1006                                                                cx,
1007                                                            );
1008                                                        },
1009                                                    )),
1010                                            )
1011                                            .child(
1012                                                ButtonLike::new("rate-negative")
1013                                                    .when(
1014                                                        *feedback_state == Some(Feedback::Negative),
1015                                                        |this| this.style(ButtonStyle::Filled),
1016                                                    )
1017                                                    .children(
1018                                                        KeyBinding::for_action(
1019                                                            &Zeta2RatePredictionNegative,
1020                                                            window,
1021                                                            cx,
1022                                                        )
1023                                                        .map(|k| k.size(TextSize::Small.rems(cx))),
1024                                                    )
1025                                                    .child(ui::Icon::new(ui::IconName::ThumbsDown))
1026                                                    .on_click(cx.listener(
1027                                                        |this, _, window, cx| {
1028                                                            this.handle_rate_negative(
1029                                                                &Zeta2RatePredictionNegative,
1030                                                                window,
1031                                                                cx,
1032                                                            );
1033                                                        },
1034                                                    )),
1035                                            ),
1036                                    )
1037                                    .into_any()
1038                            } else {
1039                                Empty.into_any_element()
1040                            },
1041                        ),
1042                ),
1043        }
1044    }
1045}
1046
1047impl Focusable for Zeta2Inspector {
1048    fn focus_handle(&self, _cx: &App) -> FocusHandle {
1049        self.focus_handle.clone()
1050    }
1051}
1052
1053impl Item for Zeta2Inspector {
1054    type Event = ();
1055
1056    fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
1057        "Zeta2 Inspector".into()
1058    }
1059}
1060
1061impl EventEmitter<()> for Zeta2Inspector {}
1062
1063impl Render for Zeta2Inspector {
1064    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1065        v_flex()
1066            .size_full()
1067            .bg(cx.theme().colors().editor_background)
1068            .child(
1069                h_flex()
1070                    .w_full()
1071                    .child(
1072                        v_flex()
1073                            .flex_1()
1074                            .p_4()
1075                            .h_full()
1076                            .justify_between()
1077                            .child(self.render_options(window, cx))
1078                            .gap_4()
1079                            .children(self.render_tabs(cx)),
1080                    )
1081                    .child(ui::vertical_divider())
1082                    .children(self.render_stats()),
1083            )
1084            .child(self.render_content(window, cx))
1085    }
1086}
1087
1088// Using same approach as commit view
1089
1090struct ExcerptMetadataFile {
1091    title: Arc<RelPath>,
1092    worktree_id: WorktreeId,
1093    path_style: PathStyle,
1094}
1095
1096impl language::File for ExcerptMetadataFile {
1097    fn as_local(&self) -> Option<&dyn language::LocalFile> {
1098        None
1099    }
1100
1101    fn disk_state(&self) -> DiskState {
1102        DiskState::New
1103    }
1104
1105    fn path(&self) -> &Arc<RelPath> {
1106        &self.title
1107    }
1108
1109    fn full_path(&self, _: &App) -> PathBuf {
1110        self.title.as_std_path().to_path_buf()
1111    }
1112
1113    fn file_name<'a>(&'a self, _: &'a App) -> &'a str {
1114        self.title.file_name().unwrap()
1115    }
1116
1117    fn path_style(&self, _: &App) -> PathStyle {
1118        self.path_style
1119    }
1120
1121    fn worktree_id(&self, _: &App) -> WorktreeId {
1122        self.worktree_id
1123    }
1124
1125    fn to_proto(&self, _: &App) -> language::proto::File {
1126        unimplemented!()
1127    }
1128
1129    fn is_private(&self) -> bool {
1130        false
1131    }
1132}
1133
1134struct ZetaContextAddon {
1135    excerpt_score_components: HashMap<editor::ExcerptId, DeclarationScoreComponents>,
1136}
1137
1138impl editor::Addon for ZetaContextAddon {
1139    fn to_any(&self) -> &dyn std::any::Any {
1140        self
1141    }
1142
1143    fn render_buffer_header_controls(
1144        &self,
1145        excerpt_info: &multi_buffer::ExcerptInfo,
1146        _window: &Window,
1147        _cx: &App,
1148    ) -> Option<AnyElement> {
1149        let score_components = self.excerpt_score_components.get(&excerpt_info.id)?.clone();
1150
1151        Some(
1152            div()
1153                .id(excerpt_info.id.to_proto() as usize)
1154                .child(ui::Icon::new(IconName::Info))
1155                .cursor(CursorStyle::PointingHand)
1156                .tooltip(move |_, cx| {
1157                    cx.new(|_| ScoreComponentsTooltip::new(&score_components))
1158                        .into()
1159                })
1160                .into_any(),
1161        )
1162    }
1163}
1164
1165struct ScoreComponentsTooltip {
1166    text: SharedString,
1167}
1168
1169impl ScoreComponentsTooltip {
1170    fn new(components: &DeclarationScoreComponents) -> Self {
1171        Self {
1172            text: format!("{:#?}", components).into(),
1173        }
1174    }
1175}
1176
1177impl Render for ScoreComponentsTooltip {
1178    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1179        div().pl_2().pt_2p5().child(
1180            div()
1181                .elevation_2(cx)
1182                .py_1()
1183                .px_2()
1184                .child(ui::Label::new(self.text.clone()).buffer_font(cx)),
1185        )
1186    }
1187}