| @@ -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")); | |||||
| } | |||||
| } | |||||