From c9953f063806ae7037d7f58733985635e51894f6 Mon Sep 17 00:00:00 2001 From: Oleksiy Syvokon Date: Thu, 4 Jun 2026 11:41:19 +0300 Subject: [PATCH] ep: Add an option to filter context by type when evaluating This allows us to collect all context once, and evaluate subsets of it later. --- crates/edit_prediction_cli/src/main.rs | 36 +++++++++- .../src/retrieve_context.rs | 65 +++++++++++++---- crates/edit_prediction_cli/src/score.rs | 71 ++++++++++++++++--- 3 files changed, 147 insertions(+), 25 deletions(-) diff --git a/crates/edit_prediction_cli/src/main.rs b/crates/edit_prediction_cli/src/main.rs index 8076d2280f686e..c079c9291aba56 100644 --- a/crates/edit_prediction_cli/src/main.rs +++ b/crates/edit_prediction_cli/src/main.rs @@ -36,7 +36,7 @@ use gaoya::minhash::{ MinHashIndex, MinHasher, MinHasher32, calculate_minhash_params, compute_minhash_similarity, }; use gpui::{AppContext as _, BackgroundExecutor, Task}; -use zeta_prompt::ZetaFormat; +use zeta_prompt::{ContextSource, ZetaFormat}; use reqwest_client::ReqwestClient; use serde::{Deserialize, Deserializer, Serialize, Serializer}; @@ -59,7 +59,9 @@ use crate::paths::{FAILED_EXAMPLES_DIR, RUN_DIR}; use crate::predict::run_prediction; use crate::progress::Progress; use crate::pull_examples::{fetch_settled_examples_after, parse_settled_after_input}; -use crate::retrieve_context::{ContextRetrievalType, run_context_retrieval}; +use crate::retrieve_context::{ + ContextRetrievalType, context_sources_for_types, run_context_retrieval, +}; use crate::score::run_scoring; use crate::split_commit::SplitCommitArgs; use crate::split_dataset::SplitArgs; @@ -274,6 +276,15 @@ impl Display for Command { if args.context_only { write!(f, " --context-only")?; } + if !args.context_types.is_empty() { + write!(f, " --type=")?; + for (index, context_type) in args.context_types.iter().enumerate() { + if index > 0 { + write!(f, ",")?; + } + write!(f, "{}", context_type)?; + } + } if args.related_context_limit != score::EVAL_RELATED_CONTEXT_TOKENS_LIMIT { write!(f, " --related-context-limit={}", args.related_context_limit)?; } @@ -337,6 +348,10 @@ struct EvalArgs { /// Only compute editable context coverage from expected patches and retrieved context. #[clap(long)] context_only: bool, + /// Only score persisted related context excerpts from these context types. + /// May be repeated or comma-delimited, e.g. `--type=current-file,edit-history`. + #[arg(long = "type", value_enum, value_delimiter = ',')] + context_types: Vec, /// Maximum number of retrieved context tokens to include when scoring. #[clap(long, default_value_t = score::EVAL_RELATED_CONTEXT_TOKENS_LIMIT)] related_context_limit: usize, @@ -348,6 +363,16 @@ struct EvalArgs { verbose: bool, } +impl EvalArgs { + fn context_source_filter(&self) -> Option> { + if self.context_types.is_empty() { + None + } else { + Some(context_sources_for_types(&self.context_types)) + } + } +} + #[derive(Clone, Copy, Default, Debug, PartialEq, Eq, Hash)] pub enum TeacherBackend { Sonnet46, @@ -1322,15 +1347,19 @@ fn main() { cx.clone(), false, None, + None, ) .await?; } Command::Eval(args) => { + let context_source_filter = + args.context_source_filter(); if args.context_only { score::run_context_coverage_scoring( example, &example_progress, Some(args.related_context_limit * 3), + context_source_filter.as_deref(), )?; } else { run_scoring( @@ -1341,6 +1370,7 @@ fn main() { cx.clone(), true, Some(args.related_context_limit * 3), + context_source_filter, ) .await?; } @@ -1495,11 +1525,13 @@ fn main() { match &command { Command::Eval(args) => { let examples = finished_examples.lock().unwrap(); + let context_source_filter = args.context_source_filter(); score::print_report( &examples, args.verbose, args.context_only, Some(args.related_context_limit * 3), + context_source_filter.as_deref(), ); if let Some(summary_path) = &args.summary_json { score::write_summary_json(&examples, summary_path)?; diff --git a/crates/edit_prediction_cli/src/retrieve_context.rs b/crates/edit_prediction_cli/src/retrieve_context.rs index 9db722827fb9b5..398cc471d79c74 100644 --- a/crates/edit_prediction_cli/src/retrieve_context.rs +++ b/crates/edit_prediction_cli/src/retrieve_context.rs @@ -50,29 +50,64 @@ impl std::fmt::Display for ContextRetrievalType { } impl ContextRetrievalType { + pub fn context_sources(self) -> Vec { + match self { + ContextRetrievalType::Lsp => vec![ContextSource::Lsp], + ContextRetrievalType::Editable => editable_context_sources(), + ContextRetrievalType::CurrentFile => vec![ContextSource::CurrentFile], + ContextRetrievalType::EditHistory => vec![ContextSource::EditHistory], + ContextRetrievalType::EditHistoryFile => vec![ContextSource::EditHistoryFile], + ContextRetrievalType::GitLog => vec![ContextSource::GitLog], + ContextRetrievalType::OracleFile => vec![ContextSource::OracleFile], + ContextRetrievalType::All => { + let mut sources = vec![ContextSource::Lsp]; + sources.extend(editable_context_sources()); + sources + } + ContextRetrievalType::None => Vec::new(), + } + } + fn includes_lsp(self) -> bool { - matches!(self, ContextRetrievalType::Lsp | ContextRetrievalType::All) + self.context_sources().contains(&ContextSource::Lsp) } fn editable_context_sources(self) -> Option> { - match self { - ContextRetrievalType::Editable | ContextRetrievalType::All => Some(vec![ - ContextSource::CursorExcerpt, - ContextSource::CurrentFile, - ContextSource::EditHistory, - ContextSource::EditHistoryFile, - ContextSource::GitLog, - ]), - ContextRetrievalType::CurrentFile => Some(vec![ContextSource::CurrentFile]), - ContextRetrievalType::EditHistory => Some(vec![ContextSource::EditHistory]), - ContextRetrievalType::EditHistoryFile => Some(vec![ContextSource::EditHistoryFile]), - ContextRetrievalType::GitLog => Some(vec![ContextSource::GitLog]), - ContextRetrievalType::OracleFile => Some(vec![ContextSource::OracleFile]), - ContextRetrievalType::Lsp | ContextRetrievalType::None => None, + let context_sources = self + .context_sources() + .into_iter() + .filter(|context_source| *context_source != ContextSource::Lsp) + .collect::>(); + if context_sources.is_empty() { + None + } else { + Some(context_sources) } } } +pub fn context_sources_for_types(context_types: &[ContextRetrievalType]) -> Vec { + let mut context_sources = Vec::new(); + for context_type in context_types { + for context_source in context_type.context_sources() { + if !context_sources.contains(&context_source) { + context_sources.push(context_source); + } + } + } + context_sources +} + +fn editable_context_sources() -> Vec { + vec![ + ContextSource::CursorExcerpt, + ContextSource::CurrentFile, + ContextSource::EditHistory, + ContextSource::EditHistoryFile, + ContextSource::GitLog, + ] +} + pub async fn run_context_retrieval( example: &mut Example, app_state: Arc, diff --git a/crates/edit_prediction_cli/src/score.rs b/crates/edit_prediction_cli/src/score.rs index dab209d749f070..2a860df0fdb001 100644 --- a/crates/edit_prediction_cli/src/score.rs +++ b/crates/edit_prediction_cli/src/score.rs @@ -17,6 +17,7 @@ use std::fs::File; use std::io::BufWriter; use std::path::Path; use std::sync::Arc; +use zeta_prompt::{ContextSource, RelatedFile}; pub const EVAL_RELATED_CONTEXT_TOKENS_LIMIT: usize = 4000; @@ -28,6 +29,7 @@ pub async fn run_scoring( cx: AsyncApp, allow_missing_predictions: bool, retrieved_context_byte_limit: Option, + context_source_filter: Option>, ) -> anyhow::Result<()> { if !(allow_missing_predictions && args.provider.is_none() && example.predictions.is_empty()) { run_prediction(example, args, app_state, example_progress, cx.clone()).await?; @@ -81,6 +83,7 @@ pub async fn run_scoring( &example_for_scoring, prompt_inputs, retrieved_context_byte_limit, + context_source_filter.as_deref(), ); let mut scores = vec![]; @@ -149,6 +152,7 @@ pub fn run_context_coverage_scoring( example: &mut Example, example_progress: &ExampleProgress, retrieved_context_byte_limit: Option, + context_source_filter: Option<&[ContextSource]>, ) -> anyhow::Result<()> { let progress = example_progress.start(Step::Score); @@ -157,7 +161,12 @@ pub fn run_context_coverage_scoring( .prompt_inputs .as_ref() .context("prompt_inputs is required for context coverage scoring")?; - let context = context_excerpts(example, prompt_inputs, retrieved_context_byte_limit); + let context = context_excerpts( + example, + prompt_inputs, + retrieved_context_byte_limit, + context_source_filter, + ); let editable_context_coverage = example .spec @@ -183,6 +192,7 @@ fn context_excerpts( _example: &Example, prompt_inputs: &zeta_prompt::ZetaPromptInput, retrieved_context_byte_limit: Option, + context_source_filter: Option<&[ContextSource]>, ) -> Vec { let mut context = Vec::new(); @@ -197,10 +207,11 @@ fn context_excerpts( } if let Some(related_files) = &prompt_inputs.related_files { + let related_files = filtered_related_files(related_files, context_source_filter); let related_files = if let Some(max_bytes) = retrieved_context_byte_limit { - limit_retrieved_context_to_bytes(related_files, max_bytes) + limit_retrieved_context_to_bytes(&related_files, max_bytes) } else { - related_files.clone() + related_files }; for related_file in &related_files { for excerpt in &related_file.excerpts { @@ -224,15 +235,48 @@ fn context_excerpts( context } +fn filtered_related_files( + related_files: &[RelatedFile], + context_source_filter: Option<&[ContextSource]>, +) -> Vec { + let Some(context_source_filter) = context_source_filter else { + return related_files.to_vec(); + }; + + related_files + .iter() + .filter_map(|related_file| { + let excerpts = related_file + .excerpts + .iter() + .filter(|excerpt| context_source_filter.contains(&excerpt.context_source)) + .cloned() + .collect::>(); + if excerpts.is_empty() { + None + } else { + Some(RelatedFile { + path: related_file.path.clone(), + max_row: related_file.max_row, + excerpts, + in_open_source_repo: related_file.in_open_source_repo, + }) + } + }) + .collect() +} + fn retrieved_context_bytes( example: &Example, retrieved_context_byte_limit: Option, + context_source_filter: Option<&[ContextSource]>, ) -> Option { let related_files = example.prompt_inputs.as_ref()?.related_files.as_ref()?; + let related_files = filtered_related_files(related_files, context_source_filter); let related_files = if let Some(max_bytes) = retrieved_context_byte_limit { - limit_retrieved_context_to_bytes(related_files, max_bytes) + limit_retrieved_context_to_bytes(&related_files, max_bytes) } else { - related_files.clone() + related_files }; Some( related_files @@ -248,6 +292,7 @@ pub fn print_report( verbose: bool, context_only: bool, retrieved_context_byte_limit: Option, + context_source_filter: Option<&[ContextSource]>, ) { const MAX_EXAMPLES_DEFAULT: usize = 20; use crate::metrics::ClassificationMetrics; @@ -255,7 +300,12 @@ pub fn print_report( const LINE_WIDTH: usize = 101; if context_only { - print_context_coverage_report(examples, verbose, retrieved_context_byte_limit); + print_context_coverage_report( + examples, + verbose, + retrieved_context_byte_limit, + context_source_filter, + ); return; } @@ -314,7 +364,9 @@ pub fn print_report( let mut skipped_lines: usize = 0; for example in examples { - if let Some(bytes) = retrieved_context_bytes(example, retrieved_context_byte_limit) { + if let Some(bytes) = + retrieved_context_bytes(example, retrieved_context_byte_limit, context_source_filter) + { retrieved_context_bytes_sum += bytes as f64; retrieved_context_bytes_count += 1; } @@ -680,6 +732,7 @@ fn print_context_coverage_report( examples: &[Example], verbose: bool, retrieved_context_byte_limit: Option, + context_source_filter: Option<&[ContextSource]>, ) { const MAX_EXAMPLES_DEFAULT: usize = 20; const LINE_WIDTH: usize = 120; @@ -721,7 +774,9 @@ fn print_context_coverage_report( let mut skipped_lines = 0; for example in examples { - if let Some(bytes) = retrieved_context_bytes(example, retrieved_context_byte_limit) { + if let Some(bytes) = + retrieved_context_bytes(example, retrieved_context_byte_limit, context_source_filter) + { retrieved_context_bytes_sum += bytes as f64; retrieved_context_bytes_count += 1; }