git_topology/embeddings/
config.rs1use 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}