mod anthropic_client;mod distill;mod example;mod format_prompt;mod git;mod headless;mod load_project;mod metrics;mod paths;mod predict;mod progress;mod pull_examples;mod reorder_patch;mod retrieve_context;mod score;mod split_commit;mod split_dataset;mod synthesize;use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};use collections::HashSet;use edit_prediction::EditPredictionStore;use futures::channel::mpsc;use futures::{SinkExt as _, StreamExt as _};use gpui::{AppContext as _, Application};use zeta_prompt::ZetaVersion; use reqwest_client::ReqwestClient;use serde::{Deserialize, Serialize};use std::fmt::Display;use std::fs::{File, OpenOptions};use std::hash::{Hash, Hasher};use std::io::{BufRead, BufReader, BufWriter, Write};use std::{path::PathBuf, sync::Arc}; 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, RUN_DIR};use crate::predict::run_prediction;use crate::progress::Progress;use crate::retrieve_context::run_context_retrieval;use crate::score::run_scoring;use crate::split_commit::SplitCommitArgs;use crate::split_dataset::SplitArgs;use crate::synthesize::{SynthesizeConfig, run_synthesize}; #[derive(Parser, Debug)]#[command(name = "ep")]struct EpArgs {    #[arg(long, default_value_t = false)]    printenv: bool,    #[clap(long, default_value_t = 10, global = true)]    max_parallelism: usize,    #[clap(long, global = true)]    limit: Option<usize>,    /// Filter examples by name    #[clap(long, global = true)]    name: Option<String>,    /// Filter examples by repository    #[clap(long, global = true)]    repo: Option<String>,    #[command(subcommand)]    command: Option<Command>,    #[clap(global = true, help = INPUTS_HELP)]    inputs: Vec<PathBuf>,    #[arg(long, short, global = true)]    output: Option<PathBuf>,    #[arg(long, short, global = true)]    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#"Inputs can be file paths or special specifiers:   path      Path to an example(s) file (.md, .json, or .jsonl)   captured-after:{timestamp}      Fetch captured examples from Snowflake after the given RFC3339 timestamp.       You can specify this multiple times and mix it with file inputs.       Required environment variables to connect to Snowflake:          EP_SNOWFLAKE_API_KEY          EP_SNOWFLAKE_BASE_URL       Optional:          EP_SNOWFLAKE_ROLE Examples:   # Predict from a file  ep predict examples.jsonl   # Predict from captured examples after a timestamp  ep predict captured-after:2025-01-01T00:00:00Z   # Mix file inputs and captured-after in the same invocation  ep predict examples.jsonl captured-after:2025-01-01T00:00:00Z"#; #[derive(Subcommand, Debug, Clone)]enum Command {    /// Parse markdown examples and output a combined .jsonl file    ParseExample,    /// Create git worktrees for each example and load file contents    LoadProject,    /// Retrieve context for input examples.    Context,    /// Generate a prompt string for a specific model    FormatPrompt(FormatPromptArgs),    /// Runs edit prediction    Predict(PredictArgs),    /// Computes a score based on actual and expected patches    Score(PredictArgs),    /// Prepares a distillation dataset by copying expected outputs to    /// predicted outputs and removing actual outputs and prompts.    Distill,    /// Print aggregated scores    Eval(PredictArgs),    /// Generate eval examples by analyzing git commits from a repository    Synthesize(SynthesizeArgs),    /// Remove git repositories and worktrees    Clean,    /// Generate an evaluation example by splitting a chronologically-ordered commit    SplitCommit(SplitCommitArgs),    /// Split a JSONL dataset into multiple files (stratified by repository_url if present)    Split(SplitArgs),} impl Display for Command {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        match self {            Command::ParseExample => write!(f, "parse-example"),            Command::LoadProject => write!(f, "load-project"),            Command::Context => write!(f, "context"),            Command::FormatPrompt(format_prompt_args) => write!(                f,                "format-prompt --prompt-format={}",                format_prompt_args                    .provider                    .to_possible_value()                    .unwrap()                    .get_name()            ),            Command::Predict(predict_args) => {                write!(                    f,                    "predict --provider={:?}",                    predict_args                        .provider                        .to_possible_value()                        .unwrap()                        .get_name()                )            }            Command::Score(predict_args) => {                write!(                    f,                    "score --provider={:?}",                    predict_args                        .provider                        .to_possible_value()                        .unwrap()                        .get_name()                )            }            Command::Distill => write!(f, "distill"),            Command::Eval(predict_args) => write!(                f,                "eval --provider={:?}",                predict_args                    .provider                    .to_possible_value()                    .unwrap()                    .get_name()            ),            Command::Synthesize(args) => {                write!(f, "synthesize --repo={}", args.repo)            }            Command::Clean => write!(f, "clean"),            Command::SplitCommit(_) => write!(f, "split-commit"),            Command::Split(_) => write!(f, "split"),        }    }} #[derive(Debug, Args, Clone)]struct FormatPromptArgs {    #[clap(long, short)]    provider: PredictionProvider,    #[clap(        long,        short,        help = "(only for --provider zeta2) A substring of a zeta_prompt::ZetaVersion variant to use",        value_parser = ZetaVersion::parse,        default_value_t = ZetaVersion::default(),    )]    version: ZetaVersion,} #[derive(Debug, Args, Clone)]struct PredictArgs {    #[clap(long, short)]    provider: PredictionProvider,    #[clap(long, default_value_t = 1)]    repetitions: usize,    #[clap(        long,        short,        help = "(only for --provider zeta2) A substring of a zeta_prompt::ZetaVersion variant to use",        value_parser = ZetaVersion::parse,    )]    version: ZetaVersion,} #[derive(Clone, Copy, Debug, PartialEq, ValueEnum, Serialize, Deserialize)]enum PredictionProvider {    Sweep,    Mercury,    Zeta1,    Zeta2,    Teacher,    TeacherNonBatching,} #[derive(Debug, Args, Clone)]struct SynthesizeArgs {    /// Repository URL (git@github.com:owner/repo or https://...)    #[clap(long)]    repo: String,     /// Number of examples to generate    #[clap(long, default_value_t = 5)]    count: usize,     /// Maximum commits to scan before giving up    #[clap(long, default_value_t = 100)]    max_commits: usize,     /// Ignore state file and reprocess all commits    #[clap(long)]    fresh: bool,} impl EpArgs {    fn output_path(&self) -> Option<PathBuf> {        if self.in_place {            if self.inputs.len() == 1 {                self.inputs.first().cloned()            } else {                panic!("--in-place requires exactly one input file")            }        } else {            self.output.clone()        }    }} async fn load_examples(    http_client: Arc<dyn http_client::HttpClient>,    args: &EpArgs,    output_path: Option<&PathBuf>,) -> anyhow::Result<Vec<Example>> {    let mut captured_after_timestamps = Vec::new();    let mut file_inputs = Vec::new();     for input in &args.inputs {        let input_string = input.to_string_lossy();        if let Some(timestamp) = pull_examples::parse_captured_after_input(input_string.as_ref()) {            captured_after_timestamps.push(timestamp.to_string());        } else {            file_inputs.push(input.clone());        }    }     let mut examples = read_example_files(&file_inputs);     Progress::global().set_total_examples(examples.len());     let remaining_limit_for_snowflake =        args.limit.map(|limit| limit.saturating_sub(examples.len()));     if let Some(0) = remaining_limit_for_snowflake {        log::info!(            "skipping captured-after inputs because --limit is already satisfied by example files"        );    } else if !captured_after_timestamps.is_empty() {        captured_after_timestamps.sort();         let max_rows_per_timestamp = remaining_limit_for_snowflake.unwrap_or(5000);         let mut captured_examples = pull_examples::fetch_captured_examples_after(            http_client,            &captured_after_timestamps,            max_rows_per_timestamp,        )        .await?;        examples.append(&mut captured_examples);    }     crate::example::sort_examples_by_repo_and_rev(&mut examples);     if let Some(name_filter) = &args.name {        examples.retain(|example| example.spec.name.contains(name_filter));    }    if let Some(repo_filter) = &args.repo {        examples.retain(|example| example.spec.repository_url.contains(repo_filter));    }     if let Some(limit) = args.limit {        if examples.len() > limit {            examples.truncate(limit);        }    }     if let Some(path) = output_path {        resume_from_output(path, &mut examples);    }     Progress::global().set_total_examples(examples.len());     Ok(examples)} fn spec_hash(spec: &edit_prediction::example_spec::ExampleSpec) -> u64 {    let mut hasher = collections::FxHasher::default();    spec.hash(&mut hasher);    hasher.finish()} fn resume_from_output(path: &PathBuf, examples: &mut Vec<Example>) {    let file = match File::open(path) {        Ok(f) => f,        Err(_) => return,    };     let input_hashes: HashSet<u64> = examples.iter().map(|e| spec_hash(&e.spec)).collect();     let reader = BufReader::new(file);    let mut kept_lines = Vec::new();    let mut kept_hashes = HashSet::default();     for line in reader.lines() {        let line = match line {            Ok(l) => l,            Err(_) => continue,        };         if let Ok(output_example) = serde_json::from_str::<Example>(&line) {            let hash = spec_hash(&output_example.spec);            if input_hashes.contains(&hash) && !kept_hashes.contains(&hash) {                kept_hashes.insert(hash);                kept_lines.push(line);            }        }    }     let total = examples.len();    let already_processed = kept_hashes.len();     eprintln!(        "Resuming: {}/{} examples already processed",        already_processed, total    );     let file = OpenOptions::new()        .write(true)        .truncate(true)        .open(path)        .expect("Failed to open output file for rewriting");    let mut writer = BufWriter::new(file);    for line in &kept_lines {        writeln!(writer, "{}", line).expect("Failed to write to output file");    }    writer.flush().expect("Failed to flush output file");     examples.retain(|e| !kept_hashes.contains(&spec_hash(&e.spec)));} fn main() {    let args = EpArgs::parse();     if args.printenv {        ::util::shell_env::print_env();        return;    }     let output = args.output_path();    let command = match &args.command {        Some(cmd) => cmd.clone(),        None => {            EpArgs::command().print_help().unwrap();            return;        }    };     match &command {        Command::Clean => {            std::fs::remove_dir_all(&*paths::DATA_DIR).unwrap();            return;        }        Command::Synthesize(synth_args) => {            let Some(output_dir) = args.output else {                panic!("output dir is required");            };            let config = SynthesizeConfig {                repo_url: synth_args.repo.clone(),                count: synth_args.count,                max_commits: synth_args.max_commits,                output_dir,                fresh: synth_args.fresh,            };            smol::block_on(async {                if let Err(e) = run_synthesize(config).await {                    eprintln!("Error: {:?}", e);                    std::process::exit(1);                }            });            return;        }        Command::SplitCommit(split_commit_args) => {            if let Err(error) =                split_commit::run_split_commit(split_commit_args, &args.inputs, output.as_ref())            {                eprintln!("{error:#}");                std::process::exit(1);            }            return;        }        Command::Split(split_args) => {            if let Err(error) = split_dataset::run_split(split_args, &args.inputs) {                eprintln!("{error:#}");                std::process::exit(1);            }            return;        }        _ => {}    }     let http_client = Arc::new(ReqwestClient::new());    let app = Application::headless().with_http_client(http_client);     app.run(move |cx| {        let app_state = Arc::new(headless::init(cx));        EditPredictionStore::global(&app_state.client, &app_state.user_store, cx);         cx.spawn(async move |cx| {            let result = async {                let mut examples =                    load_examples(app_state.client.http_client(), &args, output.as_ref()).await?;                 match &command {                    Command::Predict(args) | Command::Score(args) | Command::Eval(args) => {                        predict::sync_batches(&args.provider).await?;                    }                    _ => (),                }                 let failfast_on_single_example = examples.len() == 1;                 let output_sender: Option<mpsc::UnboundedSender<String>> =                    if args.output.is_some() || !matches!(command, Command::Eval(_)) {                        output.as_ref().map(|path| {                            let file = OpenOptions::new()                                .create(true)                                .append(true)                                .open(path)                                .expect("Failed to open output file");                            let mut writer = BufWriter::new(file);                            let (sender, mut receiver) = mpsc::unbounded::<String>();                            cx.background_spawn(async move {                                while let Some(line) = receiver.next().await {                                    writeln!(writer, "{}", line).expect("Failed to write example");                                    writer.flush().expect("Failed to flush output");                                }                            })                            .detach();                            sender                        })                    } else {                        None

Lines 1–500 of 689. Load more