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