use serde::Deserialize;
use trusty_common::memory_core::dream::DreamConfig;
use trusty_common::memory_core::semantic_consolidation::SemanticConsolidationConfig;
#[derive(Deserialize, Default, Clone)]
struct UserConfigMin {
#[serde(default)]
openrouter: OpenRouterMin,
#[serde(default)]
local_model: LocalModelMin,
}
#[derive(Deserialize, Default, Clone)]
struct OpenRouterMin {
#[serde(default)]
api_key: String,
#[serde(default)]
model: String,
}
#[derive(Deserialize, Clone)]
struct LocalModelMin {
#[serde(default = "default_local_enabled")]
enabled: bool,
#[serde(default = "default_local_base_url")]
base_url: String,
#[serde(default = "default_local_model")]
model: String,
}
fn default_local_enabled() -> bool {
true
}
fn default_local_base_url() -> String {
"http://localhost:11434".to_string()
}
fn default_local_model() -> String {
"llama3.2".to_string()
}
impl Default for LocalModelMin {
fn default() -> Self {
Self {
enabled: default_local_enabled(),
base_url: default_local_base_url(),
model: default_local_model(),
}
}
}
#[derive(Clone)]
pub struct LoadedUserConfig {
pub openrouter_api_key: String,
pub openrouter_model: String,
pub local_model: trusty_common::LocalModelConfig,
}
impl Default for LoadedUserConfig {
fn default() -> Self {
Self {
openrouter_api_key: String::new(),
openrouter_model: "anthropic/claude-3-5-sonnet".to_string(),
local_model: trusty_common::LocalModelConfig::default(),
}
}
}
pub fn load_user_config() -> Option<LoadedUserConfig> {
let home = dirs::home_dir()?;
let path = home.join(".trusty-memory").join("config.toml");
if !path.exists() {
return Some(LoadedUserConfig::default());
}
let raw = std::fs::read_to_string(&path).ok()?;
let parsed: UserConfigMin = toml::from_str(&raw).unwrap_or_default();
let model = if parsed.openrouter.model.is_empty() {
"anthropic/claude-3-5-sonnet".to_string()
} else {
parsed.openrouter.model
};
Some(LoadedUserConfig {
openrouter_api_key: parsed.openrouter.api_key,
openrouter_model: model,
local_model: trusty_common::LocalModelConfig {
enabled: parsed.local_model.enabled,
base_url: parsed.local_model.base_url,
model: parsed.local_model.model,
},
})
}
pub fn dream_config_from_user_config(cfg: &LoadedUserConfig) -> DreamConfig {
let resolved_api_key =
trusty_common::memory_core::semantic_consolidation::resolve_openrouter_api_key(
&cfg.openrouter_api_key,
);
let resolves_local = cfg.local_model.enabled && resolved_api_key.is_empty();
let model = if resolves_local {
cfg.local_model.model.clone()
} else {
cfg.openrouter_model.clone()
};
DreamConfig {
openrouter_api_key: cfg.openrouter_api_key.clone(),
local_model_enabled: cfg.local_model.enabled,
semantic: SemanticConsolidationConfig {
model,
..SemanticConsolidationConfig::default()
},
..DreamConfig::default()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[serial_test::serial]
fn dream_config_from_user_config_prefers_local_model_when_resolved() {
let _guard = EnvVarGuard::remove("OPENROUTER_API_KEY");
let cfg = LoadedUserConfig {
openrouter_api_key: String::new(),
openrouter_model: "anthropic/claude-3-5-sonnet".to_string(),
local_model: trusty_common::LocalModelConfig {
enabled: true,
base_url: "http://localhost:11434".to_string(),
model: "llama3.2".to_string(),
},
};
let dream_cfg = dream_config_from_user_config(&cfg);
assert_eq!(dream_cfg.semantic.model, "llama3.2");
assert!(dream_cfg.local_model_enabled);
assert!(dream_cfg.openrouter_api_key.is_empty());
}
#[test]
fn dream_config_from_user_config_prefers_openrouter_model_with_key() {
let cfg = LoadedUserConfig {
openrouter_api_key: "sk-test-key".to_string(),
openrouter_model: "anthropic/claude-3-5-sonnet".to_string(),
local_model: trusty_common::LocalModelConfig {
enabled: true,
base_url: "http://localhost:11434".to_string(),
model: "llama3.2".to_string(),
},
};
let dream_cfg = dream_config_from_user_config(&cfg);
assert_eq!(dream_cfg.semantic.model, "anthropic/claude-3-5-sonnet");
assert_eq!(dream_cfg.openrouter_api_key, "sk-test-key");
}
#[test]
#[serial_test::serial]
fn dream_config_from_user_config_prefers_openrouter_model_with_env_key() {
let _guard = EnvVarGuard::set("OPENROUTER_API_KEY", "sk-from-env");
let cfg = LoadedUserConfig {
openrouter_api_key: String::new(),
openrouter_model: "anthropic/claude-3-5-sonnet".to_string(),
local_model: trusty_common::LocalModelConfig {
enabled: true,
base_url: "http://localhost:11434".to_string(),
model: "llama3.2".to_string(),
},
};
let dream_cfg = dream_config_from_user_config(&cfg);
assert_eq!(
dream_cfg.semantic.model, "anthropic/claude-3-5-sonnet",
"an env-supplied OpenRouter key must resolve the OpenRouter \
model, not the local one, even though config.toml has no key"
);
}
struct EnvVarGuard {
key: &'static str,
previous: Option<String>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: &str) -> Self {
let previous = std::env::var(key).ok();
unsafe { std::env::set_var(key, value) };
Self { key, previous }
}
fn remove(key: &'static str) -> Self {
let previous = std::env::var(key).ok();
unsafe { std::env::remove_var(key) };
Self { key, previous }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match &self.previous {
Some(v) => unsafe { std::env::set_var(self.key, v) },
None => unsafe { std::env::remove_var(self.key) },
}
}
}
}