api模型检测
選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。

234 行
6.4 KiB

  1. use regex::Regex;
  2. use serde::Deserialize;
  3. use std::collections::HashMap;
  4. use std::path::{Path, PathBuf};
  5. use thiserror::Error;
  6. #[derive(Debug, Error)]
  7. pub enum ConfigError {
  8. #[error("failed to read config from {path}: {source}")]
  9. Read {
  10. path: PathBuf,
  11. #[source]
  12. source: std::io::Error,
  13. },
  14. #[error("failed to parse config from {path}: {source}")]
  15. Parse {
  16. path: PathBuf,
  17. #[source]
  18. source: serde_yaml::Error,
  19. },
  20. #[error("unknown provider: {name}")]
  21. UnknownProvider { name: String },
  22. #[error("no provider specified and config has no default_provider")]
  23. MissingProvider,
  24. #[error("missing environment variable referenced by config: {name}")]
  25. MissingEnv { name: String },
  26. }
  27. #[derive(Debug, Clone, Copy, Deserialize, PartialEq, Eq)]
  28. #[serde(rename_all = "snake_case")]
  29. pub enum ProtocolKind {
  30. Openai,
  31. Anthropic,
  32. }
  33. #[derive(Debug, Deserialize)]
  34. pub struct ProviderConfig {
  35. pub protocol: ProtocolKind,
  36. pub base_url: String,
  37. pub api_token: String,
  38. pub default_model: String,
  39. }
  40. #[derive(Debug, Deserialize)]
  41. pub struct BenchmarkConfig {
  42. #[serde(default = "default_benchmark_data_dir")]
  43. pub data_dir: String,
  44. #[serde(default)]
  45. pub aime2026: Option<DatasetConfig>,
  46. #[serde(default)]
  47. pub gpqa_diamond: Option<DatasetConfig>,
  48. }
  49. impl Default for BenchmarkConfig {
  50. fn default() -> Self {
  51. Self {
  52. data_dir: default_benchmark_data_dir(),
  53. aime2026: None,
  54. gpqa_diamond: None,
  55. }
  56. }
  57. }
  58. #[derive(Debug, Deserialize)]
  59. pub struct DatasetConfig {
  60. pub source: String,
  61. pub split: String,
  62. }
  63. #[derive(Debug, Deserialize)]
  64. pub struct AppConfig {
  65. #[serde(default)]
  66. pub default_provider: Option<String>,
  67. pub providers: HashMap<String, ProviderConfig>,
  68. #[serde(default)]
  69. pub benchmarks: BenchmarkConfig,
  70. }
  71. impl AppConfig {
  72. pub fn load(path: &Path) -> Result<Self, ConfigError> {
  73. let contents = std::fs::read_to_string(path).map_err(|source| ConfigError::Read {
  74. path: path.to_path_buf(),
  75. source,
  76. })?;
  77. let mut config: Self =
  78. serde_yaml::from_str(&contents).map_err(|source| ConfigError::Parse {
  79. path: path.to_path_buf(),
  80. source,
  81. })?;
  82. config.expand_env_refs()?;
  83. Ok(config)
  84. }
  85. pub fn provider(&self, provider: Option<&str>) -> Result<&ProviderConfig, ConfigError> {
  86. let name = match provider {
  87. Some(name) => name,
  88. None => self
  89. .default_provider
  90. .as_deref()
  91. .ok_or(ConfigError::MissingProvider)?,
  92. };
  93. self.providers
  94. .get(name)
  95. .ok_or_else(|| ConfigError::UnknownProvider {
  96. name: name.to_string(),
  97. })
  98. }
  99. fn expand_env_refs(&mut self) -> Result<(), ConfigError> {
  100. if let Some(default_provider) = &mut self.default_provider {
  101. *default_provider = expand_env_refs(default_provider)?;
  102. }
  103. for provider in self.providers.values_mut() {
  104. provider.base_url = expand_env_refs(&provider.base_url)?;
  105. provider.api_token = expand_env_refs(&provider.api_token)?;
  106. provider.default_model = expand_env_refs(&provider.default_model)?;
  107. }
  108. self.benchmarks.data_dir = expand_env_refs(&self.benchmarks.data_dir)?;
  109. if let Some(dataset) = &mut self.benchmarks.aime2026 {
  110. dataset.expand_env_refs()?;
  111. }
  112. if let Some(dataset) = &mut self.benchmarks.gpqa_diamond {
  113. dataset.expand_env_refs()?;
  114. }
  115. Ok(())
  116. }
  117. }
  118. impl DatasetConfig {
  119. fn expand_env_refs(&mut self) -> Result<(), ConfigError> {
  120. self.source = expand_env_refs(&self.source)?;
  121. self.split = expand_env_refs(&self.split)?;
  122. Ok(())
  123. }
  124. }
  125. fn default_benchmark_data_dir() -> String {
  126. "data/benchmarks".to_string()
  127. }
  128. fn expand_env_refs(value: &str) -> Result<String, ConfigError> {
  129. let regex = Regex::new(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}").expect("valid env regex");
  130. let mut expanded = String::with_capacity(value.len());
  131. let mut last_match = 0;
  132. for captures in regex.captures_iter(value) {
  133. let full_match = captures.get(0).expect("full match exists");
  134. let env_name = captures.get(1).expect("env name exists").as_str();
  135. let env_value = std::env::var(env_name).map_err(|_| ConfigError::MissingEnv {
  136. name: env_name.to_string(),
  137. })?;
  138. expanded.push_str(&value[last_match..full_match.start()]);
  139. expanded.push_str(&env_value);
  140. last_match = full_match.end();
  141. }
  142. expanded.push_str(&value[last_match..]);
  143. Ok(expanded)
  144. }
  145. #[cfg(test)]
  146. mod tests {
  147. use super::*;
  148. #[test]
  149. fn loads_provider_and_expands_env_token() {
  150. unsafe {
  151. std::env::set_var("LQ_TEST_TOKEN", "secret-token");
  152. }
  153. let temp_dir = tempfile::tempdir().expect("create temp dir");
  154. let config_path = temp_dir.path().join("config.yaml");
  155. std::fs::write(
  156. &config_path,
  157. r#"
  158. default_provider: openai
  159. providers:
  160. openai:
  161. protocol: openai
  162. base_url: https://api.openai.test/v1
  163. api_token: ${LQ_TEST_TOKEN}
  164. default_model: gpt-test
  165. benchmarks:
  166. data_dir: data/benchmarks
  167. "#,
  168. )
  169. .expect("write config");
  170. let config = AppConfig::load(&config_path).expect("load config");
  171. let provider = config.provider(None).expect("default provider");
  172. assert_eq!(provider.api_token, "secret-token");
  173. assert_eq!(provider.default_model, "gpt-test");
  174. assert_eq!(provider.protocol, ProtocolKind::Openai);
  175. }
  176. #[test]
  177. fn rejects_missing_env_token() {
  178. unsafe {
  179. std::env::remove_var("LQ_MISSING_TOKEN");
  180. }
  181. let temp_dir = tempfile::tempdir().expect("create temp dir");
  182. let config_path = temp_dir.path().join("config.yaml");
  183. std::fs::write(
  184. &config_path,
  185. r#"
  186. default_provider: openai
  187. providers:
  188. openai:
  189. protocol: openai
  190. base_url: https://api.openai.test/v1
  191. api_token: ${LQ_MISSING_TOKEN}
  192. default_model: gpt-test
  193. benchmarks:
  194. data_dir: data/benchmarks
  195. "#,
  196. )
  197. .expect("write config");
  198. let error = AppConfig::load(&config_path).expect_err("missing env should fail");
  199. assert!(error.to_string().contains("LQ_MISSING_TOKEN"));
  200. }
  201. }