zeta2_tools.rs

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