Przeglądaj źródła

fix: resolve only selected provider env refs

main
orangels 1 tydzień temu
rodzic
commit
65ce102004
2 zmienionych plików z 88 dodań i 20 usunięć
  1. +7
    -7
      src/cli.rs
  2. +81
    -13
      src/config.rs

+ 7
- 7
src/cli.rs Wyświetl plik

@@ -105,7 +105,7 @@ pub async fn dispatch(cli: Cli) -> Result<()> {
prompt,
} => {
let config = AppConfig::load(&config)?;
let provider = config.provider(provider.as_deref())?;
let provider = config.resolved_provider(provider.as_deref())?;
let request = ModelRequest {
base_url: provider.base_url.clone(),
api_token: provider.api_token.clone(),
@@ -170,7 +170,7 @@ async fn run_aime_benchmark(
) -> Result<()> {
let config = AppConfig::load(&config_path)?;
let provider_name = provider_name(&config, provider.as_deref())?;
let provider_config = config.provider(Some(&provider_name))?;
let provider_config = config.resolved_provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let loaded = benchmarks::aime::load_cases(Path::new(&config.benchmarks.data_dir))?;
let dataset = dataset_report(
@@ -184,7 +184,7 @@ async fn run_aime_benchmark(
let cases = apply_limit(loaded.cases, limit);
let started_at = Utc::now();
let started = Instant::now();
let base_request = request_template(provider_config, &model, 0.0, 1024);
let base_request = request_template(&provider_config, &model, 0.0, 1024);
let protocol = provider_config.protocol;

let results = stream::iter(cases)
@@ -251,7 +251,7 @@ async fn run_gpqa_benchmark(
) -> Result<()> {
let config = AppConfig::load(&config_path)?;
let provider_name = provider_name(&config, provider.as_deref())?;
let provider_config = config.provider(Some(&provider_name))?;
let provider_config = config.resolved_provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let loaded = benchmarks::gpqa::load_cases(Path::new(&config.benchmarks.data_dir))?;
let dataset = dataset_report(
@@ -265,7 +265,7 @@ async fn run_gpqa_benchmark(
let cases = apply_limit(loaded.cases, limit);
let started_at = Utc::now();
let started = Instant::now();
let base_request = request_template(provider_config, &model, 0.0, 1024);
let base_request = request_template(&provider_config, &model, 0.0, 1024);
let protocol = provider_config.protocol;

let results = stream::iter(cases)
@@ -339,11 +339,11 @@ async fn run_rpm(
let duration = parse_duration(&duration)?;
let config = AppConfig::load(&config_path)?;
let provider_name = provider_name(&config, provider.as_deref())?;
let provider_config = config.provider(Some(&provider_name))?;
let provider_config = config.resolved_provider(Some(&provider_name))?;
let model = model.unwrap_or_else(|| provider_config.default_model.clone());
let request = ModelRequest {
prompt,
..request_template(provider_config, &model, 0.0, 1024)
..request_template(&provider_config, &model, 0.0, 1024)
};
let schedule = rpm_start_schedule(duration, rpm);
let started_at = Utc::now();


+ 81
- 13
src/config.rs Wyświetl plik

@@ -34,7 +34,7 @@ pub enum ProtocolKind {
Anthropic,
}

#[derive(Deserialize)]
#[derive(Clone, Deserialize)]
pub struct ProviderConfig {
pub protocol: ProtocolKind,
pub base_url: String,
@@ -54,6 +54,20 @@ impl fmt::Debug for ProviderConfig {
}
}

impl ProviderConfig {
fn resolved_with_env<F, E>(&self, env: F) -> Result<Self, ConfigError>
where
F: Fn(&str) -> Result<String, E>,
{
Ok(Self {
protocol: self.protocol,
base_url: expand_env_refs_with(&self.base_url, &env)?,
api_token: expand_env_refs_with(&self.api_token, &env)?,
default_model: expand_env_refs_with(&self.default_model, &env)?,
})
}
}

#[derive(Debug, Deserialize)]
pub struct BenchmarkConfig {
#[serde(default = "default_benchmark_data_dir")]
@@ -127,7 +141,7 @@ impl AppConfig {
source,
})?;

config.expand_env_refs_with(env)?;
config.expand_non_provider_env_refs_with(env)?;

Ok(config)
}
@@ -148,7 +162,22 @@ impl AppConfig {
})
}

fn expand_env_refs_with<F, E>(&mut self, env: F) -> Result<(), ConfigError>
pub fn resolved_provider(&self, provider: Option<&str>) -> Result<ProviderConfig, ConfigError> {
self.resolved_provider_with_env(provider, |name| std::env::var(name))
}

fn resolved_provider_with_env<F, E>(
&self,
provider: Option<&str>,
env: F,
) -> Result<ProviderConfig, ConfigError>
where
F: Fn(&str) -> Result<String, E>,
{
self.provider(provider)?.resolved_with_env(env)
}

fn expand_non_provider_env_refs_with<F, E>(&mut self, env: F) -> Result<(), ConfigError>
where
F: Fn(&str) -> Result<String, E>,
{
@@ -156,12 +185,6 @@ impl AppConfig {
*default_provider = expand_env_refs_with(default_provider, &env)?;
}

for provider in self.providers.values_mut() {
provider.base_url = expand_env_refs_with(&provider.base_url, &env)?;
provider.api_token = expand_env_refs_with(&provider.api_token, &env)?;
provider.default_model = expand_env_refs_with(&provider.default_model, &env)?;
}

self.benchmarks.data_dir = expand_env_refs_with(&self.benchmarks.data_dir, &env)?;

if let Some(dataset) = &mut self.benchmarks.aime2026 {
@@ -239,7 +262,12 @@ benchmarks:
)
.expect("load config");

let provider = config.provider(None).expect("default provider");
let provider = config
.resolved_provider_with_env(None, |name| match name {
"LQ_TEST_TOKEN" => Ok("secret-token".to_string()),
_ => Err(()),
})
.expect("default provider");

assert_eq!(provider.api_token, "secret-token");
assert_eq!(provider.default_model, "gpt-test");
@@ -248,7 +276,7 @@ benchmarks:

#[test]
fn rejects_missing_env_token() {
let error = AppConfig::load_from_str_with_env(
let config = AppConfig::load_from_str_with_env(
r#"
default_provider: openai
providers:
@@ -262,11 +290,47 @@ benchmarks:
"#,
|_| Err(()),
)
.expect_err("missing env should fail");
.expect("load config should not expand provider tokens");

let error = config
.resolved_provider_with_env(None, |_| Err(()))
.expect_err("missing env should fail");

assert!(error.to_string().contains("LQ_MISSING_TOKEN"));
}

#[test]
fn resolving_one_provider_does_not_require_other_provider_token() {
let config = AppConfig::load_from_str_with_env(
r#"
default_provider: anthropic
providers:
openai:
protocol: openai
base_url: https://api.openai.test/v1
api_token: ${OPENAI_RELAY_TOKEN}
default_model: gpt-test
anthropic:
protocol: anthropic
base_url: https://api.anthropic.test
api_token: ${ANTHROPIC_RELAY_TOKEN}
default_model: claude-test
"#,
|_| Err(()),
)
.expect("load config should not require provider tokens");

let provider = config
.resolved_provider_with_env(Some("anthropic"), |name| match name {
"ANTHROPIC_RELAY_TOKEN" => Ok("anthropic-secret".to_string()),
_ => Err(()),
})
.expect("resolve anthropic only");

assert_eq!(provider.protocol, ProtocolKind::Anthropic);
assert_eq!(provider.api_token, "anthropic-secret");
}

#[test]
fn load_reads_yaml_from_path_without_env_refs() {
let temp_dir = tempfile::tempdir().expect("create temp dir");
@@ -286,7 +350,11 @@ providers:
.expect("write config");

let config = AppConfig::load(&config_path).expect("load config");
let provider = config.provider(None).expect("default provider");
let provider = config
.resolved_provider_with_env(None, |_| -> Result<String, ()> {
unreachable!("config has no env refs")
})
.expect("default provider");

assert_eq!(provider.api_token, "literal-token");
assert_eq!(provider.protocol, ProtocolKind::Anthropic);


Ładowanie…
Anuluj
Zapisz