diff --git a/src/cli.rs b/src/cli.rs index fb119e6..a493b61 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -105,7 +105,7 @@ pub async fn dispatch(cli: Cli) -> Result<()> { prompt, } => { let config = AppConfig::load(&config)?; - let provider = config.provider(provider.as_deref())?; + let provider = config.resolved_provider(provider.as_deref())?; let request = ModelRequest { base_url: provider.base_url.clone(), api_token: provider.api_token.clone(), @@ -170,7 +170,7 @@ async fn run_aime_benchmark( ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; - let provider_config = config.provider(Some(&provider_name))?; + let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?; let dataset = dataset_report( @@ -184,7 +184,7 @@ async fn run_aime_benchmark( let cases = apply_limit(loaded.cases, limit); let started_at = Utc::now(); let started = Instant::now(); - let base_request = request_template(provider_config, &model, 0.0, 1024); + let base_request = request_template(&provider_config, &model, 0.0, 1024); let protocol = provider_config.protocol; let results = stream::iter(cases) @@ -251,7 +251,7 @@ async fn run_gpqa_benchmark( ) -> Result<()> { let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; - let provider_config = config.provider(Some(&provider_name))?; + let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?; let dataset = dataset_report( @@ -265,7 +265,7 @@ async fn run_gpqa_benchmark( let cases = apply_limit(loaded.cases, limit); let started_at = Utc::now(); let started = Instant::now(); - let base_request = request_template(provider_config, &model, 0.0, 1024); + let base_request = request_template(&provider_config, &model, 0.0, 1024); let protocol = provider_config.protocol; let results = stream::iter(cases) @@ -339,11 +339,11 @@ async fn run_rpm( let duration = parse_duration(&duration)?; let config = AppConfig::load(&config_path)?; let provider_name = provider_name(&config, provider.as_deref())?; - let provider_config = config.provider(Some(&provider_name))?; + let provider_config = config.resolved_provider(Some(&provider_name))?; let model = model.unwrap_or_else(|| provider_config.default_model.clone()); let request = ModelRequest { prompt, - ..request_template(provider_config, &model, 0.0, 1024) + ..request_template(&provider_config, &model, 0.0, 1024) }; let schedule = rpm_start_schedule(duration, rpm); let started_at = Utc::now(); diff --git a/src/config.rs b/src/config.rs index 722b944..bd8ee18 100644 --- a/src/config.rs +++ b/src/config.rs @@ -34,7 +34,7 @@ pub enum ProtocolKind { Anthropic, } -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] pub struct ProviderConfig { pub protocol: ProtocolKind, pub base_url: String, @@ -54,6 +54,20 @@ impl fmt::Debug for ProviderConfig { } } +impl ProviderConfig { + fn resolved_with_env(&self, env: F) -> Result + where + F: Fn(&str) -> Result, + { + Ok(Self { + protocol: self.protocol, + base_url: expand_env_refs_with(&self.base_url, &env)?, + api_token: expand_env_refs_with(&self.api_token, &env)?, + default_model: expand_env_refs_with(&self.default_model, &env)?, + }) + } +} + #[derive(Debug, Deserialize)] pub struct BenchmarkConfig { #[serde(default = "default_benchmark_data_dir")] @@ -127,7 +141,7 @@ impl AppConfig { source, })?; - config.expand_env_refs_with(env)?; + config.expand_non_provider_env_refs_with(env)?; Ok(config) } @@ -148,7 +162,22 @@ impl AppConfig { }) } - fn expand_env_refs_with(&mut self, env: F) -> Result<(), ConfigError> + pub fn resolved_provider(&self, provider: Option<&str>) -> Result { + self.resolved_provider_with_env(provider, |name| std::env::var(name)) + } + + fn resolved_provider_with_env( + &self, + provider: Option<&str>, + env: F, + ) -> Result + where + F: Fn(&str) -> Result, + { + self.provider(provider)?.resolved_with_env(env) + } + + fn expand_non_provider_env_refs_with(&mut self, env: F) -> Result<(), ConfigError> where F: Fn(&str) -> Result, { @@ -156,12 +185,6 @@ impl AppConfig { *default_provider = expand_env_refs_with(default_provider, &env)?; } - for provider in self.providers.values_mut() { - provider.base_url = expand_env_refs_with(&provider.base_url, &env)?; - provider.api_token = expand_env_refs_with(&provider.api_token, &env)?; - provider.default_model = expand_env_refs_with(&provider.default_model, &env)?; - } - self.benchmarks.data_dir = expand_env_refs_with(&self.benchmarks.data_dir, &env)?; if let Some(dataset) = &mut self.benchmarks.aime2026 { @@ -239,7 +262,12 @@ benchmarks: ) .expect("load config"); - let provider = config.provider(None).expect("default provider"); + let provider = config + .resolved_provider_with_env(None, |name| match name { + "LQ_TEST_TOKEN" => Ok("secret-token".to_string()), + _ => Err(()), + }) + .expect("default provider"); assert_eq!(provider.api_token, "secret-token"); assert_eq!(provider.default_model, "gpt-test"); @@ -248,7 +276,7 @@ benchmarks: #[test] fn rejects_missing_env_token() { - let error = AppConfig::load_from_str_with_env( + let config = AppConfig::load_from_str_with_env( r#" default_provider: openai providers: @@ -262,11 +290,47 @@ benchmarks: "#, |_| Err(()), ) - .expect_err("missing env should fail"); + .expect("load config should not expand provider tokens"); + + let error = config + .resolved_provider_with_env(None, |_| Err(())) + .expect_err("missing env should fail"); assert!(error.to_string().contains("LQ_MISSING_TOKEN")); } + #[test] + fn resolving_one_provider_does_not_require_other_provider_token() { + let config = AppConfig::load_from_str_with_env( + r#" +default_provider: anthropic +providers: + openai: + protocol: openai + base_url: https://api.openai.test/v1 + api_token: ${OPENAI_RELAY_TOKEN} + default_model: gpt-test + anthropic: + protocol: anthropic + base_url: https://api.anthropic.test + api_token: ${ANTHROPIC_RELAY_TOKEN} + default_model: claude-test +"#, + |_| Err(()), + ) + .expect("load config should not require provider tokens"); + + let provider = config + .resolved_provider_with_env(Some("anthropic"), |name| match name { + "ANTHROPIC_RELAY_TOKEN" => Ok("anthropic-secret".to_string()), + _ => Err(()), + }) + .expect("resolve anthropic only"); + + assert_eq!(provider.protocol, ProtocolKind::Anthropic); + assert_eq!(provider.api_token, "anthropic-secret"); + } + #[test] fn load_reads_yaml_from_path_without_env_refs() { let temp_dir = tempfile::tempdir().expect("create temp dir"); @@ -286,7 +350,11 @@ providers: .expect("write config"); let config = AppConfig::load(&config_path).expect("load config"); - let provider = config.provider(None).expect("default provider"); + let provider = config + .resolved_provider_with_env(None, |_| -> Result { + unreachable!("config has no env refs") + }) + .expect("default provider"); assert_eq!(provider.api_token, "literal-token"); assert_eq!(provider.protocol, ProtocolKind::Anthropic);