Skip to main content

kyma_embed/
config.rs

1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Serialize, Deserialize)]
4#[serde(rename_all = "lowercase", tag = "provider")]
5pub enum EmbeddingProvider {
6    Fastembed {
7        id: String,
8        model_path: Option<String>,
9    },
10    Ollama {
11        id: String,
12        base_url: String,
13    },
14    #[serde(rename = "openai-compat")]
15    OpenAICompat {
16        id: String,
17        base_url: String,
18        api_key_env: Option<String>,
19    },
20    Gemini {
21        id: String,
22        api_key_env: String,
23    },
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct EmbeddingConfig {
28    pub provider: EmbeddingProvider,
29}
30
31impl EmbeddingConfig {
32    /// Defaults to fastembed bge-small-en-v1.5 (384-dim).
33    pub fn default_fastembed() -> Self {
34        Self {
35            provider: EmbeddingProvider::Fastembed {
36                id: "bge-small-en-v1.5".into(),
37                model_path: None,
38            },
39        }
40    }
41
42    /// Load from env vars: `KYMA_EMBED_PROVIDER`, `KYMA_EMBED_MODEL_ID`,
43    /// `KYMA_EMBED_BASE_URL`, `KYMA_EMBED_MODEL_PATH`. Returns the default
44    /// fastembed config when none are set.
45    pub fn from_env() -> Self {
46        let provider = std::env::var("KYMA_EMBED_PROVIDER").ok();
47        let id = std::env::var("KYMA_EMBED_MODEL_ID").ok();
48        match provider.as_deref() {
49            Some("ollama") => Self {
50                provider: EmbeddingProvider::Ollama {
51                    id: id.unwrap_or_else(|| "nomic-embed-text".into()),
52                    base_url: std::env::var("KYMA_EMBED_BASE_URL")
53                        .unwrap_or_else(|_| "http://localhost:11434".into()),
54                },
55            },
56            Some("openai-compat") => Self {
57                provider: EmbeddingProvider::OpenAICompat {
58                    id: id.unwrap_or_else(|| "text-embedding-3-small".into()),
59                    base_url: std::env::var("KYMA_EMBED_BASE_URL")
60                        .unwrap_or_else(|_| "https://api.openai.com/v1".into()),
61                    api_key_env: Some(
62                        std::env::var("KYMA_EMBED_API_KEY_ENV")
63                            .unwrap_or_else(|_| "OPENAI_API_KEY".into()),
64                    ),
65                },
66            },
67            Some("gemini") => Self {
68                provider: EmbeddingProvider::Gemini {
69                    id: id.unwrap_or_else(|| "text-embedding-004".into()),
70                    api_key_env: "GOOGLE_API_KEY".into(),
71                },
72            },
73            Some("fastembed") => Self {
74                provider: EmbeddingProvider::Fastembed {
75                    id: id.unwrap_or_else(|| "bge-small-en-v1.5".into()),
76                    model_path: std::env::var("KYMA_EMBED_MODEL_PATH").ok(),
77                },
78            },
79            _ => Self::default_fastembed(),
80        }
81    }
82}