diff --git a/src/config.rs b/src/config.rs index 8b13789..a230f3e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1 +1,233 @@ +use regex::Regex; +use serde::Deserialize; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use thiserror::Error; +#[derive(Debug, Error)] +pub enum ConfigError { + #[error("failed to read config from {path}: {source}")] + Read { + path: PathBuf, + #[source] + source: std::io::Error, + }, + #[error("failed to parse config from {path}: {source}")] + Parse { + path: PathBuf, + #[source] + source: serde_yaml::Error, + }, + #[error("unknown provider: {name}")] + UnknownProvider { name: String }, + #[error("no provider specified and config has no default_provider")] + MissingProvider, + #[error("missing environment variable referenced by config: {name}")] + MissingEnv { name: String }, +} + +#[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ProtocolKind { + Openai, + Anthropic, +} + +#[derive(Debug, Deserialize)] +pub struct ProviderConfig { + pub protocol: ProtocolKind, + pub base_url: String, + pub api_token: String, + pub default_model: String, +} + +#[derive(Debug, Deserialize)] +pub struct BenchmarkConfig { + #[serde(default = "default_benchmark_data_dir")] + pub data_dir: String, + #[serde(default)] + pub aime2026: Option, + #[serde(default)] + pub gpqa_diamond: Option, +} + +impl Default for BenchmarkConfig { + fn default() -> Self { + Self { + data_dir: default_benchmark_data_dir(), + aime2026: None, + gpqa_diamond: None, + } + } +} + +#[derive(Debug, Deserialize)] +pub struct DatasetConfig { + pub source: String, + pub split: String, +} + +#[derive(Debug, Deserialize)] +pub struct AppConfig { + #[serde(default)] + pub default_provider: Option, + pub providers: HashMap, + #[serde(default)] + pub benchmarks: BenchmarkConfig, +} + +impl AppConfig { + pub fn load(path: &Path) -> Result { + let contents = std::fs::read_to_string(path).map_err(|source| ConfigError::Read { + path: path.to_path_buf(), + source, + })?; + let mut config: Self = + serde_yaml::from_str(&contents).map_err(|source| ConfigError::Parse { + path: path.to_path_buf(), + source, + })?; + + config.expand_env_refs()?; + + Ok(config) + } + + pub fn provider(&self, provider: Option<&str>) -> Result<&ProviderConfig, ConfigError> { + let name = match provider { + Some(name) => name, + None => self + .default_provider + .as_deref() + .ok_or(ConfigError::MissingProvider)?, + }; + + self.providers + .get(name) + .ok_or_else(|| ConfigError::UnknownProvider { + name: name.to_string(), + }) + } + + fn expand_env_refs(&mut self) -> Result<(), ConfigError> { + if let Some(default_provider) = &mut self.default_provider { + *default_provider = expand_env_refs(default_provider)?; + } + + for provider in self.providers.values_mut() { + provider.base_url = expand_env_refs(&provider.base_url)?; + provider.api_token = expand_env_refs(&provider.api_token)?; + provider.default_model = expand_env_refs(&provider.default_model)?; + } + + self.benchmarks.data_dir = expand_env_refs(&self.benchmarks.data_dir)?; + + if let Some(dataset) = &mut self.benchmarks.aime2026 { + dataset.expand_env_refs()?; + } + if let Some(dataset) = &mut self.benchmarks.gpqa_diamond { + dataset.expand_env_refs()?; + } + + Ok(()) + } +} + +impl DatasetConfig { + fn expand_env_refs(&mut self) -> Result<(), ConfigError> { + self.source = expand_env_refs(&self.source)?; + self.split = expand_env_refs(&self.split)?; + Ok(()) + } +} + +fn default_benchmark_data_dir() -> String { + "data/benchmarks".to_string() +} + +fn expand_env_refs(value: &str) -> Result { + let regex = Regex::new(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}").expect("valid env regex"); + let mut expanded = String::with_capacity(value.len()); + let mut last_match = 0; + + for captures in regex.captures_iter(value) { + let full_match = captures.get(0).expect("full match exists"); + let env_name = captures.get(1).expect("env name exists").as_str(); + let env_value = std::env::var(env_name).map_err(|_| ConfigError::MissingEnv { + name: env_name.to_string(), + })?; + + expanded.push_str(&value[last_match..full_match.start()]); + expanded.push_str(&env_value); + last_match = full_match.end(); + } + + expanded.push_str(&value[last_match..]); + Ok(expanded) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn loads_provider_and_expands_env_token() { + unsafe { + std::env::set_var("LQ_TEST_TOKEN", "secret-token"); + } + + let temp_dir = tempfile::tempdir().expect("create temp dir"); + let config_path = temp_dir.path().join("config.yaml"); + std::fs::write( + &config_path, + r#" +default_provider: openai +providers: + openai: + protocol: openai + base_url: https://api.openai.test/v1 + api_token: ${LQ_TEST_TOKEN} + default_model: gpt-test +benchmarks: + data_dir: data/benchmarks +"#, + ) + .expect("write config"); + + let config = AppConfig::load(&config_path).expect("load config"); + let provider = config.provider(None).expect("default provider"); + + assert_eq!(provider.api_token, "secret-token"); + assert_eq!(provider.default_model, "gpt-test"); + assert_eq!(provider.protocol, ProtocolKind::Openai); + } + + #[test] + fn rejects_missing_env_token() { + unsafe { + std::env::remove_var("LQ_MISSING_TOKEN"); + } + + let temp_dir = tempfile::tempdir().expect("create temp dir"); + let config_path = temp_dir.path().join("config.yaml"); + std::fs::write( + &config_path, + r#" +default_provider: openai +providers: + openai: + protocol: openai + base_url: https://api.openai.test/v1 + api_token: ${LQ_MISSING_TOKEN} + default_model: gpt-test +benchmarks: + data_dir: data/benchmarks +"#, + ) + .expect("write config"); + + let error = AppConfig::load(&config_path).expect_err("missing env should fail"); + + assert!(error.to_string().contains("LQ_MISSING_TOKEN")); + } +}