ソースを参照

feat: load yaml relay config

main
orangels 1週間前
コミット
b071c99ee6
1個のファイルの変更232行の追加0行の削除
  1. +232
    -0
      src/config.rs

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

読み込み中…
キャンセル
保存