diff --git a/crates/edit_prediction_cli/src/main.rs b/crates/edit_prediction_cli/src/main.rs index 6e5b28e18f47aaa85a0d11c1754cc9c107cc18ea..819b5dff4090a84819c053751e4de51dbe40856b 100644 --- a/crates/edit_prediction_cli/src/main.rs +++ b/crates/edit_prediction_cli/src/main.rs @@ -35,7 +35,7 @@ use crate::distill::run_distill; use crate::example::{Example, group_examples_by_repo, read_example_files}; use crate::format_prompt::run_format_prompt; use crate::load_project::run_load_project; -use crate::paths::FAILED_EXAMPLES_DIR; +use crate::paths::{FAILED_EXAMPLES_DIR, RUN_DIR}; use crate::predict::run_prediction; use crate::progress::Progress; use crate::retrieve_context::run_context_retrieval; @@ -69,6 +69,21 @@ struct EpArgs { in_place: bool, #[arg(long, short, global = true)] failfast: bool, + /// How to handle failed examples in output: keep them or skip them. + /// Failed examples are always logged to the run's failed directory. + #[arg(long, global = true, default_value = "keep")] + failed: FailedHandling, +} + +/// Controls whether failed examples are included in the main output. +/// Failed examples are always logged to the run's failed/ directory regardless of this setting. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, ValueEnum)] +pub enum FailedHandling { + /// Include failed examples in the main output (default) + #[default] + Keep, + /// Exclude failed examples from the main output + Skip, } const INPUTS_HELP: &str = r#" @@ -189,7 +204,7 @@ impl Display for Command { #[derive(Debug, Args, Clone)] struct FormatPromptArgs { - #[clap(long)] + #[clap(long, short('p'))] prompt_format: PromptFormat, } @@ -533,7 +548,7 @@ fn main() { } .await; - if let Err(error) = result { + let failed = if let Err(error) = result { handle_error( error, &args, @@ -543,18 +558,25 @@ fn main() { example, ) .await; - } - - if let Some(ref mut sender) = output_sender.clone() { - let line = serde_json::to_string(example).unwrap(); - sender - .send(line) - .await - .expect("Failed to send to output writer"); - } else if args.output.is_none() && !matches!(command, Command::Eval(_)) - { - let line = serde_json::to_string(example).unwrap(); - println!("{}", line); + true + } else { + false + }; + + let should_write = !failed || args.failed == FailedHandling::Keep; + if should_write { + if let Some(ref mut sender) = output_sender.clone() { + let line = serde_json::to_string(example).unwrap(); + sender + .send(line) + .await + .expect("Failed to send to output writer"); + } else if args.output.is_none() + && !matches!(command, Command::Eval(_)) + { + let line = serde_json::to_string(example).unwrap(); + println!("{}", line); + } } } }); @@ -615,6 +637,15 @@ async fn handle_error( .await .unwrap(); + let failed_jsonl_path = RUN_DIR.join("failed.jsonl"); + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(&failed_jsonl_path) + .expect("Failed to open failed.jsonl"); + writeln!(file, "{}", serde_json::to_string(example).unwrap()) + .expect("Failed to write to failed.jsonl"); + let cursor_path = example .repo_name() .unwrap() diff --git a/crates/edit_prediction_cli/src/progress.rs b/crates/edit_prediction_cli/src/progress.rs index c1082c848f90a2e812e0572d7093f82729912af7..90080c79eb93892d2a6d32f290f975d91b09579f 100644 --- a/crates/edit_prediction_cli/src/progress.rs +++ b/crates/edit_prediction_cli/src/progress.rs @@ -6,6 +6,8 @@ use std::{ time::{Duration, Instant}, }; +use crate::paths::RUN_DIR; + use log::{Level, Log, Metadata, Record}; pub struct Progress { @@ -413,9 +415,13 @@ impl Progress { } else { 0.0 }; + let failed_jsonl_path = RUN_DIR.join("failed.jsonl"); eprintln!( - "\n{} of {} examples failed ({:.1}%)", - inner.failed_examples, total_examples, percentage + "\n{} of {} examples failed ({:.1}%)\nFailed examples: {}", + inner.failed_examples, + total_examples, + percentage, + failed_jsonl_path.display() ); } }