Skip to main content

git_topology/embeddings/
config.rs

1use anyhow::{Context, Result};
2use serde::{Deserialize, Serialize};
3use std::process::Command;
4
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub enum EmbeddingProviderType {
7    OpenAI,
8    Gemma,
9}
10
11impl std::str::FromStr for EmbeddingProviderType {
12    type Err = anyhow::Error;
13
14    fn from_str(s: &str) -> Result<Self> {
15        match s.to_lowercase().as_str() {
16            "openai" => Ok(EmbeddingProviderType::OpenAI),
17            "gemma" => Ok(EmbeddingProviderType::Gemma),
18            _ => Err(anyhow::anyhow!(
19                "Unknown provider: {}. Valid options: openai, gemma",
20                s
21            )),
22        }
23    }
24}
25
26impl std::fmt::Display for EmbeddingProviderType {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        match self {
29            EmbeddingProviderType::OpenAI => write!(f, "openai"),
30            EmbeddingProviderType::Gemma => write!(f, "gemma"),
31        }
32    }
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct EmbeddingConfig {
37    pub provider: EmbeddingProviderType,
38    pub openai: OpenAIConfig,
39    pub gemma: GemmaConfig,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct GemmaConfig {
44    pub embedding_dim: usize,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct OpenAIConfig {
49    pub api_key: Option<String>,
50    pub model: String,
51    pub max_tokens: usize,
52}
53
54impl Default for EmbeddingConfig {
55    fn default() -> Self {
56        Self {
57            provider: EmbeddingProviderType::Gemma,
58            openai: OpenAIConfig::default(),
59            gemma: GemmaConfig::default(),
60        }
61    }
62}
63
64impl Default for GemmaConfig {
65    fn default() -> Self {
66        Self { embedding_dim: 768 }
67    }
68}
69
70impl Default for OpenAIConfig {
71    fn default() -> Self {
72        Self {
73            api_key: std::env::var("OPENAI_API_KEY").ok(),
74            model: "text-embedding-3-small".to_string(),
75            max_tokens: 8000,
76        }
77    }
78}
79
80impl EmbeddingConfig {
81    pub fn get_git_config(key: &str) -> Option<String> {
82        Command::new("git")
83            .args(["config", "--get", key])
84            .output()
85            .ok()
86            .and_then(|output| {
87                if output.status.success() {
88                    String::from_utf8(output.stdout)
89                        .ok()
90                        .map(|s| s.trim().to_string())
91                } else {
92                    None
93                }
94            })
95    }
96
97    pub fn set_git_config(key: &str, value: &str) -> Result<()> {
98        let status = Command::new("git")
99            .args(["config", key, value])
100            .status()
101            .context("Failed to execute git config")?;
102
103        if !status.success() {
104            anyhow::bail!("Failed to set git config {}", key);
105        }
106
107        Ok(())
108    }
109
110    pub fn is_provider_configured() -> bool {
111        Self::get_git_config("topology.provider").is_some()
112            || std::env::var("TOPOLOGY_PROVIDER").is_ok()
113    }
114
115    pub fn load_or_default() -> Result<Self> {
116        let provider_str = Self::get_git_config("topology.provider")
117            .or_else(|| std::env::var("TOPOLOGY_PROVIDER").ok())
118            .unwrap_or_else(|| "gemma".to_string());
119
120        let provider = provider_str.parse()?;
121
122        Ok(Self {
123            provider,
124            openai: OpenAIConfig::load(),
125            gemma: GemmaConfig::load(),
126        })
127    }
128}
129
130impl OpenAIConfig {
131    fn load() -> Self {
132        Self {
133            api_key: std::env::var("OPENAI_API_KEY").ok(),
134            model: EmbeddingConfig::get_git_config("topology.openai.model")
135                .unwrap_or_else(|| "text-embedding-3-small".to_string()),
136            max_tokens: EmbeddingConfig::get_git_config("topology.openai.maxTokens")
137                .and_then(|s| s.parse().ok())
138                .unwrap_or(8000),
139        }
140    }
141}
142
143impl GemmaConfig {
144    fn load() -> Self {
145        Self {
146            embedding_dim: EmbeddingConfig::get_git_config("topology.gemma.embeddingDim")
147                .and_then(|s| s.parse().ok())
148                .unwrap_or(768),
149        }
150    }
151}