Skip to main content

zeph_memory/
embedding_registry.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Generic embedding registry backed by Qdrant.
5//!
6//! Provides deduplication through content-hash delta tracking and collection-level
7//! embedding-model change detection.
8
9use 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
22/// Boxed future returned by an embedding function.
23pub type EmbedFuture = Pin<
24    Box<dyn Future<Output = Result<Vec<f32>, Box<dyn std::error::Error + Send + Sync>>> + Send>,
25>;
26
27/// Domain type that can be stored in an [`EmbeddingRegistry`].
28///
29/// Implement this trait for any struct that should be embedded and persisted in Qdrant.
30/// The registry uses [`key`](Self::key) and [`content_hash`](Self::content_hash) to
31/// detect which items need to be re-embedded on each [`EmbeddingRegistry::sync`] call.
32pub trait Embeddable: Send + Sync {
33    /// Unique string key used for point-ID generation and delta tracking.
34    fn key(&self) -> &str;
35
36    /// BLAKE3 hex hash of all semantically relevant fields.
37    ///
38    /// When this hash changes between syncs the item's embedding is recomputed.
39    fn content_hash(&self) -> String;
40
41    /// Text that will be passed to the embedding model.
42    fn embed_text(&self) -> &str;
43
44    /// Full JSON payload to store in Qdrant alongside the vector.
45    ///
46    /// **Must** include a `"key"` field equal to [`Self::key()`] so
47    /// [`EmbeddingRegistry`] can recover items on scroll.
48    fn to_payload(&self) -> serde_json::Value;
49}
50
51/// Counters returned by [`EmbeddingRegistry::sync`].
52#[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/// Errors produced by [`EmbeddingRegistry`].
61#[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
89// Ollama appends :latest when no tag is specified; treat the two as equivalent.
90fn normalize_model_name(name: &str) -> &str {
91    name.strip_suffix(":latest").unwrap_or(name)
92}
93
94/// Probe the embedding dimension via [`crate::embed_probe::probe_vector_size`], adapting its
95/// error into [`EmbeddingRegistryError::DimensionProbe`].
96///
97/// No timeout is applied here — `EmbeddingRegistry` has never bounded this call, unlike
98/// `zeph-index`'s indexer (which has its own 15s startup timeout).
99async 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
107/// Returns `true` when any stored point uses a model name that is semantically different
108/// from `config_model` after normalizing `:latest` suffixes.
109///
110/// A missing `embedding_model` field (legacy points from pre-#3395 sessions) is treated as a
111/// mismatch: the vector was produced by an unknown model and must be regenerated.
112fn 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            // Absent field means the point was written before the model was recorded; treat as mismatch.
124            None => true,
125        })
126}
127
128/// Generic Qdrant-backed embedding registry.
129///
130/// Owns a [`QdrantOps`] instance, a collection name and a UUID namespace for
131/// deterministic point IDs (uuid v5).  The in-memory `hashes` map enables
132/// O(1) delta detection between syncs.
133///
134/// The `cached_dim` field caches the collection's vector dimension after the first successful
135/// [`sync`](Self::sync) so that [`search_raw`](Self::search_raw) can validate the query vector
136/// dimension without an extra Qdrant round-trip on every call.  When a mismatch is detected,
137/// `search_raw` returns [`EmbeddingRegistryError::DimensionProbe`] instead of silently issuing a
138/// gRPC search that would return near-zero cosine scores (Qdrant gRPC behaviour on dim mismatch).
139#[derive(Clone)]
140pub struct EmbeddingRegistry {
141    ops: QdrantOps,
142    collection: String,
143    namespace: uuid::Uuid,
144    hashes: HashMap<String, String>,
145    /// Maximum number of embedding requests dispatched concurrently during a sync.
146    pub concurrency: usize,
147    /// Vector dimension confirmed during the last successful `sync`.  Shared via `Arc` so
148    /// `Clone` works without invalidating the cached value across cloned instances.
149    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    /// Create a registry wrapping an existing [`QdrantOps`] connection.
162    #[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    /// Sync `items` into Qdrant, computing a content-hash delta to avoid
175    /// unnecessary re-embedding.  Re-creates the collection when the embedding
176    /// model changes.
177    ///
178    /// `on_progress`, when provided, is called after each successful embed+upsert with
179    /// `(completed, total)` counts so callers can display progress indicators.
180    ///
181    /// # Errors
182    ///
183    /// Returns [`EmbeddingRegistryError`] on Qdrant or embedding failures.
184    #[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            &current,
218            &existing,
219            model_changed,
220            &mut stats,
221            &mut self.hashes,
222        );
223
224        // Pre-create futures, point IDs, and payloads before taking the mutable borrow on
225        // self.hashes to avoid a double-borrow on `self`.
226        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    /// Search the collection, returning raw scored Qdrant points.
287    ///
288    /// Validates that the query vector dimension matches the collection before issuing the gRPC
289    /// call.  Qdrant gRPC silently returns near-zero cosine scores (~0.022) when dimensions
290    /// mismatch instead of returning an error — this guard prevents that silent failure.
291    ///
292    /// The dimension is checked against the cache populated by the most recent [`sync`](Self::sync)
293    /// call.  If no sync has occurred (cache is `None`) the check is skipped to avoid blocking
294    /// reads before the first sync.
295    ///
296    /// Consumers map the payloads to their domain types.
297    ///
298    /// # Errors
299    ///
300    /// Returns [`EmbeddingRegistryError::DimensionProbe`] when the query vector dimension does not
301    /// match the stored collection dimension.  Returns [`EmbeddingRegistryError::Embedding`] if the
302    /// embed function fails, or [`EmbeddingRegistryError::VectorStore`] on Qdrant search failure.
303    #[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        // Guard: Qdrant gRPC returns near-zero cosine scores when the query vector dimension
315        // does not match the stored collection dimension (issue #3418).  Check the cache first
316        // (populated by sync); fall back to a live Qdrant probe only when the cache is empty.
317        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            // Cache miss: ask Qdrant directly (first search before any sync), then populate cache.
323            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            // Safe: a Vec<f32> with 4B+ elements is impossible in practice on any 64-bit platform.
338            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    /// Retrieve stored vectors for a bounded set of keys via a single Qdrant `get_points`
402    /// round-trip (e.g. the candidate IDs returned by a prior [`Self::search_raw`] call).
403    ///
404    /// Keys with no matching point, or whose point carries no dense vector, are silently
405    /// omitted from the result — callers should treat a missing key as "vector unavailable"
406    /// rather than an error.
407    ///
408    /// # Errors
409    ///
410    /// Returns [`EmbeddingRegistryError::VectorStore`] if the underlying Qdrant call fails.
411    #[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            // Collection does not exist — probe once and create.
448            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    /// Store `dim` in the dimension cache so `search_raw` can validate without a Qdrant round-trip.
508    #[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
550/// Determine which items need embedding and update stats for unchanged ones.
551///
552/// Returns a list of `(key, hash, item)` triples that require re-embedding.  Items whose
553/// stored hash matches the current hash are counted as `unchanged` in `stats` and their
554/// hashes are pre-populated in the `hashes` map.
555fn 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/// Await each pre-created embed future and collect the resulting Qdrant points.
581///
582/// `work_items` is `(key, hash, embed_future, point_id, item_payload)` — point IDs and payloads
583/// must be pre-computed to avoid a double-borrow on the `EmbeddingRegistry` when `hashes` is
584/// mutably borrowed.
585///
586/// Processes futures with bounded concurrency (`concurrency` parameter).  Calls `on_progress`
587/// after each successful embed.  Updates `stats.added`/`stats.updated` and `hashes` in place.
588///
589/// Returns a `Vec<PointStruct>` ready for upsert, or an error if payload serialization fails.
590#[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    // Clamp concurrency to at least 1: buffer_unordered(0) silently skips all futures.
608    let concurrency = concurrency.max(1);
609
610    // Stream results as they complete so on_progress fires in real time, not after collect.
611    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    /// Validates the dimension mismatch guard in `search_raw` (issue #3418).
777    ///
778    /// When the cached collection dimension differs from the query vector dimension,
779    /// `search_raw` must return `Err(EmbeddingRegistryError::DimensionProbe)` instead of
780    /// issuing a gRPC search that would silently return near-zero cosine scores.
781    #[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        // Simulate that the collection was created with 4-dim vectors.
788        reg.set_cached_dim(4).await;
789
790        // Query with a 2-dim vector (different model / dimension).
791        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    /// Validates that `search_raw` does not reject a correctly-dimensioned query.
800    ///
801    /// When the cached dimension matches the query vector, the guard must pass and
802    /// the error (if any) comes from the Qdrant network call — not from the guard itself.
803    #[tokio::test]
804    async fn search_raw_matching_dimension_passes_guard() {
805        let ops = QdrantOps::new("http://127.0.0.1:1", None).unwrap(); // unreachable — forces network error
806        let ns = uuid::Uuid::from_bytes([0u8; 16]);
807        let reg = EmbeddingRegistry::new(ops, "test_dim_pass", ns);
808
809        // Simulate a 2-dim collection.
810        reg.set_cached_dim(2).await;
811
812        // Query with a matching 2-dim vector.
813        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        // The guard passes; the error is from the unreachable Qdrant instance.
816        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        // Unreachable Qdrant instance: if this didn't short-circuit it would error/hang.
825        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    // ── model_has_changed unit tests ──────────────────────────────────────────
857
858    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        // Root cause of #2894: stored ":latest" suffix must not trigger recreation.
869        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        // Legacy points have no embedding_model field; treat as mismatch to force recreation.
893        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    // ── concurrency guard ─────────────────────────────────────────────────────
910
911    #[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        // Clamp is applied inside sync; verify the field itself can be set to 0
918        // and the guard converts it to 1 without panicking (tested via field value).
919        assert_eq!(reg.concurrency.max(1), 1);
920    }
921
922    // ── integration tests (require live Qdrant via testcontainers) ────────────
923
924    /// Test: `on_progress` fires once per successfully embedded item with correct counts.
925    #[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    /// Test: when one embed fails, the batch continues and only successful items are upserted.
985    #[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        // Item whose embed_text contains "fail" will cause the embed_fn to return Err.
1005        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        // sync must return Ok — individual failures are warned and skipped.
1023        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    /// Validates the full dimension-mismatch guard path against a live Qdrant instance (issue #3418).
1034    ///
1035    /// Creates a collection with 4-dim vectors via `sync`, then attempts a search with a 2-dim
1036    /// query vector.  The guard in `search_raw` must return `Err(DimensionProbe)` before any
1037    /// gRPC call reaches Qdrant, preventing the silent near-zero cosine score failure.
1038    #[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        // Sync with 4-dim vectors so the collection and cache are established.
1058        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        // Search with a 2-dim query (simulates a model switch without re-sync).
1066        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}