mod anthropic_client;mod distill;mod example;mod filter_languages;mod format_prompt;mod git;mod headless; mod load_project;mod metrics;mod openai_client;mod parse_output;mod paths;mod predict;mod progress;mod prompt_assets;mod pull_examples;mod qa;mod reorder_patch;mod repair;mod retrieve_context;mod reversal_tracking;mod score;mod split_commit;mod split_dataset; mod synthesize;mod truncate_expected_patch;mod word_diff;use anyhow::Context as _;use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum};use collections::{HashMap, HashSet};use edit_prediction::EditPredictionStore;use futures::channel::mpsc;use futures::{SinkExt as _, StreamExt as _};use gaoya::minhash::{ MinHashIndex, MinHasher, MinHasher32, calculate_minhash_params, compute_minhash_similarity,};use gpui::{AppContext as _, BackgroundExecutor, Task};use zeta_prompt::ZetaFormat; use reqwest_client::ReqwestClient;use serde::{Deserialize, Deserializer, Serialize, Serializer};use std::env;use std::fmt::Display;use std::fs::{File, OpenOptions};use std::hash::{Hash, Hasher};use std::io::{BufRead, BufReader, BufWriter, Write};use std::sync::Mutex;use std::{path::PathBuf, sync::Arc}; use crate::distill::run_distill;use crate::example::{Example, group_examples_by_repo, read_example_files};use crate::filter_languages::{FilterLanguagesArgs, run_filter_languages};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::pull_examples::{fetch_settled_examples_after, parse_settled_after_input};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};use crate::truncate_expected_patch::TruncatePatchArgs; #[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, /// The limit for the number of examples to process /// Default is unlimited for processing local datasets, 5000 when pulling from snowflake #[clap(long, global = true)] limit: Option<usize>, #[clap(long, global = true)] offset: Option<usize>, /// Filter examples by name #[clap(long, global = true)] name: Option<String>, /// Filter examples by repository #[clap(long, global = true)] repo: Option<String>, /// Deduplicate by cursor position and keep at most this many examples per cluster #[clap(long, global = true)] max_duplicates: Option<usize>, #[command(subcommand)] command: Option<Command>, /// Input file paths #[clap(global = true)] 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, /// Output as markdown files instead of JSONL. When set, -o specifies a directory /// where one .md file per example will be written (named after each example). #[arg(long, short, global = true)] markdown: bool,} /// 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, /// Skip writing files SkipNoFiles,} 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. These are examples captured via the "Capture Edit Prediction Example" action. rejected-after:{timestamp} Fetch rejected edit predictions from Snowflake after the given RFC3339 timestamp. These are predictions that were shown to users but rejected (useful for DPO training). settled-after:{timestamp} Fetch settled stream examples from Snowflake after the given RFC3339 timestamp. These are examples from the edit prediction settled stream. rated-after:{timestamp} Fetch user-rated edit predictions from Snowflake after the given RFC3339 timestamp. These are predictions that users explicitly rated as positive or negative via the rate completions modal. Only zeta2 predictions are included. - Positive ratings: output becomes expected_patches - Negative ratings: output becomes rejected_patch rated-positive-after:{timestamp} Same as rated-after, but only fetches positively rated predictions. rated-negative-after:{timestamp} Same as rated-after, but only fetches negatively rated predictions. Required environment variables to connect to Snowflake: EP_SNOWFLAKE_API_KEY EP_SNOWFLAKE_BASE_URL Optional: EP_SNOWFLAKE_ROLE Examples: # Read examples from a file ep read examples.jsonl -o output.jsonl # Read captured examples after a timestamp ep read captured-after:2025-01-01T00:00:00Z -o captured.jsonl # Read rejected predictions for DPO training ep read rejected-after:2025-01-01T00:00:00Z -o rejected.jsonl # Read user-rated predictions ep read rated-after:2025-01-01T00:00:00Z -o rated.jsonl # Read settled stream examples ep read settled-after:2025-01-01T00:00:00Z -o settled.jsonl # Read only positively rated predictions ep read rated-positive-after:2025-01-01T00:00:00Z -o positive.jsonl # Read only negatively rated predictions ep read rated-negative-after:2025-01-01T00:00:00Z -o negative.jsonl # Mix multiple input sources ep predict examples.jsonl captured-after:2025-01-01T00:00:00Z"#; #[derive(Subcommand, Debug, Clone)]enum Command { /// Read examples from files or fetch from Snowflake, output as .jsonl Read(ReadArgs), /// 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), /// Parse model outputs (actual_output) into unified diffs (actual_patch). /// Requires format-prompt to have been run first. Uses provider from prompt. ParseOutput, /// 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(EvalArgs), /// 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), /// Truncate expected patch by the given criteria TruncatePatch(TruncatePatchArgs), /// Split a JSONL dataset into multiple files (stratified by repository_url if present) Split(SplitArgs), /// Filter a JSONL dataset by programming language (based on cursor_path extension) FilterLanguages(FilterLanguagesArgs), /// Import Anthropic batch results by batch IDs (useful for recovering after database loss) ImportBatch(ImportBatchArgs), /// Assess the quality of predictions using LLM-as-a-judge Qa(qa::QaArgs), /// Repair predictions that received poor QA scores by generating improved predictions Repair(repair::RepairArgs), /// Print all valid zeta formats (lowercase, one per line) PrintZetaFormats,} impl Display for Command { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Command::Read(_) => write!(f, "read"), Command::LoadProject => write!(f, "load-project"), Command::Context => write!(f, "context"), Command::FormatPrompt(args) => { write!(f, "format-prompt --provider={}", args.provider) } Command::Predict(args) => match &args.provider { Some(provider) => write!(f, "predict --provider={}", provider), None => write!(f, "predict"), }, Command::ParseOutput => write!(f, "parse-output"), Command::Score(args) => match &args.provider { Some(provider) => write!(f, "score --provider={}", provider), None => write!(f, "score"), }, Command::Distill => write!(f, "distill"), Command::Eval(args) => match &args.predict.provider { Some(provider) => write!(f, "eval --provider={}", provider), None => write!(f, "eval"), }, Command::Synthesize(args) => { write!(f, "synthesize --repos {}", args.repos.join(" ")) } Command::Clean => write!(f, "clean"), Command::SplitCommit(_) => write!(f, "split-commit"), Command::TruncatePatch(_) => write!(f, "truncate-patch"), Command::Split(_) => write!(f, "split"), Command::FilterLanguages(_) => write!(f, "filter-languages"), Command::ImportBatch(args) => { write!(f, "import-batch --batch-ids {}", args.batch_ids.join(" ")) } Command::Qa(_) => { write!(f, "qa") } Command::Repair(_) => { write!(f, "repair") } Command::PrintZetaFormats => { write!(f, "print-zeta-formats") } } }} #[derive(Debug, Args, Clone)]#[command(after_help = INPUTS_HELP)]struct ReadArgs {} #[derive(Debug, Args, Clone)]struct FormatPromptArgs { #[clap(long, short('p'), default_value_t = PredictionProvider::default())] provider: PredictionProvider,} #[derive(Debug, Args, Clone)]struct PredictArgs { #[clap(long, short('p'))] provider: Option<PredictionProvider>, #[clap(long, default_value_t = 1)] repetitions: usize, /// Only use cached responses, don't queue new requests for batching #[clap(long)] cache_only: bool, /// Wait for all batches to complete before exiting (only applies to batched providers like teacher) #[clap(long)] wait: bool,} #[derive(Debug, Args, Clone)]struct EvalArgs { #[clap(flatten)] predict: PredictArgs, /// Path to write summary scores as JSON #[clap(long)] summary_json: Option<PathBuf>, /// Print all individual example lines (default: up to 20) #[clap(long)] verbose: bool,} #[derive(Clone, Copy, Default, Debug, PartialEq, Eq, Hash)]pub enum TeacherBackend { Sonnet46, #[default] Sonnet45, Gpt52,} impl std::fmt::Display for TeacherBackend { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { TeacherBackend::Sonnet46 => write!(f, "sonnet46"), TeacherBackend::Sonnet45 => write!(f, "sonnet45"), TeacherBackend::Gpt52 => write!(f, "gpt52"), } }} impl std::str::FromStr for TeacherBackend { type Err = anyhow::Error; fn from_str(s: &str) -> Result<Self, Self::Err> { match s.to_lowercase().as_str() { "sonnet45" | "sonnet" | "claude" => Ok(TeacherBackend::Sonnet45), "sonnet46" => Ok(TeacherBackend::Sonnet46), "gpt52" | "gpt" | "openai" => Ok(TeacherBackend::Gpt52), "v0114180editableregion" => Ok(TeacherBackend::Sonnet45), _ => anyhow::bail!( "unknown teacher backend `{s}`. Valid options: sonnet45, sonnet46, gpt52" ), } }} impl TeacherBackend { pub fn model_name(&self) -> &'static str { match self { TeacherBackend::Sonnet45 => "claude-sonnet-4-5", TeacherBackend::Sonnet46 => "claude-sonnet-4-6", TeacherBackend::Gpt52 => "gpt-5.2", } }} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]enum PredictionProvider { Mercury, Zeta1, Zeta2(ZetaFormat), Baseten(ZetaFormat), Teacher(TeacherBackend), TeacherMultiRegion(TeacherBackend), TeacherNonBatching(TeacherBackend), TeacherMultiRegionNonBatching(TeacherBackend), Repair,} impl Default for PredictionProvider { fn default() -> Self { PredictionProvider::Zeta2(ZetaFormat::default()) }} impl std::fmt::Display for PredictionProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { PredictionProvider::Mercury => write!(f, "mercury"), PredictionProvider::Zeta1 => write!(f, "zeta1"), PredictionProvider::Zeta2(format) => write!(f, "zeta2:{format}"), PredictionProvider::Baseten(format) => write!(f, "baseten:{format}"), PredictionProvider::Teacher(backend) => write!(f, "teacher:{backend}"), PredictionProvider::TeacherMultiRegion(backend) => { write!(f, "teacher-multi-region:{backend}") } PredictionProvider::TeacherNonBatching(backend) => { write!(f, "teacher-non-batching:{backend}") } PredictionProvider::TeacherMultiRegionNonBatching(backend) => { write!(f, "teacher-multi-region-non-batching:{backend}") } PredictionProvider::Repair => write!(f, "repair"), } }} impl std::str::FromStr for PredictionProvider { type Err = anyhow::Error; fn from_str(s: &str) -> Result<Self, Self::Err> { let (provider, arg) = s.split_once(':').map_or((s, None), |(p, a)| (p, Some(a))); let provider_lower = provider.to_lowercase(); match provider_lower.as_str() { "mercury" => Ok(PredictionProvider::Mercury), "zeta1" => Ok(PredictionProvider::Zeta1), "zeta2" => { let format = arg.map(ZetaFormat::parse).transpose()?.unwrap_or_default(); Ok(PredictionProvider::Zeta2(format)) } "teacher" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::Teacher(backend)) } "teacher-multi-region" | "teacher_multi_region" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherMultiRegion(backend)) } "teacher-non-batching" | "teacher_non_batching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherNonBatching(backend)) } "teacher-multi-region-non-batching" | "teacher_multi_region_non_batching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherMultiRegionNonBatching(backend)) } "repair" => Ok(PredictionProvider::Repair), "baseten" => { let format = arg .map(ZetaFormat::parse) .transpose()? .unwrap_or(ZetaFormat::default()); Ok(PredictionProvider::Baseten(format)) } _ => { anyhow::bail!( "unknown provider `{provider}`. Valid options: mercury, zeta1, zeta2, zeta2:<version>, teacher, teacher:<backend>, teacher-multi-region, teacher-multi-region:<backend>, teacher-non-batching, teacher-multi-region-non-batching, repair\n\ For zeta2, you can optionally specify a version like `zeta2:ordered` or `zeta2:V0113_Ordered`.\n\ For teacher providers, you can specify a backend like `teacher:sonnet46`, `teacher-multi-region:sonnet46`, `teacher-multi-region-non-batching:sonnet46`, or `teacher:gpt52`.\n\ Available zeta versions:\n{}", ZetaFormat::options_as_string() ) } } }} impl Serialize for PredictionProvider { fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer, { serializer.serialize_str(&self.to_string()) }} impl<'de> Deserialize<'de> for PredictionProvider { fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> where D: Deserializer<'de>, { let s = String::deserialize(deserializer)?; s.parse().map_err(serde::de::Error::custom) }} #[derive(Debug, Args, Clone)]struct SynthesizeArgs { /// Repository URLs (git@github.com:owner/repo or https://...) #[clap(long, required = true, num_args = 1..)] repos: Vec<String>, /// Number of examples to generate per repository #[clap(long, default_value_t = 5)] count: usize, /// Maximum commits to scan per repository before giving up #[clap(long, default_value_t = 100)] max_commits: usize, /// Ignore state file and reprocess all commits #[clap(long)] fresh: bool,}Lines 1–500 of 1537. Load more
mod anthropic_client; mod distill; mod example; mod filter_languages; mod format_prompt; mod git; mod headless; mod load_project; mod metrics; mod openai_client; mod parse_output; mod paths; mod predict; mod progress; mod prompt_assets; mod pull_examples; mod qa; mod reorder_patch; mod repair; mod retrieve_context; mod reversal_tracking; mod score; mod split_commit; mod split_dataset; mod synthesize; mod truncate_expected_patch; mod word_diff; use anyhow::Context as _; use clap::{Args, CommandFactory, Parser, Subcommand, ValueEnum}; use collections::{HashMap, HashSet}; use edit_prediction::EditPredictionStore; use futures::channel::mpsc; use futures::{SinkExt as _, StreamExt as _}; use gaoya::minhash::{ MinHashIndex, MinHasher, MinHasher32, calculate_minhash_params, compute_minhash_similarity, }; use gpui::{AppContext as _, BackgroundExecutor, Task}; use zeta_prompt::ZetaFormat; use reqwest_client::ReqwestClient; use serde::{Deserialize, Deserializer, Serialize, Serializer}; use std::env; use std::fmt::Display; use std::fs::{File, OpenOptions}; use std::hash::{Hash, Hasher}; use std::io::{BufRead, BufReader, BufWriter, Write}; use std::sync::Mutex; use std::{path::PathBuf, sync::Arc}; use crate::distill::run_distill; use crate::example::{Example, group_examples_by_repo, read_example_files}; use crate::filter_languages::{FilterLanguagesArgs, run_filter_languages}; 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::pull_examples::{fetch_settled_examples_after, parse_settled_after_input}; 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}; use crate::truncate_expected_patch::TruncatePatchArgs; #[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, /// The limit for the number of examples to process /// Default is unlimited for processing local datasets, 5000 when pulling from snowflake #[clap(long, global = true)] limit: Option<usize>, #[clap(long, global = true)] offset: Option<usize>, /// Filter examples by name #[clap(long, global = true)] name: Option<String>, /// Filter examples by repository #[clap(long, global = true)] repo: Option<String>, /// Deduplicate by cursor position and keep at most this many examples per cluster #[clap(long, global = true)] max_duplicates: Option<usize>, #[command(subcommand)] command: Option<Command>, /// Input file paths #[clap(global = true)] 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, /// Output as markdown files instead of JSONL. When set, -o specifies a directory /// where one .md file per example will be written (named after each example). #[arg(long, short, global = true)] markdown: bool, } /// 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, /// Skip writing files SkipNoFiles, } 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. These are examples captured via the "Capture Edit Prediction Example" action. rejected-after:{timestamp} Fetch rejected edit predictions from Snowflake after the given RFC3339 timestamp. These are predictions that were shown to users but rejected (useful for DPO training). settled-after:{timestamp} Fetch settled stream examples from Snowflake after the given RFC3339 timestamp. These are examples from the edit prediction settled stream. rated-after:{timestamp} Fetch user-rated edit predictions from Snowflake after the given RFC3339 timestamp. These are predictions that users explicitly rated as positive or negative via the rate completions modal. Only zeta2 predictions are included. - Positive ratings: output becomes expected_patches - Negative ratings: output becomes rejected_patch rated-positive-after:{timestamp} Same as rated-after, but only fetches positively rated predictions. rated-negative-after:{timestamp} Same as rated-after, but only fetches negatively rated predictions. Required environment variables to connect to Snowflake: EP_SNOWFLAKE_API_KEY EP_SNOWFLAKE_BASE_URL Optional: EP_SNOWFLAKE_ROLE Examples: # Read examples from a file ep read examples.jsonl -o output.jsonl # Read captured examples after a timestamp ep read captured-after:2025-01-01T00:00:00Z -o captured.jsonl # Read rejected predictions for DPO training ep read rejected-after:2025-01-01T00:00:00Z -o rejected.jsonl # Read user-rated predictions ep read rated-after:2025-01-01T00:00:00Z -o rated.jsonl # Read settled stream examples ep read settled-after:2025-01-01T00:00:00Z -o settled.jsonl # Read only positively rated predictions ep read rated-positive-after:2025-01-01T00:00:00Z -o positive.jsonl # Read only negatively rated predictions ep read rated-negative-after:2025-01-01T00:00:00Z -o negative.jsonl # Mix multiple input sources ep predict examples.jsonl captured-after:2025-01-01T00:00:00Z "#; #[derive(Subcommand, Debug, Clone)] enum Command { /// Read examples from files or fetch from Snowflake, output as .jsonl Read(ReadArgs), /// 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), /// Parse model outputs (actual_output) into unified diffs (actual_patch). /// Requires format-prompt to have been run first. Uses provider from prompt. ParseOutput, /// 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(EvalArgs), /// 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), /// Truncate expected patch by the given criteria TruncatePatch(TruncatePatchArgs), /// Split a JSONL dataset into multiple files (stratified by repository_url if present) Split(SplitArgs), /// Filter a JSONL dataset by programming language (based on cursor_path extension) FilterLanguages(FilterLanguagesArgs), /// Import Anthropic batch results by batch IDs (useful for recovering after database loss) ImportBatch(ImportBatchArgs), /// Assess the quality of predictions using LLM-as-a-judge Qa(qa::QaArgs), /// Repair predictions that received poor QA scores by generating improved predictions Repair(repair::RepairArgs), /// Print all valid zeta formats (lowercase, one per line) PrintZetaFormats, } impl Display for Command { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Command::Read(_) => write!(f, "read"), Command::LoadProject => write!(f, "load-project"), Command::Context => write!(f, "context"), Command::FormatPrompt(args) => { write!(f, "format-prompt --provider={}", args.provider) } Command::Predict(args) => match &args.provider { Some(provider) => write!(f, "predict --provider={}", provider), None => write!(f, "predict"), }, Command::ParseOutput => write!(f, "parse-output"), Command::Score(args) => match &args.provider { Some(provider) => write!(f, "score --provider={}", provider), None => write!(f, "score"), }, Command::Distill => write!(f, "distill"), Command::Eval(args) => match &args.predict.provider { Some(provider) => write!(f, "eval --provider={}", provider), None => write!(f, "eval"), }, Command::Synthesize(args) => { write!(f, "synthesize --repos {}", args.repos.join(" ")) } Command::Clean => write!(f, "clean"), Command::SplitCommit(_) => write!(f, "split-commit"), Command::TruncatePatch(_) => write!(f, "truncate-patch"), Command::Split(_) => write!(f, "split"), Command::FilterLanguages(_) => write!(f, "filter-languages"), Command::ImportBatch(args) => { write!(f, "import-batch --batch-ids {}", args.batch_ids.join(" ")) } Command::Qa(_) => { write!(f, "qa") } Command::Repair(_) => { write!(f, "repair") } Command::PrintZetaFormats => { write!(f, "print-zeta-formats") } } } } #[derive(Debug, Args, Clone)] #[command(after_help = INPUTS_HELP)] struct ReadArgs {} #[derive(Debug, Args, Clone)] struct FormatPromptArgs { #[clap(long, short('p'), default_value_t = PredictionProvider::default())] provider: PredictionProvider, } #[derive(Debug, Args, Clone)] struct PredictArgs { #[clap(long, short('p'))] provider: Option<PredictionProvider>, #[clap(long, default_value_t = 1)] repetitions: usize, /// Only use cached responses, don't queue new requests for batching #[clap(long)] cache_only: bool, /// Wait for all batches to complete before exiting (only applies to batched providers like teacher) #[clap(long)] wait: bool, } #[derive(Debug, Args, Clone)] struct EvalArgs { #[clap(flatten)] predict: PredictArgs, /// Path to write summary scores as JSON #[clap(long)] summary_json: Option<PathBuf>, /// Print all individual example lines (default: up to 20) #[clap(long)] verbose: bool, } #[derive(Clone, Copy, Default, Debug, PartialEq, Eq, Hash)] pub enum TeacherBackend { Sonnet46, #[default] Sonnet45, Gpt52, } impl std::fmt::Display for TeacherBackend { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { TeacherBackend::Sonnet46 => write!(f, "sonnet46"), TeacherBackend::Sonnet45 => write!(f, "sonnet45"), TeacherBackend::Gpt52 => write!(f, "gpt52"), } } } impl std::str::FromStr for TeacherBackend { type Err = anyhow::Error; fn from_str(s: &str) -> Result<Self, Self::Err> { match s.to_lowercase().as_str() { "sonnet45" | "sonnet" | "claude" => Ok(TeacherBackend::Sonnet45), "sonnet46" => Ok(TeacherBackend::Sonnet46), "gpt52" | "gpt" | "openai" => Ok(TeacherBackend::Gpt52), "v0114180editableregion" => Ok(TeacherBackend::Sonnet45), _ => anyhow::bail!( "unknown teacher backend `{s}`. Valid options: sonnet45, sonnet46, gpt52" ), } } } impl TeacherBackend { pub fn model_name(&self) -> &'static str { match self { TeacherBackend::Sonnet45 => "claude-sonnet-4-5", TeacherBackend::Sonnet46 => "claude-sonnet-4-6", TeacherBackend::Gpt52 => "gpt-5.2", } } } #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] enum PredictionProvider { Mercury, Zeta1, Zeta2(ZetaFormat), Baseten(ZetaFormat), Teacher(TeacherBackend), TeacherMultiRegion(TeacherBackend), TeacherNonBatching(TeacherBackend), TeacherMultiRegionNonBatching(TeacherBackend), Repair, } impl Default for PredictionProvider { fn default() -> Self { PredictionProvider::Zeta2(ZetaFormat::default()) } } impl std::fmt::Display for PredictionProvider { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { PredictionProvider::Mercury => write!(f, "mercury"), PredictionProvider::Zeta1 => write!(f, "zeta1"), PredictionProvider::Zeta2(format) => write!(f, "zeta2:{format}"), PredictionProvider::Baseten(format) => write!(f, "baseten:{format}"), PredictionProvider::Teacher(backend) => write!(f, "teacher:{backend}"), PredictionProvider::TeacherMultiRegion(backend) => { write!(f, "teacher-multi-region:{backend}") } PredictionProvider::TeacherNonBatching(backend) => { write!(f, "teacher-non-batching:{backend}") } PredictionProvider::TeacherMultiRegionNonBatching(backend) => { write!(f, "teacher-multi-region-non-batching:{backend}") } PredictionProvider::Repair => write!(f, "repair"), } } } impl std::str::FromStr for PredictionProvider { type Err = anyhow::Error; fn from_str(s: &str) -> Result<Self, Self::Err> { let (provider, arg) = s.split_once(':').map_or((s, None), |(p, a)| (p, Some(a))); let provider_lower = provider.to_lowercase(); match provider_lower.as_str() { "mercury" => Ok(PredictionProvider::Mercury), "zeta1" => Ok(PredictionProvider::Zeta1), "zeta2" => { let format = arg.map(ZetaFormat::parse).transpose()?.unwrap_or_default(); Ok(PredictionProvider::Zeta2(format)) } "teacher" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::Teacher(backend)) } "teacher-multi-region" | "teacher_multi_region" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherMultiRegion(backend)) } "teacher-non-batching" | "teacher_non_batching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherNonBatching(backend)) } "teacher-multi-region-non-batching" | "teacher_multi_region_non_batching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherMultiRegionNonBatching(backend)) } "repair" => Ok(PredictionProvider::Repair), "baseten" => { let format = arg .map(ZetaFormat::parse) .transpose()? .unwrap_or(ZetaFormat::default()); Ok(PredictionProvider::Baseten(format)) } _ => { anyhow::bail!( "unknown provider `{provider}`. Valid options: mercury, zeta1, zeta2, zeta2:<version>, teacher, teacher:<backend>, teacher-multi-region, teacher-multi-region:<backend>, teacher-non-batching, teacher-multi-region-non-batching, repair\n\ For zeta2, you can optionally specify a version like `zeta2:ordered` or `zeta2:V0113_Ordered`.\n\ For teacher providers, you can specify a backend like `teacher:sonnet46`, `teacher-multi-region:sonnet46`, `teacher-multi-region-non-batching:sonnet46`, or `teacher:gpt52`.\n\ Available zeta versions:\n{}", ZetaFormat::options_as_string() ) } } } } impl Serialize for PredictionProvider { fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer, { serializer.serialize_str(&self.to_string()) } } impl<'de> Deserialize<'de> for PredictionProvider { fn deserialize<D>(deserializer: D) -> Result<Self, D::Error> where D: Deserializer<'de>, { let s = String::deserialize(deserializer)?; s.parse().map_err(serde::de::Error::custom) } } #[derive(Debug, Args, Clone)] struct SynthesizeArgs { /// Repository URLs (git@github.com:owner/repo or https://...) #[clap(long, required = true, num_args = 1..)] repos: Vec<String>, /// Number of examples to generate per repository #[clap(long, default_value_t = 5)] count: usize, /// Maximum commits to scan per repository before giving up #[clap(long, default_value_t = 100)] max_commits: usize, /// Ignore state file and reprocess all commits #[clap(long)] fresh: bool, }