use serde::Deserialize;
use serde::Serialize;
use std::collections::HashMap;
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct EmbeddingsConfig {
#[serde(default)]
pub enabled: bool,
#[serde(default = "default_provider")]
pub provider: ProviderSelection,
#[serde(skip_serializing_if = "Option::is_none")]
pub openai: Option<OpenAIConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub gemini: Option<GeminiConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub voyage: Option<VoyageConfig>,
#[serde(default)]
pub cache: CacheConfig,
}
impl Default for EmbeddingsConfig {
fn default() -> Self {
Self {
enabled: false, provider: ProviderSelection::Auto,
openai: None,
gemini: None,
voyage: None,
cache: CacheConfig::default(),
}
}
}
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum ProviderSelection {
Auto,
OpenAI,
Gemini,
Voyage,
}
const fn default_provider() -> ProviderSelection {
ProviderSelection::Auto
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OpenAIConfig {
#[serde(default = "default_openai_model")]
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub dimensions: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub api_endpoint: Option<String>,
}
fn default_openai_model() -> String {
"text-embedding-3-small".to_string()
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct GeminiConfig {
#[serde(default = "default_gemini_model")]
pub model: String,
#[serde(default = "default_task_type")]
pub task_type: String,
}
fn default_gemini_model() -> String {
"gemini-embedding-001".to_string()
}
fn default_task_type() -> String {
"retrieval_document".to_string()
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct VoyageConfig {
#[serde(default = "default_voyage_model")]
pub model: String,
#[serde(default = "default_input_type")]
pub input_type: String,
}
fn default_voyage_model() -> String {
"voyage-3.5".to_string()
}
fn default_input_type() -> String {
"document".to_string()
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct CacheConfig {
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default = "default_cache_size")]
pub max_size_mb: usize,
#[serde(default = "default_cache_ttl")]
pub ttl_seconds: u64,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
enabled: true,
max_size_mb: 500,
ttl_seconds: 3600, }
}
}
const fn default_true() -> bool {
true
}
const fn default_cache_size() -> usize {
500
}
const fn default_cache_ttl() -> u64 {
3600
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IntelligenceMode {
Light,
Medium,
Hard,
}
impl IntelligenceMode {
pub fn model_config(&self, provider: &str) -> ProviderConfig {
match (self, provider) {
(Self::Light, "openai") => ProviderConfig {
model: "text-embedding-3-small".to_string(),
dimensions: Some(256),
},
(Self::Medium, "openai") => ProviderConfig {
model: "text-embedding-3-small".to_string(),
dimensions: Some(1536),
},
(Self::Hard, "openai") => ProviderConfig {
model: "text-embedding-3-large".to_string(),
dimensions: Some(3072),
},
(Self::Light, "gemini") => ProviderConfig {
model: "gemini-embedding-001".to_string(),
dimensions: Some(256),
},
(Self::Medium, "gemini") => ProviderConfig {
model: "gemini-embedding-001".to_string(),
dimensions: Some(768),
},
(Self::Hard, "gemini") => ProviderConfig {
model: "gemini-embedding-exp-03-07".to_string(),
dimensions: Some(1536),
},
(Self::Light, "voyage") => ProviderConfig {
model: "voyage-3.5-lite".to_string(),
dimensions: None, },
(Self::Medium, "voyage") => ProviderConfig {
model: "voyage-3.5".to_string(),
dimensions: None,
},
(Self::Hard, "voyage") => ProviderConfig {
model: "voyage-3-large".to_string(),
dimensions: None,
},
_ => ProviderConfig {
model: "text-embedding-3-small".to_string(),
dimensions: Some(1536),
},
}
}
}
#[derive(Debug, Clone)]
pub struct ProviderConfig {
pub model: String,
pub dimensions: Option<usize>,
}
pub fn load_embeddings_config() -> EmbeddingsConfig {
if let Ok(config_str) = std::fs::read_to_string(
dirs::home_dir()
.unwrap_or_default()
.join(".agcodex")
.join("config.toml"),
) && let Ok(config) = toml::from_str::<HashMap<String, toml::Value>>(&config_str)
&& let Some(embeddings) = config.get("embeddings")
&& let Ok(config) = embeddings.clone().try_into::<EmbeddingsConfig>()
{
return config;
}
EmbeddingsConfig::default()
}
pub fn embeddings_available() -> bool {
let config = load_embeddings_config();
config.enabled && has_any_embedding_key()
}
pub fn has_any_embedding_key() -> bool {
std::env::var("OPENAI_EMBEDDING_KEY").is_ok()
|| std::env::var("GEMINI_API_KEY").is_ok()
|| std::env::var("VOYAGE_API_KEY").is_ok()
}
pub fn get_embedding_api_key(provider: &str) -> Option<String> {
match provider {
"openai" => std::env::var("OPENAI_EMBEDDING_KEY").ok(),
"gemini" => std::env::var("GEMINI_API_KEY").ok(),
"voyage" => std::env::var("VOYAGE_API_KEY").ok(),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config_is_disabled() {
let config = EmbeddingsConfig::default();
assert!(!config.enabled);
}
#[test]
fn test_intelligence_mode_mapping() {
let light_openai = IntelligenceMode::Light.model_config("openai");
assert_eq!(light_openai.model, "text-embedding-3-small");
assert_eq!(light_openai.dimensions, Some(256));
let hard_openai = IntelligenceMode::Hard.model_config("openai");
assert_eq!(hard_openai.model, "text-embedding-3-large");
assert_eq!(hard_openai.dimensions, Some(3072));
let medium_gemini = IntelligenceMode::Medium.model_config("gemini");
assert_eq!(medium_gemini.model, "gemini-embedding-001");
assert_eq!(medium_gemini.dimensions, Some(768));
}
}