use crate::providers::Provider;
use clap::Parser;
#[derive(Parser)]
#[command(name = "every-other-token")]
#[command(version = "4.0.0")]
#[command(about = "A real-time token stream mutator for LLM interpretability research")]
pub struct Args {
#[arg(default_value = "")]
pub prompt: String,
#[arg(default_value = "reverse")]
pub transform: String,
#[arg(default_value = "gpt-3.5-turbo")]
pub model: String,
#[arg(long, value_enum, default_value = "openai")]
pub provider: Provider,
#[arg(long, short)]
pub visual: bool,
#[arg(long)]
pub heatmap: bool,
#[arg(long)]
pub orchestrator: bool,
#[arg(long)]
pub web: bool,
#[arg(long, default_value = "8888")]
pub port: u16,
#[arg(long)]
pub research: bool,
#[arg(long, default_value = "10")]
pub runs: u32,
#[arg(long, default_value = "research_output.json")]
pub output: String,
#[arg(long)]
pub system_a: Option<String>,
#[arg(long, default_value = "5")]
pub top_logprobs: u8,
#[arg(long)]
pub system_b: Option<String>,
#[arg(long)]
pub db: Option<String>,
#[arg(long)]
pub significance: bool,
#[arg(long)]
pub heatmap_export: Option<String>,
#[arg(long, default_value = "0.0")]
pub heatmap_min_confidence: f32,
#[arg(long, default_value = "position")]
pub heatmap_sort_by: String,
#[arg(long)]
pub record: Option<String>,
#[arg(long)]
pub replay: Option<String>,
#[arg(long)]
pub rate: Option<f64>,
#[arg(long)]
pub seed: Option<u64>,
#[arg(long)]
pub log_db: Option<String>,
#[arg(long)]
pub baseline: bool,
#[arg(long)]
pub prompt_file: Option<String>,
#[arg(long)]
pub diff_terminal: bool,
#[arg(long)]
pub json_stream: bool,
#[arg(long, value_name = "SHELL")]
pub completions: Option<clap_complete::Shell>,
#[cfg(feature = "helix-bridge")]
#[arg(long)]
pub helix_url: Option<String>,
#[arg(long)]
pub rate_range: Option<String>,
#[arg(long)]
pub dry_run: bool,
#[arg(long)]
pub template: Option<String>,
#[arg(long)]
pub min_confidence: Option<f64>,
#[arg(long, default_value = "json")]
pub format: String,
#[arg(long, default_value = "5")]
pub collapse_window: usize,
#[arg(long, default_value = "http://localhost:3000")]
pub orchestrator_url: String,
#[arg(long, default_value = "3")]
pub max_retries: u32,
#[arg(long, default_value = "4096")]
pub anthropic_max_tokens: u32,
#[arg(long)]
pub synonym_file: Option<String>,
#[arg(long)]
pub api_key: Option<String>,
#[arg(long, default_value = "1.0")]
pub replay_speed: f64,
#[arg(long, default_value = "120")]
pub timeout: u64,
#[arg(long)]
pub export_timeseries: Option<String>,
#[arg(long)]
pub json_schema: bool,
#[arg(long)]
pub list_models: Option<String>,
#[arg(long)]
pub validate_config: bool,
}
pub fn resolve_model(provider: &Provider, model: &str) -> String {
match provider {
Provider::Anthropic if model == "gpt-3.5-turbo" => "claude-sonnet-4-6".to_string(),
Provider::Mock => "mock-fixture-v1".to_string(),
_ => model.to_string(),
}
}
const KNOWN_OPENAI_MODELS: &[&str] = &[
"gpt-3.5-turbo",
"gpt-3.5-turbo-0125",
"gpt-4",
"gpt-4-turbo",
"gpt-4o",
"gpt-4o-mini",
"gpt-4.1",
"gpt-4.1-mini",
"o1",
"o1-mini",
"o3",
"o3-mini",
];
const KNOWN_ANTHROPIC_MODELS: &[&str] = &[
"claude-3-haiku-20240307",
"claude-3-sonnet-20240229",
"claude-3-opus-20240229",
"claude-3-5-sonnet-20241022",
"claude-3-5-haiku-20241022",
"claude-haiku-4-5-20251001",
"claude-sonnet-4-6",
"claude-opus-4-6",
];
pub fn validate_model(provider: &Provider, model: &str) {
let known: &[&str] = match provider {
Provider::Openai => KNOWN_OPENAI_MODELS,
Provider::Anthropic => KNOWN_ANTHROPIC_MODELS,
Provider::Mock => return,
};
if !known.contains(&model) {
eprintln!(
"[warn] '{}' is not in the known {} model list — verify the model name is correct",
model, provider
);
}
}
pub fn parse_rate_range(s: &str) -> Option<(f64, f64)> {
let sep = s.rfind('-')?;
let min = s[..sep].parse::<f64>().ok()?;
let max = s[sep + 1..].parse::<f64>().ok()?;
if min <= max && min >= 0.0 && max <= 1.0 {
Some((min, max))
} else {
None
}
}
pub fn apply_template(template: &str, prompt: &str) -> String {
template.split("{input}").collect::<Vec<_>>().join(prompt)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolve_model_anthropic_default_swap() {
assert_eq!(
resolve_model(&Provider::Anthropic, "gpt-3.5-turbo"),
"claude-sonnet-4-6"
);
}
#[test]
fn test_resolve_model_anthropic_explicit_model_kept() {
assert_eq!(
resolve_model(&Provider::Anthropic, "claude-haiku-4-5-20251001"),
"claude-haiku-4-5-20251001"
);
}
#[test]
fn test_resolve_model_openai_default_kept() {
assert_eq!(
resolve_model(&Provider::Openai, "gpt-3.5-turbo"),
"gpt-3.5-turbo"
);
}
#[test]
fn test_resolve_model_openai_explicit_model_kept() {
assert_eq!(resolve_model(&Provider::Openai, "gpt-4"), "gpt-4");
}
#[test]
fn test_args_parse_minimal() {
let args = Args::parse_from(["eot", "hello world"]);
assert_eq!(args.prompt, "hello world");
assert_eq!(args.transform, "reverse");
assert_eq!(args.model, "gpt-3.5-turbo");
assert_eq!(args.provider, Provider::Openai);
assert!(!args.visual);
assert!(!args.heatmap);
assert!(!args.orchestrator);
assert!(!args.web);
assert_eq!(args.port, 8888);
assert_eq!(args.top_logprobs, 5);
}
#[test]
fn test_args_parse_full() {
let args = Args::parse_from([
"eot",
"test prompt",
"uppercase",
"gpt-4",
"--provider",
"anthropic",
"--visual",
"--heatmap",
"--orchestrator",
"--web",
"--port",
"9000",
]);
assert_eq!(args.prompt, "test prompt");
assert_eq!(args.transform, "uppercase");
assert_eq!(args.model, "gpt-4");
assert_eq!(args.provider, Provider::Anthropic);
assert!(args.visual);
assert!(args.heatmap);
assert!(args.orchestrator);
assert!(args.web);
assert_eq!(args.port, 9000);
}
#[test]
fn test_args_parse_provider_openai() {
let args = Args::parse_from(["eot", "prompt", "--provider", "openai"]);
assert_eq!(args.provider, Provider::Openai);
}
#[test]
fn test_args_parse_provider_anthropic() {
let args = Args::parse_from(["eot", "prompt", "--provider", "anthropic"]);
assert_eq!(args.provider, Provider::Anthropic);
}
#[test]
fn test_args_parse_short_visual() {
let args = Args::parse_from(["eot", "prompt", "-v"]);
assert!(args.visual);
}
#[test]
fn test_args_default_port() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.port, 8888);
}
#[test]
fn test_args_custom_port() {
let args = Args::parse_from(["eot", "prompt", "--port", "3000"]);
assert_eq!(args.port, 3000);
}
#[test]
fn test_args_research_flag_default_false() {
let args = Args::parse_from(["eot", "prompt"]);
assert!(!args.research);
}
#[test]
fn test_args_research_flag_set() {
let args = Args::parse_from(["eot", "prompt", "--research"]);
assert!(args.research);
}
#[test]
fn test_args_runs_default_one() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.runs, 10);
}
#[test]
fn test_args_runs_custom() {
let args = Args::parse_from(["eot", "prompt", "--runs", "50"]);
assert_eq!(args.runs, 50);
}
#[test]
fn test_args_output_default_none() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.output, "research_output.json");
}
#[test]
fn test_args_output_custom() {
let args = Args::parse_from(["eot", "prompt", "--output", "results.json"]);
assert_eq!(args.output, "results.json");
}
#[test]
fn test_args_system_prompt_default_none() {
let args = Args::parse_from(["eot", "prompt"]);
assert!(args.system_a.is_none());
}
#[test]
fn test_args_system_prompt_set() {
let args = Args::parse_from(["eot", "prompt", "--system-a", "Be concise."]);
assert_eq!(args.system_a.as_deref(), Some("Be concise."));
}
#[test]
fn test_args_research_with_runs_and_output() {
let args = Args::parse_from([
"eot",
"test prompt",
"--research",
"--runs",
"100",
"--output",
"out.json",
]);
assert!(args.research);
assert_eq!(args.runs, 100);
assert_eq!(args.output, "out.json");
}
#[test]
fn test_args_research_does_not_require_web() {
let args = Args::parse_from(["eot", "prompt", "--research"]);
assert!(!args.web);
assert!(args.research);
}
#[test]
fn test_args_parse_research_flag() {
let args = Args::parse_from(["eot", "prompt", "--research"]);
assert!(args.research);
assert_eq!(args.runs, 10);
assert_eq!(args.output, "research_output.json");
}
#[test]
fn test_args_parse_research_custom_runs() {
let args = Args::parse_from(["eot", "prompt", "--research", "--runs", "50"]);
assert_eq!(args.runs, 50);
}
#[test]
fn test_args_parse_research_custom_output() {
let args = Args::parse_from(["eot", "prompt", "--research", "--output", "out.json"]);
assert_eq!(args.output, "out.json");
}
#[test]
fn test_args_parse_system_a() {
let args = Args::parse_from(["eot", "prompt", "--system-a", "Be concise."]);
assert_eq!(args.system_a, Some("Be concise.".to_string()));
}
#[test]
fn test_args_parse_system_b() {
let args = Args::parse_from(["eot", "prompt", "--system-b", "Be verbose."]);
assert_eq!(args.system_b, Some("Be verbose.".to_string()));
}
#[cfg(feature = "helix-bridge")]
#[test]
fn test_args_helix_url_default_none() {
let args = Args::parse_from(["eot", "prompt"]);
assert!(args.helix_url.is_none());
}
#[cfg(feature = "helix-bridge")]
#[test]
fn test_args_helix_url_set() {
let args = Args::parse_from(["eot", "prompt", "--helix-url", "http://127.0.0.1:8080"]);
assert_eq!(args.helix_url.as_deref(), Some("http://127.0.0.1:8080"));
}
#[test]
fn test_parse_rate_range_valid() {
assert_eq!(parse_rate_range("0.3-0.7"), Some((0.3, 0.7)));
}
#[test]
fn test_parse_rate_range_equal() {
assert_eq!(parse_rate_range("0.5-0.5"), Some((0.5, 0.5)));
}
#[test]
fn test_parse_rate_range_invalid() {
assert_eq!(parse_rate_range("invalid"), None);
}
#[test]
fn test_parse_rate_range_min_greater_than_max() {
assert_eq!(parse_rate_range("0.8-0.2"), None);
}
#[test]
fn test_parse_rate_range_scientific_notation() {
let result = parse_rate_range("1e-3-0.5");
assert!(result.is_some());
let (min, max) = result.unwrap();
assert!((min - 0.001).abs() < 1e-9);
assert!((max - 0.5).abs() < 1e-9);
}
#[test]
fn test_parse_rate_range_no_separator_returns_none() {
assert_eq!(parse_rate_range("0.5"), None);
}
#[test]
fn test_apply_template_with_placeholder() {
assert_eq!(apply_template("Answer: {input}", "hello"), "Answer: hello");
}
#[test]
fn test_apply_template_no_placeholder() {
assert_eq!(apply_template("No placeholder", "hello"), "No placeholder");
}
#[test]
fn test_args_dry_run_flag() {
let args = Args::parse_from(["eot", "prompt", "--dry-run"]);
assert!(args.dry_run);
}
#[test]
fn test_args_min_confidence() {
let args = Args::parse_from(["eot", "prompt", "--min-confidence", "0.8"]);
assert_eq!(args.min_confidence, Some(0.8));
}
#[test]
fn test_args_collapse_window() {
let args = Args::parse_from(["eot", "prompt", "--collapse-window", "10"]);
assert_eq!(args.collapse_window, 10);
}
#[test]
fn test_args_format_jsonl() {
let args = Args::parse_from(["eot", "prompt", "--format", "jsonl"]);
assert_eq!(args.format, "jsonl");
}
#[test]
fn test_validate_model_known_openai_no_warn() {
validate_model(&Provider::Openai, "gpt-4");
validate_model(&Provider::Openai, "gpt-3.5-turbo");
validate_model(&Provider::Openai, "gpt-4o");
}
#[test]
fn test_validate_model_known_anthropic_no_warn() {
validate_model(&Provider::Anthropic, "claude-sonnet-4-6");
validate_model(&Provider::Anthropic, "claude-opus-4-6");
validate_model(&Provider::Anthropic, "claude-haiku-4-5-20251001");
}
#[test]
fn test_validate_model_unknown_does_not_panic() {
validate_model(&Provider::Openai, "gpt-9-turbo-ultra");
validate_model(&Provider::Anthropic, "claude-99");
}
#[test]
fn test_validate_model_mock_always_silent() {
validate_model(&Provider::Mock, "any-model-string");
}
#[test]
fn test_known_openai_models_nonempty() {
assert!(!KNOWN_OPENAI_MODELS.is_empty());
}
#[test]
fn test_known_anthropic_models_nonempty() {
assert!(!KNOWN_ANTHROPIC_MODELS.is_empty());
}
#[test]
fn test_known_openai_models_contain_gpt4() {
assert!(KNOWN_OPENAI_MODELS.contains(&"gpt-4"));
}
#[test]
fn test_known_anthropic_models_contain_sonnet() {
assert!(KNOWN_ANTHROPIC_MODELS.contains(&"claude-sonnet-4-6"));
}
#[test]
fn test_apply_template_prompt_with_placeholder_not_reexpanded() {
let result = apply_template("Q: {input}", "what is {input}?");
assert_eq!(result, "Q: what is {input}?");
}
#[test]
fn test_apply_template_multiple_placeholders() {
let result = apply_template("{input} and {input}", "hello");
assert_eq!(result, "hello and hello");
}
#[test]
fn test_args_orchestrator_url_default() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.orchestrator_url, "http://localhost:3000");
}
#[test]
fn test_args_orchestrator_url_custom() {
let args = Args::parse_from([
"eot",
"prompt",
"--orchestrator-url",
"http://10.0.0.1:9000",
]);
assert_eq!(args.orchestrator_url, "http://10.0.0.1:9000");
}
#[test]
fn test_args_max_retries_default() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.max_retries, 3);
}
#[test]
fn test_args_max_retries_custom() {
let args = Args::parse_from(["eot", "prompt", "--max-retries", "5"]);
assert_eq!(args.max_retries, 5);
}
#[test]
fn test_args_max_retries_zero() {
let args = Args::parse_from(["eot", "prompt", "--max-retries", "0"]);
assert_eq!(args.max_retries, 0);
}
#[test]
fn test_args_timeout_default() {
let args = Args::parse_from(["eot", "prompt"]);
assert_eq!(args.timeout, 120);
}
#[test]
fn test_args_timeout_custom() {
let args = Args::parse_from(["eot", "prompt", "--timeout", "60"]);
assert_eq!(args.timeout, 60);
}
#[test]
fn test_args_timeout_zero_disables() {
let args = Args::parse_from(["eot", "prompt", "--timeout", "0"]);
assert_eq!(args.timeout, 0);
}
#[test]
fn test_validate_config_flag_exists() {
let args = Args::parse_from(["eot", "prompt"]);
assert!(!args.validate_config, "validate_config should default to false");
let args2 = Args::parse_from(["eot", "prompt", "--validate-config"]);
assert!(args2.validate_config);
}
#[test]
fn test_list_models_openai_includes_gpt4() {
let openai_models = ["gpt-3.5-turbo", "gpt-4", "gpt-4o", "gpt-4o-mini", "gpt-4-turbo"];
assert!(openai_models.contains(&"gpt-4"), "openai list should include gpt-4");
}
#[test]
fn test_list_models_flag_accepts_openai() {
let args = Args::parse_from(["eot", "prompt", "--list-models", "openai"]);
assert_eq!(args.list_models.as_deref(), Some("openai"));
}
#[test]
fn test_list_models_flag_accepts_all() {
let args = Args::parse_from(["eot", "prompt", "--list-models", "all"]);
assert_eq!(args.list_models.as_deref(), Some("all"));
}
#[test]
fn test_json_schema_flag_outputs_valid_json() {
const RESEARCH_SCHEMA: &str = include_str!("../docs/research-schema.json");
let result = serde_json::from_str::<serde_json::Value>(RESEARCH_SCHEMA);
assert!(result.is_ok(), "research-schema.json must be valid JSON");
}
#[test]
fn test_record_path_unwritable_detected() {
let bad_path = "/nonexistent_dir_eot_test/output.json";
let result = std::fs::OpenOptions::new()
.create(true)
.write(true)
.open(bad_path);
assert!(result.is_err(), "opening a path in a nonexistent dir should fail");
}
#[test]
fn test_export_timeseries_flag_default_none() {
let args = Args::parse_from(["eot", "prompt"]);
assert!(args.export_timeseries.is_none());
}
#[test]
fn test_export_timeseries_flag_set() {
let args = Args::parse_from(["eot", "prompt", "--export-timeseries", "out.csv"]);
assert_eq!(args.export_timeseries.as_deref(), Some("out.csv"));
}
}