| @@ -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<DatasetConfig>, | |||
| #[serde(default)] | |||
| pub gpqa_diamond: Option<DatasetConfig>, | |||
| } | |||
| 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<String>, | |||
| pub providers: HashMap<String, ProviderConfig>, | |||
| #[serde(default)] | |||
| pub benchmarks: BenchmarkConfig, | |||
| } | |||
| impl AppConfig { | |||
| pub fn load(path: &Path) -> Result<Self, ConfigError> { | |||
| 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<String, ConfigError> { | |||
| 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")); | |||
| } | |||
| } | |||