use serde::Deserialize;
use std::collections::HashMap;
use std::path::PathBuf;
use crate::error::LlmError;
#[derive(Debug, Deserialize)]
pub struct RobitConfig {
pub default_model: Option<String>,
pub providers: HashMap<String, ProviderConfig>,
pub app: Option<AppConfig>,
#[serde(default)]
pub channels: Option<ChannelsConfig>,
}
#[derive(Debug, Deserialize)]
pub struct ProviderConfig {
pub name: Option<String>,
pub base_url: String,
pub api_key: String,
pub models: Vec<ModelConfig>,
}
#[derive(Debug, Deserialize)]
pub struct ModelConfig {
pub id: String,
pub name: Option<String>,
pub context_window: Option<u64>,
pub max_output_tokens: Option<u64>,
pub temperature: Option<f32>,
pub max_tokens: Option<u32>,
pub supports_images: Option<bool>,
pub supports_tools: Option<bool>,
}
#[derive(Debug, Deserialize, Default)]
pub struct AppConfig {
pub log_level: Option<String>,
pub max_steps: Option<usize>,
pub enabled_tools: Option<Vec<String>>,
pub enabled_skills: Option<Vec<String>>,
pub context: Option<ContextConfig>,
pub retry: Option<RetryConfig>,
pub auto_approve: Option<bool>,
pub global_storage: Option<bool>,
pub bot: Option<BotConfig>,
}
#[derive(Debug, Deserialize)]
pub struct ContextConfig {
pub max_output_lines: Option<usize>,
pub max_output_bytes: Option<usize>,
pub reserve_ratio: Option<f32>,
pub compression_token_threshold: Option<usize>,
pub compression_enabled: Option<bool>,
}
#[derive(Debug, Deserialize)]
pub struct RetryConfig {
pub max_retries: Option<u32>,
pub initial_backoff_ms: Option<u64>,
pub max_backoff_ms: Option<u64>,
}
#[derive(Debug, Deserialize, Default)]
pub struct ChannelsConfig {
pub qq_bot: Option<QqBotConfig>,
}
#[derive(Debug, Deserialize, Clone)]
pub struct QqBotConfig {
pub app_id: String,
pub app_secret: String,
}
#[derive(Debug, Deserialize, Default)]
pub struct BotConfig {
pub confirm_timeout_secs: Option<u64>,
pub session_timeout_minutes: Option<u64>,
pub confirm_keywords: Option<ConfirmKeywordsConfig>,
}
#[derive(Debug, Deserialize, Clone, Default)]
pub struct ConfirmKeywordsConfig {
pub approve: Option<Vec<String>>,
pub reject: Option<Vec<String>>,
}
#[derive(Debug, Clone)]
pub struct ResolvedModel {
pub profile_name: String,
pub model_id: String,
pub base_url: String,
pub api_key: String,
pub max_tokens: Option<u32>,
pub temperature: Option<f32>,
pub context_window: Option<u64>,
pub supports_images: bool,
pub supports_tools: bool,
}
fn robit_home() -> Result<PathBuf, LlmError> {
let home = dirs::home_dir()
.ok_or_else(|| LlmError::ConfigError("Cannot determine home directory".to_string()))?;
Ok(home.join(".robit"))
}
fn resolve_env_var(value: &str) -> String {
if let Some(var_name) = value.strip_prefix("${").and_then(|s| s.strip_suffix('}')) {
std::env::var(var_name).unwrap_or_else(|_| value.to_string())
} else {
value.to_string()
}
}
pub fn load_config(workdir: Option<&std::path::Path>) -> Result<RobitConfig, LlmError> {
load_env();
let path = find_config_path(workdir)?;
let content = std::fs::read_to_string(&path)
.map_err(|e| LlmError::ConfigError(format!("Failed to read {}: {}", path.display(), e)))?;
let mut config: RobitConfig = toml::from_str(&content)
.map_err(|e| LlmError::ConfigError(format!("Failed to parse config.toml: {}", e)))?;
for provider in config.providers.values_mut() {
provider.api_key = resolve_env_var(&provider.api_key);
}
Ok(config)
}
pub fn load_env() {
if let Ok(robit_dir) = robit_home() {
let env_path = robit_dir.join(".env");
if env_path.exists() {
let _ = dotenvy::from_path(&env_path);
}
}
}
fn find_config_path(workdir: Option<&std::path::Path>) -> Result<PathBuf, LlmError> {
if let Some(workdir) = workdir {
let local_path = workdir.join(".robit").join("config.toml");
if local_path.exists() {
return Ok(local_path);
}
}
if let Ok(cwd) = std::env::current_dir() {
let local_path = cwd.join(".robit").join("config.toml");
if local_path.exists() {
return Ok(local_path);
}
}
let global_path = robit_home()?.join("config.toml");
if global_path.exists() {
return Ok(global_path);
}
Err(LlmError::ConfigError(format!(
"Configuration file config.toml not found.\n\
Please create one of the following:\n\
- Project-local: .robit/config.toml\n\
- Global: {}",
global_path.display()
)))
}
pub fn resolve_profile(
config: &RobitConfig,
provider_name: Option<&str>,
) -> Result<ResolvedModel, LlmError> {
let (provider_key, model_id) = if let Some(name) = provider_name {
let provider = config.providers.get(name).ok_or_else(|| {
LlmError::ConfigError(format!(
"Provider '{}' is not defined in config.toml. Available providers: {:?}",
name,
config.providers.keys().collect::<Vec<_>>()
))
})?;
let first_model = provider.models.first().ok_or_else(|| {
LlmError::ConfigError(format!("Provider '{}' has no models defined", name))
})?;
(name.to_string(), first_model.id.clone())
} else if let Some(ref default_model) = config.default_model {
parse_default_model(default_model)?
} else {
let (key, provider) = config.providers.iter().next().ok_or_else(|| {
LlmError::ConfigError("No providers defined in config.toml".to_string())
})?;
let first_model = provider.models.first().ok_or_else(|| {
LlmError::ConfigError(format!("Provider '{}' has no models defined", key))
})?;
(key.clone(), first_model.id.clone())
};
let provider = config.providers.get(&provider_key).ok_or_else(|| {
LlmError::ConfigError(format!(
"Provider '{}' is not defined in config.toml. Available providers: {:?}",
provider_key,
config.providers.keys().collect::<Vec<_>>()
))
})?;
let model = provider
.models
.iter()
.find(|m| m.id == model_id)
.ok_or_else(|| {
let available: Vec<&str> = provider.models.iter().map(|m| m.id.as_str()).collect();
LlmError::ConfigError(format!(
"Model '{}' not found in provider '{}'. Available models: {:?}",
model_id, provider_key, available
))
})?;
if provider.api_key.is_empty() || provider.api_key.starts_with("${") {
return Err(LlmError::ConfigError(format!(
"Provider '{}' API key is not configured or the environment variable is not set",
provider_key
)));
}
Ok(ResolvedModel {
profile_name: provider_key,
model_id: model.id.clone(),
base_url: provider.base_url.clone(),
api_key: provider.api_key.clone(),
max_tokens: model.max_tokens,
temperature: model.temperature,
context_window: model.context_window,
supports_images: model.supports_images.unwrap_or(false),
supports_tools: model.supports_tools.unwrap_or(false),
})
}
fn parse_default_model(default_model: &str) -> Result<(String, String), LlmError> {
let parts: Vec<&str> = default_model.splitn(2, '/').collect();
if parts.len() != 2 || parts[0].is_empty() || parts[1].is_empty() {
return Err(LlmError::ConfigError(format!(
"Invalid default_model '{}' format, expected 'provider/model' (e.g. 'deepseek/deepseek-chat')",
default_model
)));
}
Ok((parts[0].to_string(), parts[1].to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_resolve_env_var_with_env_set() {
std::env::set_var("ROBIT_TEST_KEY", "test-value-123");
assert_eq!(resolve_env_var("${ROBIT_TEST_KEY}"), "test-value-123");
std::env::remove_var("ROBIT_TEST_KEY");
}
#[test]
fn test_resolve_env_var_without_env() {
assert_eq!(
resolve_env_var("${ROBIT_NONEXISTENT_KEY}"),
"${ROBIT_NONEXISTENT_KEY}"
);
}
#[test]
fn test_resolve_env_var_plain_string() {
assert_eq!(resolve_env_var("plain-key"), "plain-key");
}
#[test]
fn test_parse_robit_config() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
name = "DeepSeek"
base_url = "https://api.deepseek.com"
api_key = "sk-test-key"
[[providers.deepseek.models]]
id = "deepseek-chat"
name = "DeepSeek Chat"
context_window = 65536
max_output_tokens = 8192
temperature = 0.0
max_tokens = 4096
[[providers.deepseek.models]]
id = "deepseek-reasoner"
name = "DeepSeek Reasoner"
context_window = 65536
temperature = 0.6
[providers.qwen]
name = "通义千问"
base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
api_key = "sk-qwen-key"
[[providers.qwen.models]]
id = "qwen-max"
name = "Qwen Max"
context_window = 32768
[app]
log_level = "DEBUG"
max_steps = 10
global_storage = true
[app.context]
max_output_lines = 500
reserve_ratio = 0.2
[app.retry]
max_retries = 3
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
assert_eq!(
config.default_model.as_deref(),
Some("deepseek/deepseek-chat")
);
assert_eq!(config.providers.len(), 2);
let ds = &config.providers["deepseek"];
assert_eq!(ds.name.as_deref(), Some("DeepSeek"));
assert_eq!(ds.base_url, "https://api.deepseek.com");
assert_eq!(ds.api_key, "sk-test-key");
assert_eq!(ds.models.len(), 2);
assert_eq!(ds.models[0].id, "deepseek-chat");
assert_eq!(ds.models[0].context_window, Some(65536));
assert_eq!(ds.models[0].temperature, Some(0.0));
assert_eq!(ds.models[0].max_tokens, Some(4096));
assert_eq!(ds.models[1].id, "deepseek-reasoner");
assert_eq!(ds.models[1].temperature, Some(0.6));
let qw = &config.providers["qwen"];
assert_eq!(qw.name.as_deref(), Some("通义千问"));
assert_eq!(qw.models.len(), 1);
assert_eq!(qw.models[0].id, "qwen-max");
let app = config.app.as_ref().unwrap();
assert_eq!(app.log_level.as_deref(), Some("DEBUG"));
assert_eq!(app.max_steps, Some(10));
assert_eq!(app.global_storage, Some(true));
assert!(app.context.is_some());
assert_eq!(app.context.as_ref().unwrap().max_output_lines, Some(500));
assert!(app.retry.is_some());
assert_eq!(app.retry.as_ref().unwrap().max_retries, Some(3));
}
#[test]
fn test_parse_config_minimal() {
let toml_str = r#"
[providers.default]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.default.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
assert!(config.default_model.is_none());
assert!(config.app.is_none());
assert_eq!(config.providers.len(), 1);
}
#[test]
fn test_resolve_profile_from_default_model() {
let config = make_test_config();
let resolved = resolve_profile(&config, None).unwrap();
assert_eq!(resolved.profile_name, "deepseek");
assert_eq!(resolved.model_id, "deepseek-chat");
assert_eq!(resolved.base_url, "https://api.deepseek.com");
assert_eq!(resolved.api_key, "sk-test");
assert_eq!(resolved.context_window, Some(65536));
assert_eq!(resolved.temperature, Some(0.0));
assert_eq!(resolved.max_tokens, Some(4096));
}
#[test]
fn test_resolve_profile_explicit_provider() {
let config = make_test_config();
let resolved = resolve_profile(&config, Some("qwen")).unwrap();
assert_eq!(resolved.profile_name, "qwen");
assert_eq!(resolved.model_id, "qwen-max");
assert_eq!(
resolved.base_url,
"https://dashscope.aliyuncs.com/compatible-mode/v1"
);
}
#[test]
fn test_resolve_profile_first_available() {
let toml_str = r#"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let resolved = resolve_profile(&config, None).unwrap();
assert_eq!(resolved.profile_name, "deepseek");
assert_eq!(resolved.model_id, "deepseek-chat");
}
#[test]
fn test_resolve_profile_not_found() {
let config = make_test_config();
let result = resolve_profile(&config, Some("nonexistent"));
assert!(result.is_err());
}
#[test]
fn test_resolve_profile_model_not_found() {
let toml_str = r#"
default_model = "deepseek/nonexistent-model"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let result = resolve_profile(&config, None);
assert!(result.is_err());
}
#[test]
fn test_resolve_profile_invalid_default_model_format() {
let toml_str = r#"
default_model = "invalid-no-slash"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let result = resolve_profile(&config, None);
assert!(result.is_err());
assert!(result
.unwrap_err()
.to_string()
.contains("Invalid default_model"));
}
#[test]
fn test_resolve_profile_empty_api_key() {
let toml_str = r#"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = ""
[[providers.deepseek.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let result = resolve_profile(&config, None);
assert!(result.is_err());
}
#[test]
fn test_parse_enabled_skills() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
[app]
enabled_skills = ["code-review", "refactor"]
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let app = config.app.as_ref().unwrap();
assert!(app.enabled_skills.is_some());
let skills = app.enabled_skills.as_ref().unwrap();
assert_eq!(skills.len(), 2);
assert_eq!(skills[0], "code-review");
assert_eq!(skills[1], "refactor");
}
#[test]
fn test_parse_enabled_tools() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
[app]
enabled_tools = ["read", "bash", "edit", "write", "grep", "find", "ls"]
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let app = config.app.as_ref().unwrap();
assert!(app.enabled_tools.is_some());
let tools = app.enabled_tools.as_ref().unwrap();
assert_eq!(tools.len(), 7);
assert_eq!(tools[0], "read");
assert_eq!(tools[1], "bash");
assert_eq!(tools[2], "edit");
assert_eq!(tools[3], "write");
assert_eq!(tools[4], "grep");
assert_eq!(tools[5], "find");
assert_eq!(tools[6], "ls");
}
#[test]
fn test_parse_auto_approve() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
[app]
auto_approve = true
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let app = config.app.as_ref().unwrap();
assert_eq!(app.auto_approve, Some(true));
}
#[test]
fn test_parse_auto_approve_default_none() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
[app]
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let app = config.app.as_ref().unwrap();
assert_eq!(app.auto_approve, None);
}
fn make_test_config() -> RobitConfig {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
context_window = 65536
temperature = 0.0
max_tokens = 4096
[providers.qwen]
base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
api_key = "sk-qwen-test"
[[providers.qwen.models]]
id = "qwen-max"
context_window = 32768
"#;
toml::from_str(toml_str).unwrap()
}
#[test]
fn test_parse_channels_and_bot_sections() {
let toml_str = r#"
default_model = "deepseek/deepseek-chat"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
[channels.qq_bot]
app_id = "123456789"
app_secret = "secret-value"
[app.bot]
confirm_timeout_secs = 60
session_timeout_minutes = 30
[app.bot.confirm_keywords]
approve = ["确认", "yes"]
reject = ["取消", "no"]
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
let qq = config
.channels
.as_ref()
.and_then(|c| c.qq_bot.as_ref())
.expect("qq_bot config missing");
assert_eq!(qq.app_id, "123456789");
assert_eq!(qq.app_secret, "secret-value");
let bot = config.app.as_ref().unwrap().bot.as_ref().unwrap();
assert_eq!(bot.confirm_timeout_secs, Some(60));
assert_eq!(bot.session_timeout_minutes, Some(30));
let kw = bot.confirm_keywords.as_ref().unwrap();
assert_eq!(kw.approve.as_ref().unwrap(), &vec!["确认".to_string(), "yes".to_string()]);
assert_eq!(kw.reject.as_ref().unwrap(), &vec!["取消".to_string(), "no".to_string()]);
}
#[test]
fn test_config_without_channels_still_parses() {
let toml_str = r#"
[providers.deepseek]
base_url = "https://api.deepseek.com"
api_key = "sk-test"
[[providers.deepseek.models]]
id = "deepseek-chat"
"#;
let config: RobitConfig = toml::from_str(toml_str).unwrap();
assert!(config.channels.is_none());
assert!(config.app.is_none() || config.app.as_ref().unwrap().bot.is_none());
}
}