klieo_memory_qdrant/
client.rs1use crate::error::store_err;
4use dashmap::DashMap;
5use klieo_core::error::MemoryError;
6use qdrant_client::{
7 qdrant::{CreateCollectionBuilder, Distance, VectorParamsBuilder},
8 Qdrant,
9};
10use secrecy::{ExposeSecret, SecretString};
11use std::sync::Arc;
12use tokio::sync::OnceCell;
13
14#[derive(Clone, Debug)]
21#[non_exhaustive]
22pub struct QdrantConfig {
23 url: String,
24 api_key: Option<SecretString>,
25 collection_prefix: String,
27 allow_plaintext_remote: bool,
32 embedder_id: String,
38}
39
40impl QdrantConfig {
41 pub fn new(url: impl Into<String>) -> Self {
43 Self {
44 url: url.into(),
45 api_key: None,
46 collection_prefix: "klieo_facts".to_string(),
47 allow_plaintext_remote: false,
48 embedder_id: "default".to_string(),
49 }
50 }
51
52 pub fn with_api_key(mut self, key: impl Into<SecretString>) -> Self {
54 self.api_key = Some(key.into());
55 self
56 }
57
58 pub fn with_collection_prefix(mut self, prefix: impl Into<String>) -> Self {
60 self.collection_prefix = prefix.into();
61 self
62 }
63
64 pub fn allow_plaintext_remote(mut self) -> Self {
69 self.allow_plaintext_remote = true;
70 self
71 }
72
73 pub fn with_embedder_id(mut self, id: impl Into<String>) -> Self {
79 self.embedder_id = id.into();
80 self
81 }
82
83 pub fn embedder_id(&self) -> &str {
86 &self.embedder_id
87 }
88
89 pub(crate) fn url(&self) -> &str {
90 &self.url
91 }
92
93 pub(crate) fn api_key(&self) -> Option<&str> {
94 self.api_key
95 .as_ref()
96 .map(|s: &SecretString| s.expose_secret())
97 }
98
99 pub(crate) fn collection_prefix(&self) -> &str {
100 &self.collection_prefix
101 }
102}
103
104#[derive(Clone)]
113pub(crate) struct QdrantHandle {
114 pub(crate) client: Arc<Qdrant>,
115 pub(crate) collection_prefix: String,
116 bootstrapped: Arc<DashMap<String, Arc<OnceCell<()>>>>,
117}
118
119impl QdrantHandle {
120 pub(crate) fn connect(cfg: &QdrantConfig) -> Result<Self, MemoryError> {
121 check_plaintext_scheme(cfg.url(), cfg.allow_plaintext_remote)?;
127 let mut builder = Qdrant::from_url(cfg.url());
128 if let Some(key) = cfg.api_key() {
129 builder = builder.api_key(key.to_string());
130 }
131 let client = builder.build().map_err(store_err)?;
132 Ok(Self {
133 client: Arc::new(client),
134 collection_prefix: cfg.collection_prefix().to_string(),
135 bootstrapped: Arc::new(DashMap::new()),
136 })
137 }
138
139 pub(crate) async fn ensure_collection(
147 &self,
148 name: &str,
149 vector_dim: u64,
150 ) -> Result<(), MemoryError> {
151 let key = format!("{name}/{vector_dim}");
152 let cell = self
153 .bootstrapped
154 .entry(key)
155 .or_insert_with(|| Arc::new(OnceCell::new()))
156 .clone();
157 cell.get_or_try_init(|| self.bootstrap_collection(name, vector_dim))
158 .await?;
159 Ok(())
160 }
161
162 async fn bootstrap_collection(&self, name: &str, vector_dim: u64) -> Result<(), MemoryError> {
163 let exists = self
164 .client
165 .collection_exists(name)
166 .await
167 .map_err(store_err)?;
168 if exists {
169 return self.verify_existing_dim(name, vector_dim).await;
173 }
174 let req = CreateCollectionBuilder::new(name)
175 .vectors_config(VectorParamsBuilder::new(vector_dim, Distance::Cosine));
176 self.client
177 .create_collection(req)
178 .await
179 .map_err(store_err)?;
180 Ok(())
181 }
182
183 async fn verify_existing_dim(&self, name: &str, expected_dim: u64) -> Result<(), MemoryError> {
187 let info = self.client.collection_info(name).await.map_err(store_err)?;
188 let actual = info
189 .result
190 .and_then(|i| i.config)
191 .and_then(|c| c.params)
192 .and_then(|p| p.vectors_config)
193 .and_then(|vc| vc.config)
194 .and_then(single_vector_dim);
195 match actual {
196 Some(actual) if actual != expected_dim => Err(MemoryError::Embedding(format!(
197 "qdrant collection `{name}` has dimension {actual}, but the configured \
198 embedder produces {expected_dim}-dim vectors; a model change needs a new \
199 collection or a re-index"
200 ))),
201 _ => Ok(()),
202 }
203 }
204}
205
206fn single_vector_dim(config: qdrant_client::qdrant::vectors_config::Config) -> Option<u64> {
210 use qdrant_client::qdrant::vectors_config::Config;
211 match config {
212 Config::Params(params) => Some(params.size),
213 Config::ParamsMap(_) => None,
214 }
215}
216
217fn check_plaintext_scheme(url: &str, allow_plaintext_remote: bool) -> Result<(), MemoryError> {
220 if !url.starts_with("http://") {
221 return Ok(());
222 }
223 if allow_plaintext_remote {
224 return Ok(());
225 }
226 let after_scheme = url.trim_start_matches("http://");
227 let authority = after_scheme.split('/').next().unwrap_or("");
228 let host = if let Some(end) = authority.find(']') {
231 &authority[..=end]
232 } else {
233 authority.split(':').next().unwrap_or("")
234 };
235 if matches!(host, "localhost" | "127.0.0.1" | "[::1]" | "::1") {
236 return Ok(());
237 }
238 Err(MemoryError::Store(format!(
239 "QdrantConfig URL `{url}` is plaintext http:// to a non-loopback host; \
240 the api-key is sent in cleartext. Use https://, switch to a loopback \
241 binding, or set QdrantConfig::allow_plaintext_remote() if a sealed \
242 pod network (mTLS/sidecar) handles the encryption."
243 )))
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249 use qdrant_client::qdrant::vectors_config::Config;
250 use qdrant_client::qdrant::{VectorParams, VectorParamsMap};
251 use std::sync::atomic::{AtomicUsize, Ordering};
252
253 #[test]
254 fn single_vector_dim_reads_unnamed_vector_size() {
255 let config = Config::Params(VectorParams {
256 size: 768,
257 ..Default::default()
258 });
259 assert_eq!(single_vector_dim(config), Some(768));
260 }
261
262 #[test]
263 fn single_vector_dim_skips_named_vector_collections() {
264 let config = Config::ParamsMap(VectorParamsMap {
265 map: Default::default(),
266 });
267 assert_eq!(single_vector_dim(config), None);
268 }
269
270 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
274 async fn once_cell_runs_init_exactly_once_across_concurrent_callers() {
275 let counter = Arc::new(AtomicUsize::new(0));
276 let cell: Arc<OnceCell<()>> = Arc::new(OnceCell::new());
277
278 let mut handles = Vec::new();
279 for _ in 0..16 {
280 let cell = cell.clone();
281 let counter = counter.clone();
282 handles.push(tokio::spawn(async move {
283 cell.get_or_try_init(|| async {
284 counter.fetch_add(1, Ordering::SeqCst);
285 tokio::task::yield_now().await;
286 Ok::<_, MemoryError>(())
287 })
288 .await
289 .unwrap();
290 }));
291 }
292 for handle in handles {
293 handle.await.unwrap();
294 }
295 assert_eq!(counter.load(Ordering::SeqCst), 1);
296 }
297
298 #[tokio::test]
299 async fn bootstrapped_keys_are_collection_plus_dim() {
300 let key_a = format!("{}/{}", "klieo_facts:dev", 768u64);
301 let key_b = format!("{}/{}", "klieo_facts:dev", 1024u64);
302 assert_ne!(key_a, key_b);
303 }
304
305 #[test]
308 fn check_plaintext_scheme_rejects_remote_http() {
309 assert!(check_plaintext_scheme("http://qdrant.example.com:6333", false).is_err());
310 assert!(check_plaintext_scheme("http://10.0.0.1:6333", false).is_err());
311 }
312
313 #[test]
314 fn check_plaintext_scheme_allows_loopback_http() {
315 assert!(check_plaintext_scheme("http://localhost:6333", false).is_ok());
316 assert!(check_plaintext_scheme("http://127.0.0.1:6333", false).is_ok());
317 assert!(check_plaintext_scheme("http://[::1]:6333", false).is_ok());
318 }
319
320 #[test]
321 fn check_plaintext_scheme_allows_https() {
322 assert!(check_plaintext_scheme("https://qdrant.example.com:6333", false).is_ok());
323 }
324
325 #[test]
326 fn check_plaintext_scheme_honours_opt_in() {
327 assert!(check_plaintext_scheme("http://qdrant.example.com:6333", true).is_ok());
328 }
329}