1use std::collections::HashMap;
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13
14use tokio::sync::RwLock;
15
16use futures::StreamExt as _;
17use qdrant_client::qdrant::{PointStruct, value::Kind};
18
19use crate::QdrantOps;
20use crate::vector_store::{VectorStore, VectorStoreError};
21
22pub type EmbedFuture = Pin<
24 Box<dyn Future<Output = Result<Vec<f32>, Box<dyn std::error::Error + Send + Sync>>> + Send>,
25>;
26
27pub trait Embeddable: Send + Sync {
33 fn key(&self) -> &str;
35
36 fn content_hash(&self) -> String;
40
41 fn embed_text(&self) -> &str;
43
44 fn to_payload(&self) -> serde_json::Value;
49}
50
51#[derive(Debug, Default, Clone)]
53pub struct SyncStats {
54 pub added: usize,
55 pub updated: usize,
56 pub removed: usize,
57 pub unchanged: usize,
58}
59
60#[derive(Debug, thiserror::Error)]
62#[non_exhaustive]
63pub enum EmbeddingRegistryError {
64 #[error("vector store error: {0}")]
65 VectorStore(#[from] VectorStoreError),
66
67 #[error("embedding error: {0}")]
68 Embedding(String),
69
70 #[error("serialization error: {0}")]
71 Serialization(String),
72
73 #[error("dimension probe failed: {0}")]
74 DimensionProbe(String),
75}
76
77impl From<Box<qdrant_client::QdrantError>> for EmbeddingRegistryError {
78 fn from(e: Box<qdrant_client::QdrantError>) -> Self {
79 Self::VectorStore(VectorStoreError::Collection(e.to_string()))
80 }
81}
82
83impl From<serde_json::Error> for EmbeddingRegistryError {
84 fn from(e: serde_json::Error) -> Self {
85 Self::Serialization(e.to_string())
86 }
87}
88
89fn normalize_model_name(name: &str) -> &str {
91 name.strip_suffix(":latest").unwrap_or(name)
92}
93
94async fn probe_dimension(
100 embed_fn: &impl Fn(&str) -> EmbedFuture,
101) -> Result<u64, EmbeddingRegistryError> {
102 crate::embed_probe::probe_vector_size(embed_fn("dimension probe"), None)
103 .await
104 .map_err(|e| EmbeddingRegistryError::DimensionProbe(e.to_string()))
105}
106
107fn model_has_changed(
113 existing: &HashMap<String, HashMap<String, String>>,
114 config_model: &str,
115) -> bool {
116 if config_model.is_empty() {
117 return false;
118 }
119 existing
120 .values()
121 .any(|stored| match stored.get("embedding_model") {
122 Some(m) => normalize_model_name(m) != normalize_model_name(config_model),
123 None => true,
125 })
126}
127
128#[derive(Clone)]
140pub struct EmbeddingRegistry {
141 ops: QdrantOps,
142 collection: String,
143 namespace: uuid::Uuid,
144 hashes: HashMap<String, String>,
145 pub concurrency: usize,
147 cached_dim: Arc<RwLock<Option<u64>>>,
150}
151
152impl std::fmt::Debug for EmbeddingRegistry {
153 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154 f.debug_struct("EmbeddingRegistry")
155 .field("collection", &self.collection)
156 .finish_non_exhaustive()
157 }
158}
159
160impl EmbeddingRegistry {
161 #[must_use]
163 pub fn new(ops: QdrantOps, collection: impl Into<String>, namespace: uuid::Uuid) -> Self {
164 Self {
165 ops,
166 collection: collection.into(),
167 namespace,
168 hashes: HashMap::new(),
169 concurrency: 4,
170 cached_dim: Arc::new(RwLock::new(None)),
171 }
172 }
173
174 #[tracing::instrument(name = "memory.embed_registry.sync", skip_all, err)]
185 pub async fn sync<T: Embeddable>(
186 &mut self,
187 items: &[T],
188 embedding_model: &str,
189 embed_fn: impl Fn(&str) -> EmbedFuture,
190 on_progress: Option<Box<dyn Fn(usize, usize) + Send>>,
191 ) -> Result<SyncStats, EmbeddingRegistryError> {
192 let mut stats = SyncStats::default();
193
194 self.ensure_collection(&embed_fn).await?;
195
196 let existing = self
197 .ops
198 .scroll_all(&self.collection, "key")
199 .await
200 .map_err(|e| {
201 EmbeddingRegistryError::VectorStore(VectorStoreError::Scroll(e.to_string()))
202 })?;
203
204 let mut current: HashMap<String, (String, &T)> = HashMap::with_capacity(items.len());
205 for item in items {
206 current.insert(item.key().to_owned(), (item.content_hash(), item));
207 }
208
209 let model_changed = model_has_changed(&existing, embedding_model);
210
211 if model_changed {
212 tracing::warn!("embedding model changed to '{embedding_model}', recreating collection");
213 self.recreate_collection(&embed_fn).await?;
214 }
215
216 let work_items = build_work_set(
217 ¤t,
218 &existing,
219 model_changed,
220 &mut stats,
221 &mut self.hashes,
222 );
223
224 let work_with_futures: Vec<(String, String, EmbedFuture, String, serde_json::Value)> =
227 work_items
228 .into_iter()
229 .map(|(key, hash, item)| {
230 let text = item.embed_text().to_owned();
231 let fut = embed_fn(&text);
232 let point_id = self.point_id(&key);
233 let payload = item.to_payload();
234 (key, hash, fut, point_id, payload)
235 })
236 .collect();
237
238 let points_to_upsert = embed_and_collect_points(
239 work_with_futures,
240 on_progress,
241 &existing,
242 embedding_model,
243 self.concurrency,
244 &mut stats,
245 &mut self.hashes,
246 )
247 .await?;
248
249 if !points_to_upsert.is_empty() {
250 self.ops
251 .upsert(&self.collection, points_to_upsert)
252 .await
253 .map_err(|e| {
254 EmbeddingRegistryError::VectorStore(VectorStoreError::Upsert(e.to_string()))
255 })?;
256 }
257
258 let orphan_ids: Vec<qdrant_client::qdrant::PointId> = existing
259 .keys()
260 .filter(|key| !current.contains_key(*key))
261 .map(|key| qdrant_client::qdrant::PointId::from(self.point_id(key).as_str()))
262 .collect();
263
264 if !orphan_ids.is_empty() {
265 stats.removed = orphan_ids.len();
266 self.ops
267 .delete_by_ids(&self.collection, orphan_ids)
268 .await
269 .map_err(|e| {
270 EmbeddingRegistryError::VectorStore(VectorStoreError::Delete(e.to_string()))
271 })?;
272 }
273
274 tracing::info!(
275 added = stats.added,
276 updated = stats.updated,
277 removed = stats.removed,
278 unchanged = stats.unchanged,
279 collection = &self.collection,
280 "embeddings synced"
281 );
282
283 Ok(stats)
284 }
285
286 #[tracing::instrument(name = "memory.embed_registry.search_raw", skip_all, err)]
304 pub async fn search_raw(
305 &self,
306 query: &str,
307 limit: usize,
308 embed_fn: impl Fn(&str) -> EmbedFuture,
309 ) -> Result<Vec<crate::ScoredVectorPoint>, EmbeddingRegistryError> {
310 let query_vec = embed_fn(query)
311 .await
312 .map_err(|e| EmbeddingRegistryError::Embedding(e.to_string()))?;
313
314 let collection_dim: Option<u64> = *self.cached_dim.read().await;
318
319 let collection_dim = if collection_dim.is_some() {
320 collection_dim
321 } else {
322 let probed = self
324 .ops
325 .get_collection_vector_size(&self.collection)
326 .await
327 .map_err(|e| {
328 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
329 })?;
330 if let Some(d) = probed {
331 self.set_cached_dim(d).await;
332 }
333 probed
334 };
335
336 if let Some(stored_dim) = collection_dim {
337 let query_dim = query_vec.len() as u64;
339 if query_dim != stored_dim {
340 return Err(EmbeddingRegistryError::DimensionProbe(format!(
341 "query vector dimension {query_dim} does not match collection '{}' \
342 dimension {stored_dim}; re-run sync to rebuild the collection",
343 self.collection
344 )));
345 }
346 }
347
348 let Ok(limit_u64) = u64::try_from(limit) else {
349 return Ok(Vec::new());
350 };
351
352 let results = self
353 .ops
354 .search(&self.collection, query_vec, limit_u64, None)
355 .await
356 .map_err(|e| {
357 EmbeddingRegistryError::VectorStore(VectorStoreError::Search(e.to_string()))
358 })?;
359
360 let scored: Vec<crate::ScoredVectorPoint> = results
361 .into_iter()
362 .map(|point| {
363 let payload: HashMap<String, serde_json::Value> = point
364 .payload
365 .into_iter()
366 .filter_map(|(k, v)| {
367 let json_val = match v.kind? {
368 Kind::StringValue(s) => serde_json::Value::String(s),
369 Kind::IntegerValue(i) => serde_json::Value::Number(i.into()),
370 Kind::BoolValue(b) => serde_json::Value::Bool(b),
371 Kind::DoubleValue(d) => {
372 serde_json::Number::from_f64(d).map(serde_json::Value::Number)?
373 }
374 _ => return None,
375 };
376 Some((k, json_val))
377 })
378 .collect();
379
380 let id = match point.id.and_then(|pid| pid.point_id_options) {
381 Some(qdrant_client::qdrant::point_id::PointIdOptions::Uuid(u)) => u,
382 Some(qdrant_client::qdrant::point_id::PointIdOptions::Num(n)) => n.to_string(),
383 None => String::new(),
384 };
385
386 crate::ScoredVectorPoint {
387 id,
388 score: point.score,
389 payload,
390 }
391 })
392 .collect();
393
394 Ok(scored)
395 }
396
397 fn point_id(&self, key: &str) -> String {
398 uuid::Uuid::new_v5(&self.namespace, key.as_bytes()).to_string()
399 }
400
401 #[tracing::instrument(name = "memory.embed_registry.get_vectors_by_keys", skip_all, err)]
412 pub async fn get_vectors_by_keys(
413 &self,
414 keys: &[String],
415 ) -> Result<HashMap<String, Vec<f32>>, EmbeddingRegistryError> {
416 if keys.is_empty() {
417 return Ok(HashMap::new());
418 }
419 let id_to_key: HashMap<String, String> =
420 keys.iter().map(|k| (self.point_id(k), k.clone())).collect();
421 let ids: Vec<String> = id_to_key.keys().cloned().collect();
422 let points = self.ops.get_points(&self.collection, ids).await?;
423 Ok(points
424 .into_iter()
425 .filter_map(|p| id_to_key.get(&p.id).map(|k| (k.clone(), p.vector)))
426 .collect())
427 }
428
429 #[tracing::instrument(
430 name = "memory.embed_registry.ensure_collection",
431 skip_all,
432 err,
433 level = "debug"
434 )]
435 async fn ensure_collection(
436 &self,
437 embed_fn: &impl Fn(&str) -> EmbedFuture,
438 ) -> Result<(), EmbeddingRegistryError> {
439 if !self
440 .ops
441 .collection_exists(&self.collection)
442 .await
443 .map_err(|e| {
444 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
445 })?
446 {
447 let vector_size = probe_dimension(embed_fn).await?;
449 self.ops
450 .ensure_collection(&self.collection, vector_size)
451 .await
452 .map_err(|e| {
453 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
454 })?;
455 tracing::info!(
456 collection = &self.collection,
457 dimensions = vector_size,
458 "created Qdrant collection"
459 );
460 self.set_cached_dim(vector_size).await;
461 return Ok(());
462 }
463
464 let existing_size = self
465 .ops
466 .get_collection_vector_size(&self.collection)
467 .await
468 .map_err(|e| {
469 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
470 })?;
471
472 let vector_size = probe_dimension(embed_fn).await?;
473
474 if existing_size == Some(vector_size) {
475 self.set_cached_dim(vector_size).await;
476 return Ok(());
477 }
478
479 tracing::warn!(
480 collection = &self.collection,
481 existing = ?existing_size,
482 required = vector_size,
483 "vector dimension mismatch, recreating collection"
484 );
485 self.ops
486 .delete_collection(&self.collection)
487 .await
488 .map_err(|e| {
489 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
490 })?;
491 self.ops
492 .ensure_collection(&self.collection, vector_size)
493 .await
494 .map_err(|e| {
495 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
496 })?;
497 tracing::info!(
498 collection = &self.collection,
499 dimensions = vector_size,
500 "created Qdrant collection"
501 );
502 self.set_cached_dim(vector_size).await;
503
504 Ok(())
505 }
506
507 #[tracing::instrument(
509 name = "memory.embed_registry.set_cached_dim",
510 skip_all,
511 level = "debug"
512 )]
513 async fn set_cached_dim(&self, dim: u64) {
514 *self.cached_dim.write().await = Some(dim);
515 }
516
517 #[tracing::instrument(
518 name = "memory.embed_registry.recreate_collection",
519 skip_all,
520 err,
521 level = "debug"
522 )]
523 async fn recreate_collection(
524 &self,
525 embed_fn: &impl Fn(&str) -> EmbedFuture,
526 ) -> Result<(), EmbeddingRegistryError> {
527 if self
528 .ops
529 .collection_exists(&self.collection)
530 .await
531 .map_err(|e| {
532 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
533 })?
534 {
535 self.ops
536 .delete_collection(&self.collection)
537 .await
538 .map_err(|e| {
539 EmbeddingRegistryError::VectorStore(VectorStoreError::Collection(e.to_string()))
540 })?;
541 tracing::info!(
542 collection = &self.collection,
543 "deleted collection for recreation"
544 );
545 }
546 self.ensure_collection(embed_fn).await
547 }
548}
549
550fn build_work_set<'a, T: Embeddable>(
556 current: &HashMap<String, (String, &'a T)>,
557 existing: &HashMap<String, HashMap<String, String>>,
558 model_changed: bool,
559 stats: &mut SyncStats,
560 hashes: &mut HashMap<String, String>,
561) -> Vec<(String, String, &'a T)> {
562 let mut work_items: Vec<(String, String, &'a T)> = Vec::new();
563 for (key, (hash, item)) in current {
564 let needs_update = if let Some(stored) = existing.get(key) {
565 model_changed || stored.get("content_hash").is_some_and(|h| h != hash)
566 } else {
567 true
568 };
569
570 if needs_update {
571 work_items.push((key.clone(), hash.clone(), *item));
572 } else {
573 stats.unchanged += 1;
574 hashes.insert(key.clone(), hash.clone());
575 }
576 }
577 work_items
578}
579
580#[tracing::instrument(
591 name = "memory.embed_registry.embed_and_collect_points",
592 skip_all,
593 err,
594 level = "debug"
595)]
596#[allow(clippy::too_many_arguments)]
597async fn embed_and_collect_points(
598 work_items: Vec<(String, String, EmbedFuture, String, serde_json::Value)>,
599 on_progress: Option<Box<dyn Fn(usize, usize) + Send>>,
600 existing: &HashMap<String, HashMap<String, String>>,
601 embedding_model: &str,
602 concurrency: usize,
603 stats: &mut SyncStats,
604 hashes: &mut HashMap<String, String>,
605) -> Result<Vec<PointStruct>, EmbeddingRegistryError> {
606 let total = work_items.len();
607 let concurrency = concurrency.max(1);
609
610 let mut stream =
612 futures::stream::iter(work_items.into_iter().map(
613 |(key, hash, fut, point_id, payload)| async move {
614 (key, hash, fut.await, point_id, payload)
615 },
616 ))
617 .buffer_unordered(concurrency);
618
619 let mut points_to_upsert = Vec::new();
620 let mut completed: usize = 0;
621 while let Some((key, hash, result, point_id, mut payload)) = stream.next().await {
622 let vector = match result {
623 Ok(v) => v,
624 Err(e) => {
625 tracing::warn!("failed to embed item '{key}': {e:#}");
626 continue;
627 }
628 };
629
630 if let Some(obj) = payload.as_object_mut() {
631 obj.insert(
632 "content_hash".into(),
633 serde_json::Value::String(hash.clone()),
634 );
635 obj.insert(
636 "embedding_model".into(),
637 serde_json::Value::String(embedding_model.to_owned()),
638 );
639 }
640 let payload_map = QdrantOps::json_to_payload(payload)?;
641
642 points_to_upsert.push(PointStruct::new(point_id, vector, payload_map));
643
644 if existing.contains_key(&key) {
645 stats.updated += 1;
646 } else {
647 stats.added += 1;
648 }
649 hashes.insert(key, hash);
650
651 completed += 1;
652 if let Some(ref cb) = on_progress {
653 cb(completed, total);
654 }
655 }
656 Ok(points_to_upsert)
657}
658
659#[cfg(test)]
660mod tests {
661 use super::*;
662
663 #[test]
664 fn normalize_no_suffix() {
665 assert_eq!(normalize_model_name("foo"), "foo");
666 }
667
668 #[test]
669 fn normalize_strips_latest() {
670 assert_eq!(normalize_model_name("foo:latest"), "foo");
671 }
672
673 #[test]
674 fn normalize_other_tag_unchanged() {
675 assert_eq!(normalize_model_name("foo:v2"), "foo:v2");
676 }
677
678 struct TestItem {
679 k: String,
680 text: String,
681 }
682
683 impl Embeddable for TestItem {
684 fn key(&self) -> &str {
685 &self.k
686 }
687
688 fn content_hash(&self) -> String {
689 let mut hasher = blake3::Hasher::new();
690 hasher.update(self.text.as_bytes());
691 hasher.finalize().to_hex().to_string()
692 }
693
694 fn embed_text(&self) -> &str {
695 &self.text
696 }
697
698 fn to_payload(&self) -> serde_json::Value {
699 serde_json::json!({"key": self.k, "text": self.text})
700 }
701 }
702
703 fn make_item(k: &str, text: &str) -> TestItem {
704 TestItem {
705 k: k.into(),
706 text: text.into(),
707 }
708 }
709
710 #[test]
711 fn registry_new_valid_url() {
712 let ops = QdrantOps::new("http://localhost:6334", None).unwrap();
713 let ns = uuid::Uuid::from_bytes([0u8; 16]);
714 let reg = EmbeddingRegistry::new(ops, "test_col", ns);
715 let dbg = format!("{reg:?}");
716 assert!(dbg.contains("EmbeddingRegistry"));
717 assert!(dbg.contains("test_col"));
718 }
719
720 #[test]
721 fn embeddable_content_hash_deterministic() {
722 let item = make_item("key", "some text");
723 assert_eq!(item.content_hash(), item.content_hash());
724 }
725
726 #[test]
727 fn embeddable_content_hash_changes() {
728 let a = make_item("key", "text a");
729 let b = make_item("key", "text b");
730 assert_ne!(a.content_hash(), b.content_hash());
731 }
732
733 #[test]
734 fn embeddable_payload_contains_key() {
735 let item = make_item("my-key", "desc");
736 let payload = item.to_payload();
737 assert_eq!(payload["key"], "my-key");
738 }
739
740 #[test]
741 fn sync_stats_default() {
742 let s = SyncStats::default();
743 assert_eq!(s.added, 0);
744 assert_eq!(s.updated, 0);
745 assert_eq!(s.removed, 0);
746 assert_eq!(s.unchanged, 0);
747 }
748
749 #[test]
750 fn sync_stats_debug() {
751 let s = SyncStats {
752 added: 1,
753 updated: 2,
754 removed: 3,
755 unchanged: 4,
756 };
757 let dbg = format!("{s:?}");
758 assert!(dbg.contains("added"));
759 }
760
761 #[tokio::test]
762 async fn search_raw_embed_fail_returns_error() {
763 let ops = QdrantOps::new("http://localhost:6334", None).unwrap();
764 let ns = uuid::Uuid::from_bytes([0u8; 16]);
765 let reg = EmbeddingRegistry::new(ops, "test", ns);
766 let embed_fn = |_: &str| -> EmbedFuture {
767 Box::pin(async {
768 Err(Box::new(std::io::Error::other("fail"))
769 as Box<dyn std::error::Error + Send + Sync>)
770 })
771 };
772 let result = reg.search_raw("query", 5, embed_fn).await;
773 assert!(result.is_err());
774 }
775
776 #[tokio::test]
782 async fn search_raw_dimension_mismatch_returns_error() {
783 let ops = QdrantOps::new("http://localhost:6334", None).unwrap();
784 let ns = uuid::Uuid::from_bytes([0u8; 16]);
785 let reg = EmbeddingRegistry::new(ops, "test_dim_guard", ns);
786
787 reg.set_cached_dim(4).await;
789
790 let embed_fn = |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0_f32, 0.0]) }) };
792 let result = reg.search_raw("query", 5, embed_fn).await;
793 assert!(
794 matches!(result, Err(EmbeddingRegistryError::DimensionProbe(_))),
795 "expected DimensionProbe error on dimension mismatch, got: {result:?}"
796 );
797 }
798
799 #[tokio::test]
804 async fn search_raw_matching_dimension_passes_guard() {
805 let ops = QdrantOps::new("http://127.0.0.1:1", None).unwrap(); let ns = uuid::Uuid::from_bytes([0u8; 16]);
807 let reg = EmbeddingRegistry::new(ops, "test_dim_pass", ns);
808
809 reg.set_cached_dim(2).await;
811
812 let embed_fn = |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0_f32, 0.0]) }) };
814 let result = reg.search_raw("query", 5, embed_fn).await;
815 assert!(
817 !matches!(result, Err(EmbeddingRegistryError::DimensionProbe(_))),
818 "guard must not fire when dimensions match"
819 );
820 }
821
822 #[tokio::test]
823 async fn get_vectors_by_keys_empty_input_short_circuits() {
824 let ops = QdrantOps::new("http://127.0.0.1:1", None).unwrap();
826 let ns = uuid::Uuid::from_bytes([0u8; 16]);
827 let reg = EmbeddingRegistry::new(ops, "test_empty_keys", ns);
828 let result = reg.get_vectors_by_keys(&[]).await.unwrap();
829 assert!(result.is_empty());
830 }
831
832 #[tokio::test]
833 async fn get_vectors_by_keys_unreachable_qdrant_errors() {
834 let ops = QdrantOps::new("http://127.0.0.1:1", None).unwrap();
835 let ns = uuid::Uuid::from_bytes([0u8; 16]);
836 let reg = EmbeddingRegistry::new(ops, "test_vec_fetch", ns);
837 let keys = vec!["skill-a".to_string(), "skill-b".to_string()];
838 let result = reg.get_vectors_by_keys(&keys).await;
839 assert!(
840 matches!(result, Err(EmbeddingRegistryError::VectorStore(_))),
841 "expected VectorStore error, got: {result:?}"
842 );
843 }
844
845 #[tokio::test]
846 async fn sync_with_unreachable_qdrant_fails() {
847 let ops = QdrantOps::new("http://127.0.0.1:1", None).unwrap();
848 let ns = uuid::Uuid::from_bytes([0u8; 16]);
849 let mut reg = EmbeddingRegistry::new(ops, "test", ns);
850 let items = vec![make_item("k", "text")];
851 let embed_fn = |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![0.1_f32, 0.2]) }) };
852 let result = reg.sync(&items, "model", embed_fn, None).await;
853 assert!(result.is_err());
854 }
855
856 fn make_existing(model: &str) -> HashMap<String, HashMap<String, String>> {
859 let mut point = HashMap::new();
860 point.insert("embedding_model".to_owned(), model.to_owned());
861 let mut map = HashMap::new();
862 map.insert("k1".to_owned(), point);
863 map
864 }
865
866 #[test]
867 fn model_has_changed_latest_vs_bare_is_false() {
868 let existing = make_existing("nomic-embed-text-v2-moe:latest");
870 assert!(!model_has_changed(&existing, "nomic-embed-text-v2-moe"));
871 }
872
873 #[test]
874 fn model_has_changed_same_model_is_false() {
875 let existing = make_existing("nomic-embed-text-v2-moe");
876 assert!(!model_has_changed(&existing, "nomic-embed-text-v2-moe"));
877 }
878
879 #[test]
880 fn model_has_changed_different_model_is_true() {
881 let existing = make_existing("all-minilm");
882 assert!(model_has_changed(&existing, "nomic-embed-text-v2-moe"));
883 }
884
885 #[test]
886 fn model_has_changed_empty_existing_is_false() {
887 assert!(!model_has_changed(&HashMap::new(), "any-model"));
888 }
889
890 #[test]
891 fn model_has_changed_absent_field_with_config_model_is_true() {
892 let mut point = HashMap::new();
894 point.insert("content_hash".to_owned(), "abc".to_owned());
895 let mut map = HashMap::new();
896 map.insert("k1".to_owned(), point);
897 assert!(model_has_changed(&map, "nomic-embed-text-v2-moe"));
898 }
899
900 #[test]
901 fn model_has_changed_absent_field_with_empty_config_model_is_false() {
902 let mut point = HashMap::new();
903 point.insert("content_hash".to_owned(), "abc".to_owned());
904 let mut map = HashMap::new();
905 map.insert("k1".to_owned(), point);
906 assert!(!model_has_changed(&map, ""));
907 }
908
909 #[test]
912 fn concurrency_zero_clamped_to_one() {
913 let ops = QdrantOps::new("http://localhost:6334", None).unwrap();
914 let ns = uuid::Uuid::from_bytes([0u8; 16]);
915 let mut reg = EmbeddingRegistry::new(ops, "test", ns);
916 reg.concurrency = 0;
917 assert_eq!(reg.concurrency.max(1), 1);
920 }
921
922 #[tokio::test]
926 #[ignore = "requires Docker for Qdrant"]
927 async fn on_progress_called_once_per_successful_embed() {
928 use std::sync::{
929 Arc,
930 atomic::{AtomicUsize, Ordering},
931 };
932 use testcontainers::GenericImage;
933 use testcontainers::core::{ContainerPort, WaitFor};
934 use testcontainers::runners::AsyncRunner;
935
936 let container = GenericImage::new("qdrant/qdrant", "v1.16.0")
937 .with_wait_for(WaitFor::message_on_stdout("gRPC listening"))
938 .with_wait_for(WaitFor::seconds(1))
939 .with_exposed_port(ContainerPort::Tcp(6334))
940 .start()
941 .await
942 .unwrap();
943 let port = container.get_host_port_ipv4(6334).await.unwrap();
944 let ops = QdrantOps::new(&format!("http://127.0.0.1:{port}"), None).unwrap();
945 let ns = uuid::Uuid::new_v4();
946 let mut reg = EmbeddingRegistry::new(ops, "test_progress", ns);
947
948 let items = [
949 make_item("a", "alpha"),
950 make_item("b", "beta"),
951 make_item("c", "gamma"),
952 ];
953 let call_count = Arc::new(AtomicUsize::new(0));
954 let last_done = Arc::new(AtomicUsize::new(0));
955 let last_total = Arc::new(AtomicUsize::new(0));
956 let cc = Arc::clone(&call_count);
957 let ld = Arc::clone(&last_done);
958 let lt = Arc::clone(&last_total);
959
960 let embed_fn =
961 |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![0.1_f32, 0.2, 0.3, 0.4]) }) };
962 let on_progress: Option<Box<dyn Fn(usize, usize) + Send>> =
963 Some(Box::new(move |completed, total| {
964 cc.fetch_add(1, Ordering::SeqCst);
965 ld.store(completed, Ordering::SeqCst);
966 lt.store(total, Ordering::SeqCst);
967 }));
968
969 let stats = reg
970 .sync(&items, "test-model", embed_fn, on_progress)
971 .await
972 .unwrap();
973 let n = stats.added + stats.updated;
974
975 assert_eq!(
976 call_count.load(Ordering::SeqCst),
977 n,
978 "on_progress call count"
979 );
980 assert_eq!(last_done.load(Ordering::SeqCst), n, "last completed");
981 assert_eq!(last_total.load(Ordering::SeqCst), n, "total");
982 }
983
984 #[tokio::test]
986 #[ignore = "requires Docker for Qdrant"]
987 async fn partial_embed_failure_skips_failed_item() {
988 use testcontainers::GenericImage;
989 use testcontainers::core::{ContainerPort, WaitFor};
990 use testcontainers::runners::AsyncRunner;
991
992 let container = GenericImage::new("qdrant/qdrant", "v1.16.0")
993 .with_wait_for(WaitFor::message_on_stdout("gRPC listening"))
994 .with_wait_for(WaitFor::seconds(1))
995 .with_exposed_port(ContainerPort::Tcp(6334))
996 .start()
997 .await
998 .unwrap();
999 let port = container.get_host_port_ipv4(6334).await.unwrap();
1000 let ops = QdrantOps::new(&format!("http://127.0.0.1:{port}"), None).unwrap();
1001 let ns = uuid::Uuid::new_v4();
1002 let mut reg = EmbeddingRegistry::new(ops, "test_partial", ns);
1003
1004 let items = [
1006 make_item("ok1", "ok text"),
1007 make_item("fail", "fail text"),
1008 make_item("ok2", "ok text 2"),
1009 ];
1010
1011 let embed_fn = |text: &str| -> EmbedFuture {
1012 if text.contains("fail") {
1013 Box::pin(async {
1014 Err(Box::new(std::io::Error::other("injected failure"))
1015 as Box<dyn std::error::Error + Send + Sync>)
1016 })
1017 } else {
1018 Box::pin(async { Ok(vec![0.1_f32, 0.2, 0.3, 0.4]) })
1019 }
1020 };
1021
1022 let stats = reg
1024 .sync(&items, "test-model", embed_fn, None)
1025 .await
1026 .unwrap();
1027 assert_eq!(
1028 stats.added, 2,
1029 "two items should be upserted, failed one skipped"
1030 );
1031 }
1032
1033 #[tokio::test]
1039 #[ignore = "requires Docker for Qdrant"]
1040 async fn search_raw_dimension_mismatch_returns_error_live() {
1041 use testcontainers::GenericImage;
1042 use testcontainers::core::{ContainerPort, WaitFor};
1043 use testcontainers::runners::AsyncRunner;
1044
1045 let container = GenericImage::new("qdrant/qdrant", "v1.16.0")
1046 .with_wait_for(WaitFor::message_on_stdout("gRPC listening"))
1047 .with_wait_for(WaitFor::seconds(1))
1048 .with_exposed_port(ContainerPort::Tcp(6334))
1049 .start()
1050 .await
1051 .unwrap();
1052 let port = container.get_host_port_ipv4(6334).await.unwrap();
1053 let ops = QdrantOps::new(&format!("http://127.0.0.1:{port}"), None).unwrap();
1054 let ns = uuid::Uuid::new_v4();
1055 let mut reg = EmbeddingRegistry::new(ops, "test_dim_guard_live", ns);
1056
1057 let items = [make_item("a", "alpha")];
1059 let embed_fn_4d =
1060 |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0_f32, 0.0, 0.0, 0.0]) }) };
1061 reg.sync(&items, "model-4d", embed_fn_4d, None)
1062 .await
1063 .unwrap();
1064
1065 let embed_fn_2d = |_: &str| -> EmbedFuture { Box::pin(async { Ok(vec![1.0_f32, 0.0]) }) };
1067 let result = reg.search_raw("query", 5, embed_fn_2d).await;
1068 assert!(
1069 matches!(result, Err(EmbeddingRegistryError::DimensionProbe(_))),
1070 "dimension mismatch must return DimensionProbe error, not silent near-zero scores; got: {result:?}"
1071 );
1072 }
1073}