Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
osyvokon merged 1 commit into from
Jun 4, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions crates/edit_prediction_cli/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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;
Expand Down Expand Up @@ -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)?;
}
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -1341,6 +1370,7 @@ fn main() {
cx.clone(),
true,
Some(args.related_context_limit * 3),
context_source_filter,
)
.await?;
}
Expand Down Expand Up @@ -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)?;
Expand Down
65 changes: 50 additions & 15 deletions crates/edit_prediction_cli/src/retrieve_context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
Expand Down
71 changes: 63 additions & 8 deletions crates/edit_prediction_cli/src/score.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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?;
Expand Down Expand Up @@ -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![];
Expand Down Expand Up @@ -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);

Expand All @@ -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
Expand All @@ -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();

Expand All @@ -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 {
Expand All @@ -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
Expand All @@ -248,14 +292,20 @@ 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;

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;
}

Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down
Loading