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 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::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::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). 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 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,} #[derive(Debug, Args, Clone)]struct EvalArgs { #[clap(flatten)] predict: PredictArgs, /// Path to write summary scores as JSON #[clap(long)] summary_json: Option<PathBuf>,} #[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 { Sweep, Mercury, Zeta1, Zeta2(ZetaFormat), Teacher(TeacherBackend), TeacherNonBatching(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::Sweep => write!(f, "sweep"), PredictionProvider::Mercury => write!(f, "mercury"), PredictionProvider::Zeta1 => write!(f, "zeta1"), PredictionProvider::Zeta2(format) => write!(f, "zeta2:{format}"), PredictionProvider::Teacher(backend) => write!(f, "teacher:{backend}"), PredictionProvider::TeacherNonBatching(backend) => { write!(f, "teacher-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() { "sweep" => Ok(PredictionProvider::Sweep), "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-non-batching" | "teacher_non_batching" | "teachernonbatching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherNonBatching(backend)) } "repair" => Ok(PredictionProvider::Repair), _ => { anyhow::bail!( "unknown provider `{provider}`. Valid options: sweep, mercury, zeta1, zeta2, zeta2:<version>, teacher, teacher:<backend>, teacher-non-batching, repair\n\ For zeta2, you can optionally specify a version like `zeta2:ordered` or `zeta2:V0113_Ordered`.\n\ For teacher, you can specify a backend like `teacher: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,} #[derive(Debug, Args, Clone)]struct ImportBatchArgs { /// Batch IDs to import (e.g., msgbatch_xxx for Anthropic, batch_xxx for OpenAI) #[clap(long, required = true, num_args = 1..)] batch_ids: Vec<String>, /// Which provider's batches to import (anthropic or openai) #[clap(long, default_value = "anthropic")] provider: BatchProvider,} #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)]enum BatchProvider { Anthropic, Openai,} 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() } }} /// Minimum Zed version required for Snowflake queries./// This version introduced the current request schema with predicted edits in the edit/// history, and open source repos distinguished.const MIN_CAPTURE_VERSION: pull_examples::MinCaptureVersion = pull_examples::MinCaptureVersion { minor: 224, patch: 1,}; fn deduplicate_examples(examples: &mut Vec<Example>, max_per_cluster: usize) { let total_before_exact = examples.len(); let mut seen_positions = HashSet::default(); examples.retain(|example| seen_positions.insert(example.spec.cursor_position.clone())); log::info!( "exact duplicate filter: {total_before_exact} examples → {} examples ({} removed)",Lines 1–500 of 1368. 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 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::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::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). 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 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, } #[derive(Debug, Args, Clone)] struct EvalArgs { #[clap(flatten)] predict: PredictArgs, /// Path to write summary scores as JSON #[clap(long)] summary_json: Option<PathBuf>, } #[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 { Sweep, Mercury, Zeta1, Zeta2(ZetaFormat), Teacher(TeacherBackend), TeacherNonBatching(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::Sweep => write!(f, "sweep"), PredictionProvider::Mercury => write!(f, "mercury"), PredictionProvider::Zeta1 => write!(f, "zeta1"), PredictionProvider::Zeta2(format) => write!(f, "zeta2:{format}"), PredictionProvider::Teacher(backend) => write!(f, "teacher:{backend}"), PredictionProvider::TeacherNonBatching(backend) => { write!(f, "teacher-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() { "sweep" => Ok(PredictionProvider::Sweep), "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-non-batching" | "teacher_non_batching" | "teachernonbatching" => { let backend = arg .map(|a| a.parse()) .transpose()? .unwrap_or(TeacherBackend::default()); Ok(PredictionProvider::TeacherNonBatching(backend)) } "repair" => Ok(PredictionProvider::Repair), _ => { anyhow::bail!( "unknown provider `{provider}`. Valid options: sweep, mercury, zeta1, zeta2, zeta2:<version>, teacher, teacher:<backend>, teacher-non-batching, repair\n\ For zeta2, you can optionally specify a version like `zeta2:ordered` or `zeta2:V0113_Ordered`.\n\ For teacher, you can specify a backend like `teacher: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, } #[derive(Debug, Args, Clone)] struct ImportBatchArgs { /// Batch IDs to import (e.g., msgbatch_xxx for Anthropic, batch_xxx for OpenAI) #[clap(long, required = true, num_args = 1..)] batch_ids: Vec<String>, /// Which provider's batches to import (anthropic or openai) #[clap(long, default_value = "anthropic")] provider: BatchProvider, } #[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] enum BatchProvider { Anthropic, Openai, } 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() } } } /// Minimum Zed version required for Snowflake queries. /// This version introduced the current request schema with predicted edits in the edit /// history, and open source repos distinguished. const MIN_CAPTURE_VERSION: pull_examples::MinCaptureVersion = pull_examples::MinCaptureVersion { minor: 224, patch: 1, }; fn deduplicate_examples(examples: &mut Vec<Example>, max_per_cluster: usize) { let total_before_exact = examples.len(); let mut seen_positions = HashSet::default(); examples.retain(|example| seen_positions.insert(example.spec.cursor_position.clone())); log::info!( "exact duplicate filter: {total_before_exact} examples → {} examples ({} removed)",