ep: Add an option to filter context by type when evaluating (#58523)

This allows us to collect all context once, and evaluate subsets of it
later.

Closes #ISSUE

Release Notes:

- N/A

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<ContextRetrievalType>,
     /// 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<Vec<ContextSource>> {
+        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)?;

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<ContextSource> {
+        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<Vec<ContextSource>> {
-        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::<Vec<_>>();
+        if context_sources.is_empty() {
+            None
+        } else {
+            Some(context_sources)
         }
     }
 }
 
+pub fn context_sources_for_types(context_types: &[ContextRetrievalType]) -> Vec<ContextSource> {
+    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<ContextSource> {
+    vec![
+        ContextSource::CursorExcerpt,
+        ContextSource::CurrentFile,
+        ContextSource::EditHistory,
+        ContextSource::EditHistoryFile,
+        ContextSource::GitLog,
+    ]
+}
+
 pub async fn run_context_retrieval(
     example: &mut Example,
     app_state: Arc<EpAppState>,

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<usize>,
+    context_source_filter: Option<Vec<ContextSource>>,
 ) -> 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<usize>,
+    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<usize>,
+    context_source_filter: Option<&[ContextSource]>,
 ) -> Vec<Excerpt> {
     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<RelatedFile> {
+    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::<Vec<_>>();
+            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<usize>,
+    context_source_filter: Option<&[ContextSource]>,
 ) -> Option<usize> {
     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<usize>,
+    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<usize>,
+    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;
         }