From eda7a49f0185bec1f899ba3f3aa652152588a0f2 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Mon, 20 Oct 2025 12:26:41 -0300 Subject: [PATCH] zeta2: Max retrieved definitions option (#40515) Release Notes: - N/A --- .../src/edit_prediction_context.rs | 13 ++++++++----- crates/zeta2/src/zeta2.rs | 2 +- crates/zeta2_tools/src/zeta2_tools.rs | 14 ++++++++++++++ crates/zeta_cli/src/main.rs | 6 +++--- 4 files changed, 26 insertions(+), 9 deletions(-) diff --git a/crates/edit_prediction_context/src/edit_prediction_context.rs b/crates/edit_prediction_context/src/edit_prediction_context.rs index 34941f93b4017dfe8e96802b5078daf5652ed22a..f52a2259cf83ff992a904b6b5c9b3ceea7c0a71e 100644 --- a/crates/edit_prediction_context/src/edit_prediction_context.rs +++ b/crates/edit_prediction_context/src/edit_prediction_context.rs @@ -27,9 +27,9 @@ pub use predict_edits_v3::Line; #[derive(Clone, Debug, PartialEq)] pub struct EditPredictionContextOptions { pub use_imports: bool, - pub use_references: bool, pub excerpt: EditPredictionExcerptOptions, pub score: EditPredictionScoreOptions, + pub max_retrieved_declarations: u8, } #[derive(Clone, Debug)] @@ -118,7 +118,7 @@ impl EditPredictionContext { )?; let excerpt_text = excerpt.text(buffer); - let declarations = if options.use_references + let declarations = if options.max_retrieved_declarations > 0 && let Some(index_state) = index_state { let excerpt_occurrences = @@ -136,7 +136,7 @@ impl EditPredictionContext { let references = get_references(&excerpt, &excerpt_text, buffer); - scored_declarations( + let mut declarations = scored_declarations( &options.score, &index_state, &excerpt, @@ -146,7 +146,10 @@ impl EditPredictionContext { references, cursor_offset_in_file, buffer, - ) + ); + // TODO [zeta2] if we need this when we ship, we should probably do it in a smarter way + declarations.truncate(options.max_retrieved_declarations as usize); + declarations } else { vec![] }; @@ -200,7 +203,6 @@ mod tests { buffer_snapshot, EditPredictionContextOptions { use_imports: true, - use_references: true, excerpt: EditPredictionExcerptOptions { max_bytes: 60, min_bytes: 10, @@ -209,6 +211,7 @@ mod tests { score: EditPredictionScoreOptions { omit_excerpt_overlaps: true, }, + max_retrieved_declarations: u8::MAX, }, Some(index.clone()), cx, diff --git a/crates/zeta2/src/zeta2.rs b/crates/zeta2/src/zeta2.rs index 6e92443b8b4c91ce6bb168d3c87ce4ddce49bc35..7c945d6ebbe9c994977adbcc72c6d8fc175930d4 100644 --- a/crates/zeta2/src/zeta2.rs +++ b/crates/zeta2/src/zeta2.rs @@ -48,7 +48,7 @@ const MAX_EVENT_COUNT: usize = 16; pub const DEFAULT_CONTEXT_OPTIONS: EditPredictionContextOptions = EditPredictionContextOptions { use_imports: true, - use_references: false, + max_retrieved_declarations: 0, excerpt: EditPredictionExcerptOptions { max_bytes: 512, min_bytes: 128, diff --git a/crates/zeta2_tools/src/zeta2_tools.rs b/crates/zeta2_tools/src/zeta2_tools.rs index 005a3f1e48b56332d83aa127781f6dd5d95daa0d..0ac4fb2162ca632618df0c2b0d256b2fd7c30742 100644 --- a/crates/zeta2_tools/src/zeta2_tools.rs +++ b/crates/zeta2_tools/src/zeta2_tools.rs @@ -68,6 +68,7 @@ pub struct Zeta2Inspector { min_excerpt_bytes_input: Entity, cursor_context_ratio_input: Entity, max_prompt_bytes_input: Entity, + max_retrieved_declarations: Entity, active_view: ActiveView, zeta: Entity, _active_editor_subscription: Option, @@ -133,6 +134,7 @@ impl Zeta2Inspector { min_excerpt_bytes_input: Self::number_input("Min Excerpt Bytes", window, cx), cursor_context_ratio_input: Self::number_input("Cursor Context Ratio", window, cx), max_prompt_bytes_input: Self::number_input("Max Prompt Bytes", window, cx), + max_retrieved_declarations: Self::number_input("Max Retrieved Definitions", window, cx), zeta: zeta.clone(), _active_editor_subscription: None, _update_state_task: Task::ready(()), @@ -170,6 +172,13 @@ impl Zeta2Inspector { self.max_prompt_bytes_input.update(cx, |input, cx| { input.set_text(options.max_prompt_bytes.to_string(), window, cx); }); + self.max_retrieved_declarations.update(cx, |input, cx| { + input.set_text( + options.context.max_retrieved_declarations.to_string(), + window, + cx, + ); + }); cx.notify(); } @@ -246,6 +255,10 @@ impl Zeta2Inspector { cx, ), }, + max_retrieved_declarations: number_input_value( + &this.max_retrieved_declarations, + cx, + ), ..zeta_options.context }; @@ -536,6 +549,7 @@ impl Zeta2Inspector { h_flex() .gap_2() .items_end() + .child(self.max_retrieved_declarations.clone()) .child(self.max_prompt_bytes_input.clone()) .child(self.render_prompt_format_dropdown(window, cx)), ), diff --git a/crates/zeta_cli/src/main.rs b/crates/zeta_cli/src/main.rs index 75b859d2f55d99cc37c961455ddcdb86a5f49351..149b13719f2075143d81c164e8d91bbdaca17384 100644 --- a/crates/zeta_cli/src/main.rs +++ b/crates/zeta_cli/src/main.rs @@ -94,8 +94,8 @@ struct Zeta2Args { file_indexing_parallelism: usize, #[arg(long, default_value_t = false)] disable_imports_gathering: bool, - #[arg(long, default_value_t = false)] - disable_reference_retrieval: bool, + #[arg(long, default_value_t = u8::MAX)] + max_retrieved_definitions: u8, } #[derive(clap::ValueEnum, Default, Debug, Clone)] @@ -302,7 +302,7 @@ impl Zeta2Args { fn to_options(&self, omit_excerpt_overlaps: bool) -> zeta2::ZetaOptions { zeta2::ZetaOptions { context: EditPredictionContextOptions { - use_references: !self.disable_reference_retrieval, + max_retrieved_declarations: self.max_retrieved_definitions, use_imports: !self.disable_imports_gathering, excerpt: EditPredictionExcerptOptions { max_bytes: self.max_excerpt_bytes,