zeta2_tools.rs

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