llm_embedding/
config.rs

1//! Configuration for embedding providers.
2//!
3//! Aligned with linkrs `graphdb-embedding`: the API key is optional (local
4//! servers such as Ollama need none), the expected vector dimension is part
5//! of the config so misconfigured deployments fail fast, and a preprocessor
6//! selects query/document prompt prefixes.
7
8use serde::{Deserialize, Serialize};
9
10use crate::error::EmbeddingError;
11use crate::preprocessor::PreprocessorConfig;
12
13/// Configuration for an OpenAI-compatible embedding endpoint.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct EmbeddingConfig {
16    /// Full embeddings endpoint URL (e.g. `https://api.openai.com/v1/embeddings`).
17    pub base_url: String,
18    /// Bearer API key; absent for keyless local servers.
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub api_key: Option<String>,
21    /// Embedding model name (e.g. `text-embedding-3-small`).
22    pub model: String,
23    /// Request timeout in seconds.
24    #[serde(default = "default_timeout_secs")]
25    pub timeout_secs: u64,
26    /// Expected vector dimension; required so startup validates the wiring.
27    #[serde(default, skip_serializing_if = "Option::is_none")]
28    pub dimension: Option<usize>,
29    /// Text preprocessing applied before sending the request.
30    #[serde(default)]
31    pub preprocessor: PreprocessorConfig,
32}
33
34fn default_timeout_secs() -> u64 {
35    30
36}
37
38impl EmbeddingConfig {
39    /// Creates a config for an explicit endpoint URL and model.
40    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Self {
41        Self {
42            base_url: base_url.into(),
43            api_key: None,
44            model: model.into(),
45            timeout_secs: default_timeout_secs(),
46            dimension: None,
47            preprocessor: PreprocessorConfig::default(),
48        }
49    }
50
51    /// OpenAI endpoint preset (`https://api.openai.com/v1/embeddings`).
52    pub fn openai(api_key: impl Into<String>, model: impl Into<String>) -> Self {
53        Self::new("https://api.openai.com/v1/embeddings", model).with_api_key(api_key)
54    }
55
56    /// Ollama endpoint preset (`http://localhost:11434/api/embed`).
57    pub fn ollama(model: impl Into<String>) -> Self {
58        Self::new("http://localhost:11434/api/embed", model)
59    }
60
61    /// Azure OpenAI endpoint preset; `resource_url` already contains the
62    /// deployment path and `api-version` query (e.g.
63    /// `https://{resource}.openai.azure.com/openai/deployments/{id}/embeddings?api-version=2024-02-01`).
64    pub fn azure(
65        resource_url: impl Into<String>,
66        api_key: impl Into<String>,
67        model: impl Into<String>,
68    ) -> Self {
69        Self::new(resource_url, model).with_api_key(api_key)
70    }
71
72    /// Gemini OpenAI-compatibility endpoint preset
73    /// (`https://generativelanguage.googleapis.com/v1beta/openai/embeddings`).
74    pub fn gemini(api_key: impl Into<String>, model: impl Into<String>) -> Self {
75        Self::new(
76            "https://generativelanguage.googleapis.com/v1beta/openai/embeddings",
77            model,
78        )
79        .with_api_key(api_key)
80    }
81
82    /// Sets the bearer API key.
83    pub fn with_api_key(mut self, api_key: impl Into<String>) -> Self {
84        self.api_key = Some(api_key.into());
85        self
86    }
87
88    /// Sets the request timeout in seconds.
89    pub fn with_timeout(mut self, timeout_secs: u64) -> Self {
90        self.timeout_secs = timeout_secs;
91        self
92    }
93
94    /// Sets the expected vector dimension.
95    pub fn with_dimension(mut self, dimension: usize) -> Self {
96        self.dimension = Some(dimension);
97        self
98    }
99
100    /// Sets the text preprocessor.
101    pub fn with_preprocessor(mut self, preprocessor: PreprocessorConfig) -> Self {
102        self.preprocessor = preprocessor;
103        self
104    }
105
106    /// Validates the wiring before any request is sent.
107    pub fn validate(&self) -> crate::error::Result<()> {
108        if self.base_url.is_empty() {
109            return Err(EmbeddingError::Config("base_url is required".into()));
110        }
111        if self.model.is_empty() {
112            return Err(EmbeddingError::Config("model is required".into()));
113        }
114        url::Url::parse(&self.base_url)
115            .map_err(|err| EmbeddingError::Config(format!("invalid base_url: {err}")))?;
116        match self.dimension {
117            Some(dimension) if dimension > 0 => Ok(()),
118            _ => Err(EmbeddingError::Config(
119                "dimension is required for startup validation; set it with EmbeddingConfig::with_dimension()".into(),
120            )),
121        }
122    }
123}
124
125impl Default for EmbeddingConfig {
126    fn default() -> Self {
127        Self {
128            base_url: "http://localhost:11434/api/embed".into(),
129            api_key: None,
130            model: "all-minilm".into(),
131            timeout_secs: default_timeout_secs(),
132            dimension: None,
133            preprocessor: PreprocessorConfig::default(),
134        }
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::preprocessor::PreprocessorConfig;
142
143    #[test]
144    fn presets_point_at_expected_endpoints() {
145        assert!(EmbeddingConfig::openai("sk-test", "text-embedding-3-small")
146            .base_url
147            .contains("api.openai.com"));
148        assert!(EmbeddingConfig::ollama("nomic-embed-text")
149            .base_url
150            .contains("11434"));
151        assert!(EmbeddingConfig::gemini("key", "text-embedding-004")
152            .base_url
153            .contains("generativelanguage"));
154        let azure =
155            EmbeddingConfig::azure("https://example.openai.azure.com/x", "key", "deployment");
156        assert_eq!(azure.api_key.as_deref(), Some("key"));
157    }
158
159    #[test]
160    fn validate_accepts_complete_config() {
161        let config =
162            EmbeddingConfig::new("http://localhost:11434/api/embed", "model").with_dimension(384);
163        assert!(config.validate().is_ok());
164    }
165
166    #[test]
167    fn validate_rejects_missing_fields() {
168        let missing_url = EmbeddingConfig::new("", "model").with_dimension(384);
169        assert!(missing_url.validate().is_err());
170
171        let missing_model = EmbeddingConfig::new("http://example.com", "").with_dimension(384);
172        assert!(missing_model.validate().is_err());
173
174        let bad_url = EmbeddingConfig::new("not-a-url", "model").with_dimension(384);
175        assert!(bad_url
176            .validate()
177            .unwrap_err()
178            .to_string()
179            .contains("base_url"));
180
181        let missing_dimension = EmbeddingConfig::new("http://example.com", "model");
182        assert!(missing_dimension
183            .validate()
184            .unwrap_err()
185            .to_string()
186            .contains("dimension"));
187    }
188
189    #[test]
190    fn builder_sets_preprocessor() {
191        let config = EmbeddingConfig::new("http://example.com", "model").with_preprocessor(
192            PreprocessorConfig::Prefix {
193                prefix: "search_document: ".into(),
194            },
195        );
196        assert!(matches!(
197            config.preprocessor,
198            PreprocessorConfig::Prefix { .. }
199        ));
200    }
201}
llm_embedding/config.rs

llm_embedding/
config.rs
