zeph_memory/
embedding_store.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Qdrant-backed embedding store for message vector search.
5//!
6//! [`EmbeddingStore`] owns a [`VectorStore`] implementation (Qdrant in production,
7//! [`crate::db_vector_store::DbVectorStore`] in tests) and exposes typed `embed` /
8//! `search` / `delete` operations used by [`crate::semantic::SemanticMemory`].
9//!
10//! Message vectors are stored in the `zeph_conversations` Qdrant collection with a
11//! payload that includes `message_id`, `conversation_id`, `role`, and `category`.
12
13pub use qdrant_client::qdrant::Filter;
14use zeph_db::DbPool;
15#[allow(unused_imports)]
16use zeph_db::sql;
17
18use crate::db_vector_store::DbVectorStore;
19use crate::error::MemoryError;
20use crate::qdrant_ops::QdrantOps;
21use crate::types::{ConversationId, MessageId};
22use crate::vector_store::{FieldCondition, FieldValue, VectorFilter, VectorPoint, VectorStore};
23
24/// Distinguishes regular messages from summaries when storing embeddings.
25///
26/// The kind is encoded in the Qdrant payload so search filters can restrict
27/// results to one category.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29#[non_exhaustive]
30pub enum MessageKind {
31    /// A normal conversation message.
32    Regular,
33    /// A compression summary generated by the summarization subsystem.
34    Summary,
35}
36
37impl MessageKind {
38    #[must_use]
39    pub fn is_summary(self) -> bool {
40        matches!(self, Self::Summary)
41    }
42}
43
44const COLLECTION_NAME: &str = "zeph_conversations";
45
46/// Ensure a Qdrant collection exists with cosine distance vectors.
47///
48/// Idempotent: no-op if the collection already exists.
49///
50/// # Errors
51///
52/// Returns an error if Qdrant cannot be reached or collection creation fails.
53#[tracing::instrument(name = "memory.embed_store.ensure_collection", skip_all)]
54pub async fn ensure_qdrant_collection(
55    ops: &QdrantOps,
56    collection: &str,
57    vector_size: u64,
58) -> Result<(), Box<qdrant_client::QdrantError>> {
59    ops.ensure_collection(collection, vector_size).await
60}
61
62/// Typed wrapper over a [`VectorStore`] backend for conversation message embeddings.
63///
64/// Constructed via [`EmbeddingStore::new`] (Qdrant URL + optional API key) or
65/// [`EmbeddingStore::with_store`] (custom backend for testing).
66pub struct EmbeddingStore {
67    ops: Box<dyn VectorStore>,
68    collection: String,
69    pool: DbPool,
70    /// Stable per-database identity (#5742), combined with `conversation_id` in Qdrant
71    /// filters/payloads to disambiguate conversations across databases sharing one Qdrant
72    /// instance. Empty string in all constructors by default — set via
73    /// [`Self::with_db_instance_id`] at production bootstrap.
74    db_instance_id: String,
75}
76
77impl std::fmt::Debug for EmbeddingStore {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        f.debug_struct("EmbeddingStore")
80            .field("collection", &self.collection)
81            .finish_non_exhaustive()
82    }
83}
84
85/// Optional filters applied to a vector similarity search.
86#[derive(Debug)]
87pub struct SearchFilter {
88    /// Restrict results to a single conversation. `None` searches across all conversations.
89    pub conversation_id: Option<ConversationId>,
90    /// Restrict by message role (`"user"` / `"assistant"`). `None` returns all roles.
91    pub role: Option<String>,
92    /// Restrict by category payload field (category-aware memory, #2428).
93    /// When `Some`, Qdrant search is restricted to vectors with a matching `category` payload.
94    pub category: Option<String>,
95}
96
97/// A single result returned by [`EmbeddingStore::search`].
98#[derive(Debug)]
99pub struct SearchResult {
100    /// Database row ID of the matching message.
101    pub message_id: MessageId,
102    /// Conversation the message belongs to.
103    pub conversation_id: ConversationId,
104    /// Cosine similarity score in `[0, 1]`.
105    pub score: f32,
106}
107
108/// Extra Qdrant payload fields specific to one of the `store*` variants (#5486).
109///
110/// Passed to [`EmbeddingStore::store_impl`], which owns the point-id/dimensions/base-payload
111/// construction and `embeddings_metadata` upsert shared by [`EmbeddingStore::store`],
112/// [`EmbeddingStore::store_with_tool_context`], and [`EmbeddingStore::store_with_category`].
113enum StoreExtra<'a> {
114    /// No extra payload fields (plain [`EmbeddingStore::store`]).
115    None,
116    /// `category` field for category-aware memory (#2428).
117    Category(Option<&'a str>),
118    /// Tool execution metadata fields.
119    ToolContext {
120        tool_name: &'a str,
121        exit_code: Option<i32>,
122        timestamp: Option<&'a str>,
123    },
124}
125
126impl EmbeddingStore {
127    /// Create a new `EmbeddingStore` connected to the given Qdrant URL with optional API key.
128    ///
129    /// `api_key` is forwarded to [`QdrantOps::new`]. The `pool` is used for `SQLite` metadata
130    /// operations on the `embeddings_metadata` table (which must already exist via sqlx
131    /// migrations).
132    ///
133    /// # Errors
134    ///
135    /// Returns an error if the Qdrant client cannot be created.
136    pub fn new(url: &str, api_key: Option<&str>, pool: DbPool) -> Result<Self, MemoryError> {
137        let ops = QdrantOps::new(url, api_key).map_err(MemoryError::Qdrant)?;
138
139        Ok(Self {
140            ops: Box::new(ops),
141            collection: COLLECTION_NAME.into(),
142            pool,
143            db_instance_id: String::new(),
144        })
145    }
146
147    /// Create a new `EmbeddingStore` backed by `SQLite` for vector storage.
148    ///
149    /// Uses the same pool for both vector data and metadata. No external Qdrant required.
150    #[must_use]
151    pub fn new_sqlite(pool: DbPool) -> Self {
152        let ops = DbVectorStore::new(pool.clone());
153        Self {
154            ops: Box::new(ops),
155            collection: COLLECTION_NAME.into(),
156            pool,
157            db_instance_id: String::new(),
158        }
159    }
160
161    /// Create an `EmbeddingStore` backed by an arbitrary [`VectorStore`] implementation.
162    ///
163    /// Intended for testing: inject a pre-configured or mock store without requiring
164    /// an external Qdrant instance.
165    #[must_use]
166    pub fn with_store(store: Box<dyn VectorStore>, pool: DbPool) -> Self {
167        Self {
168            ops: store,
169            collection: COLLECTION_NAME.into(),
170            pool,
171            db_instance_id: String::new(),
172        }
173    }
174
175    /// Attach this database's stable [`db_instance_id`](Self::db_instance_id) (#5742).
176    ///
177    /// Production bootstrap paths chain this after construction so every Qdrant write/read
178    /// this store performs is scoped to the owning physical database. Test constructors leave
179    /// it at the default empty string, which is fine as long as a given test never mixes two
180    /// different `db_instance_id` values while asserting on isolation semantics.
181    #[must_use]
182    pub fn with_db_instance_id(mut self, id: impl Into<String>) -> Self {
183        self.db_instance_id = id.into();
184        self
185    }
186
187    /// This store's stable per-database identity (#5742).
188    #[must_use]
189    pub fn db_instance_id(&self) -> &str {
190        &self.db_instance_id
191    }
192
193    /// Return `true` if the backing store is reachable and healthy.
194    #[tracing::instrument(name = "memory.embed_store.health_check", skip_all)]
195    pub async fn health_check(&self) -> bool {
196        self.ops.health_check().await.unwrap_or(false)
197    }
198
199    /// Ensure the collection exists in Qdrant with the given vector size.
200    ///
201    /// Idempotent: no-op if the collection already exists.
202    ///
203    /// # Errors
204    ///
205    /// Returns an error if Qdrant cannot be reached or collection creation fails.
206    #[tracing::instrument(name = "memory.embed_store.ensure_collection", skip_all)]
207    pub async fn ensure_collection(&self, vector_size: u64) -> Result<(), MemoryError> {
208        self.ops
209            .ensure_collection(&self.collection, vector_size)
210            .await?;
211        // Create keyword indexes for the fields used in filtered recall so Qdrant can satisfy
212        // filter conditions in O(log n) instead of scanning all payload documents.
213        self.ops
214            .create_keyword_indexes(&self.collection, &["category", "conversation_id", "role"])
215            .await?;
216        Ok(())
217    }
218
219    /// Ensure the collection exists with a vector dimension matching an already-computed
220    /// `vector`.
221    ///
222    /// Callers that already hold an embedding (e.g. from embedding real content for search or
223    /// storage) use this instead of duplicating `vector.len() as u64` followed by a call to
224    /// [`Self::ensure_collection`].
225    ///
226    /// # Errors
227    ///
228    /// Returns an error if Qdrant cannot be reached or collection creation fails.
229    pub async fn ensure_collection_for_vector(&self, vector: &[f32]) -> Result<(), MemoryError> {
230        // Safe: a Vec<f32> with 4B+ elements is impossible in practice on any 64-bit platform.
231        self.ensure_collection(vector.len() as u64).await
232    }
233
234    /// Ensure a named collection exists with a vector dimension matching an already-computed
235    /// `vector`.
236    ///
237    /// Callers that already hold an embedding for a non-default collection (e.g. graph entity
238    /// resolution, session summaries) use this instead of duplicating `vector.len() as u64`
239    /// followed by a call to [`Self::ensure_named_collection`].
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if Qdrant cannot be reached or collection creation fails.
244    pub async fn ensure_named_collection_for_vector(
245        &self,
246        name: &str,
247        vector: &[f32],
248    ) -> Result<(), MemoryError> {
249        // Safe: a Vec<f32> with 4B+ elements is impossible in practice on any 64-bit platform.
250        self.ensure_named_collection(name, vector.len() as u64)
251            .await
252    }
253
254    /// Store a vector in Qdrant with additional tool execution metadata as payload fields.
255    ///
256    /// Metadata fields (`tool_name`, `exit_code`, `timestamp`) are stored as Qdrant payload
257    /// alongside the standard fields. This allows filtering and scoring by tool context
258    /// without corrupting the embedding vector with text prefixes.
259    ///
260    /// # Errors
261    ///
262    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
263    #[tracing::instrument(name = "memory.embed_store.store_with_tool_context", skip_all)]
264    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
265    pub async fn store_with_tool_context(
266        &self,
267        message_id: MessageId,
268        conversation_id: ConversationId,
269        role: &str,
270        vector: Vec<f32>,
271        kind: MessageKind,
272        model: &str,
273        chunk_index: u32,
274        tool_name: &str,
275        exit_code: Option<i32>,
276        timestamp: Option<&str>,
277        trust_level: Option<&str>,
278    ) -> Result<String, MemoryError> {
279        self.store_impl(
280            message_id,
281            conversation_id,
282            role,
283            vector,
284            kind,
285            model,
286            chunk_index,
287            trust_level,
288            StoreExtra::ToolContext {
289                tool_name,
290                exit_code,
291                timestamp,
292            },
293        )
294        .await
295    }
296
297    /// Store a vector in Qdrant and persist metadata to `SQLite`.
298    ///
299    /// `chunk_index` is 0 for single-vector messages and increases for each chunk
300    /// when a long message is split into multiple embeddings.
301    ///
302    /// Returns the UUID of the newly created Qdrant point.
303    ///
304    /// # Errors
305    ///
306    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
307    #[tracing::instrument(name = "memory.embed_store.store", skip_all)]
308    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
309    pub async fn store(
310        &self,
311        message_id: MessageId,
312        conversation_id: ConversationId,
313        role: &str,
314        vector: Vec<f32>,
315        kind: MessageKind,
316        model: &str,
317        chunk_index: u32,
318        trust_level: Option<&str>,
319    ) -> Result<String, MemoryError> {
320        self.store_impl(
321            message_id,
322            conversation_id,
323            role,
324            vector,
325            kind,
326            model,
327            chunk_index,
328            trust_level,
329            StoreExtra::None,
330        )
331        .await
332    }
333
334    /// Store a vector with an optional category tag in the Qdrant payload.
335    ///
336    /// Identical to [`Self::store`] but adds a `category` field to the payload when provided.
337    /// Used by category-aware memory (#2428) to enable category-filtered recall.
338    ///
339    /// Note: when `category` is `None` no `category` field is written to the Qdrant payload.
340    /// Memories stored before category-aware recall was enabled therefore won't match a
341    /// category filter — this is intentional (no silent false-positives), but a backfill
342    /// pass is needed if retrospective categorization is desired.
343    ///
344    /// # Errors
345    ///
346    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
347    #[tracing::instrument(name = "memory.embed_store.store_with_category", skip_all)]
348    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
349    pub async fn store_with_category(
350        &self,
351        message_id: MessageId,
352        conversation_id: ConversationId,
353        role: &str,
354        vector: Vec<f32>,
355        kind: MessageKind,
356        model: &str,
357        chunk_index: u32,
358        category: Option<&str>,
359        trust_level: Option<&str>,
360    ) -> Result<String, MemoryError> {
361        self.store_impl(
362            message_id,
363            conversation_id,
364            role,
365            vector,
366            kind,
367            model,
368            chunk_index,
369            trust_level,
370            StoreExtra::Category(category),
371        )
372        .await
373    }
374
375    /// Shared point-id/dimensions/base-payload construction and `embeddings_metadata` upsert
376    /// for [`Self::store`], [`Self::store_with_tool_context`], and [`Self::store_with_category`]
377    /// (#5486). `extra` supplies the payload fields that differ between the three callers.
378    /// `trust_level` is the write-time provenance tier (issue #6490); `None` omits the payload
379    /// field entirely rather than writing a placeholder, matching the `category` field's
380    /// no-silent-false-positive convention.
381    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
382    async fn store_impl(
383        &self,
384        message_id: MessageId,
385        conversation_id: ConversationId,
386        role: &str,
387        vector: Vec<f32>,
388        kind: MessageKind,
389        model: &str,
390        chunk_index: u32,
391        trust_level: Option<&str>,
392        extra: StoreExtra<'_>,
393    ) -> Result<String, MemoryError> {
394        let point_id = uuid::Uuid::new_v4().to_string();
395        let dimensions = i64::try_from(vector.len())?;
396
397        let mut payload = std::collections::HashMap::from([
398            ("message_id".to_owned(), serde_json::json!(message_id.0)),
399            (
400                "conversation_id".to_owned(),
401                serde_json::json!(conversation_id.0),
402            ),
403            (
404                "db_instance_id".to_owned(),
405                serde_json::json!(self.db_instance_id),
406            ),
407            ("role".to_owned(), serde_json::json!(role)),
408            (
409                "is_summary".to_owned(),
410                serde_json::json!(kind.is_summary()),
411            ),
412        ]);
413        if let Some(trust) = trust_level {
414            payload.insert("trust_level".to_owned(), serde_json::json!(trust));
415        }
416        match extra {
417            StoreExtra::None => {}
418            StoreExtra::Category(category) => {
419                if let Some(cat) = category {
420                    payload.insert("category".to_owned(), serde_json::json!(cat));
421                }
422            }
423            StoreExtra::ToolContext {
424                tool_name,
425                exit_code,
426                timestamp,
427            } => {
428                payload.insert("tool_name".to_owned(), serde_json::json!(tool_name));
429                if let Some(code) = exit_code {
430                    payload.insert("exit_code".to_owned(), serde_json::json!(code));
431                }
432                if let Some(ts) = timestamp {
433                    payload.insert("timestamp".to_owned(), serde_json::json!(ts));
434                }
435            }
436        }
437
438        let point = VectorPoint {
439            id: point_id.clone(),
440            vector,
441            payload,
442        };
443
444        self.ops.upsert(&self.collection, vec![point]).await?;
445
446        let chunk_index_i64 = i64::from(chunk_index);
447        zeph_db::query(sql!(
448            "INSERT INTO embeddings_metadata \
449             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
450             VALUES (?, ?, ?, ?, ?) \
451             ON CONFLICT(message_id, chunk_index, model) DO UPDATE SET \
452             qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions"
453        ))
454        .bind(message_id)
455        .bind(chunk_index_i64)
456        .bind(&point_id)
457        .bind(dimensions)
458        .bind(model)
459        .execute(&self.pool)
460        .await?;
461
462        Ok(point_id)
463    }
464
465    /// Search for similar vectors in Qdrant, returning up to `limit` results.
466    ///
467    /// `limit` is clamped to `[1, `[`MAX_SEARCH_LIMIT`](crate::MAX_SEARCH_LIMIT)`]` before
468    /// being forwarded to Qdrant (issue #6553) — the bound is enforced here rather than
469    /// relying on every caller to clamp before calling. A one-shot `tracing::warn!` fires
470    /// the first time this actually reduces the requested value.
471    ///
472    /// # Errors
473    ///
474    /// Returns an error if the Qdrant search fails.
475    #[tracing::instrument(name = "memory.embed_store.search", skip_all)]
476    pub async fn search(
477        &self,
478        query_vector: &[f32],
479        limit: usize,
480        filter: Option<SearchFilter>,
481    ) -> Result<Vec<SearchResult>, MemoryError> {
482        static CLAMP_WARNED: std::sync::atomic::AtomicBool =
483            std::sync::atomic::AtomicBool::new(false);
484        crate::warn_if_search_limit_clamped("EmbeddingStore::search", limit, &CLAMP_WARNED);
485        let limit = limit.clamp(1, crate::MAX_SEARCH_LIMIT);
486        let limit_u64 = u64::try_from(limit)?;
487
488        let vector_filter = filter.as_ref().and_then(|f| {
489            let mut must = Vec::new();
490            if let Some(cid) = f.conversation_id {
491                must.push(FieldCondition {
492                    field: "conversation_id".into(),
493                    value: FieldValue::Integer(cid.0),
494                });
495                must.push(FieldCondition {
496                    field: "db_instance_id".into(),
497                    value: FieldValue::Text(self.db_instance_id.clone()),
498                });
499            }
500            if let Some(ref role) = f.role {
501                must.push(FieldCondition {
502                    field: "role".into(),
503                    value: FieldValue::Text(role.clone()),
504                });
505            }
506            if let Some(ref category) = f.category {
507                must.push(FieldCondition {
508                    field: "category".into(),
509                    value: FieldValue::Text(category.clone()),
510                });
511            }
512            if must.is_empty() {
513                None
514            } else {
515                Some(VectorFilter {
516                    must,
517                    must_not: vec![],
518                })
519            }
520        });
521
522        let results = self
523            .ops
524            .search(
525                &self.collection,
526                query_vector.to_vec(),
527                limit_u64,
528                vector_filter,
529            )
530            .await?;
531
532        // Deduplicate by message_id, keeping the chunk with the highest score.
533        // A single message may produce multiple Qdrant points (one per chunk).
534        let mut best: std::collections::HashMap<MessageId, SearchResult> =
535            std::collections::HashMap::new();
536        for point in results {
537            let Some(message_id) = point
538                .payload
539                .get("message_id")
540                .and_then(serde_json::Value::as_i64)
541            else {
542                continue;
543            };
544            let Some(conversation_id) = point
545                .payload
546                .get("conversation_id")
547                .and_then(serde_json::Value::as_i64)
548            else {
549                continue;
550            };
551            let message_id = MessageId(message_id);
552            let entry = best.entry(message_id).or_insert(SearchResult {
553                message_id,
554                conversation_id: ConversationId(conversation_id),
555                score: f32::NEG_INFINITY,
556            });
557            if point.score > entry.score {
558                entry.score = point.score;
559            }
560        }
561
562        let mut search_results: Vec<SearchResult> = best.into_values().collect();
563        search_results.sort_by(|a, b| {
564            b.score
565                .partial_cmp(&a.score)
566                .unwrap_or(std::cmp::Ordering::Equal)
567        });
568        search_results.truncate(limit);
569
570        Ok(search_results)
571    }
572
573    /// Check whether a named collection exists in the vector store.
574    ///
575    /// # Errors
576    ///
577    /// Returns an error if the store backend cannot be reached.
578    #[tracing::instrument(name = "memory.embed_store.collection_exists", skip_all)]
579    pub async fn collection_exists(&self, name: &str) -> Result<bool, MemoryError> {
580        self.ops.collection_exists(name).await.map_err(Into::into)
581    }
582
583    /// Ensure a named collection exists in Qdrant with the given vector size.
584    ///
585    /// # Errors
586    ///
587    /// Returns an error if Qdrant cannot be reached or collection creation fails.
588    #[tracing::instrument(name = "memory.embed_store.ensure_named_collection", skip_all)]
589    pub async fn ensure_named_collection(
590        &self,
591        name: &str,
592        vector_size: u64,
593    ) -> Result<(), MemoryError> {
594        self.ops.ensure_collection(name, vector_size).await?;
595        Ok(())
596    }
597
598    /// Store a vector in a named Qdrant collection with arbitrary payload.
599    ///
600    /// Returns the UUID of the newly created point.
601    ///
602    /// # Errors
603    ///
604    /// Returns an error if the Qdrant upsert fails.
605    #[tracing::instrument(name = "memory.embed_store.store_to_collection", skip_all)]
606    pub async fn store_to_collection(
607        &self,
608        collection: &str,
609        payload: serde_json::Value,
610        vector: Vec<f32>,
611    ) -> Result<String, MemoryError> {
612        let point_id = uuid::Uuid::new_v4().to_string();
613        let payload_map: std::collections::HashMap<String, serde_json::Value> =
614            serde_json::from_value(payload)?;
615        let point = VectorPoint {
616            id: point_id.clone(),
617            vector,
618            payload: payload_map,
619        };
620        self.ops.upsert(collection, vec![point]).await?;
621        Ok(point_id)
622    }
623
624    /// Upsert a vector into a named collection, reusing an existing point ID.
625    ///
626    /// Use this when updating an existing entity to avoid orphaned Qdrant points.
627    ///
628    /// # Errors
629    ///
630    /// Returns an error if the Qdrant upsert fails.
631    #[tracing::instrument(name = "memory.embed_store.upsert_to_collection", skip_all)]
632    pub async fn upsert_to_collection(
633        &self,
634        collection: &str,
635        point_id: &str,
636        payload: serde_json::Value,
637        vector: Vec<f32>,
638    ) -> Result<(), MemoryError> {
639        let payload_map: std::collections::HashMap<String, serde_json::Value> =
640            serde_json::from_value(payload)?;
641        let point = VectorPoint {
642            id: point_id.to_owned(),
643            vector,
644            payload: payload_map,
645        };
646        self.ops.upsert(collection, vec![point]).await?;
647        Ok(())
648    }
649
650    /// Search a named Qdrant collection, returning scored points with payloads.
651    ///
652    /// `limit` is clamped to `[1, `[`MAX_SEARCH_LIMIT`](crate::MAX_SEARCH_LIMIT)`]` before
653    /// being forwarded to Qdrant (issue #6553) — the bound is enforced here rather than
654    /// relying on every caller to clamp before calling. A one-shot `tracing::warn!` fires
655    /// the first time this actually reduces the requested value.
656    ///
657    /// # Errors
658    ///
659    /// Returns an error if the Qdrant search fails.
660    #[tracing::instrument(name = "memory.embed_store.search_collection", skip_all)]
661    pub async fn search_collection(
662        &self,
663        collection: &str,
664        query_vector: &[f32],
665        limit: usize,
666        filter: Option<VectorFilter>,
667    ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
668        static CLAMP_WARNED: std::sync::atomic::AtomicBool =
669            std::sync::atomic::AtomicBool::new(false);
670        crate::warn_if_search_limit_clamped(
671            "EmbeddingStore::search_collection",
672            limit,
673            &CLAMP_WARNED,
674        );
675        let limit = limit.clamp(1, crate::MAX_SEARCH_LIMIT);
676        let limit_u64 = u64::try_from(limit)?;
677        let results = self
678            .ops
679            .search(collection, query_vector.to_vec(), limit_u64, filter)
680            .await?;
681        Ok(results)
682    }
683
684    /// Enumerate `(point_id, entity_id)` pairs for all points in `collection` that carry
685    /// an `entity_id_str` payload field.
686    ///
687    /// `entity_id_str` is a string mirror of the i64 `entity_id` written alongside the numeric
688    /// field at embedding time. The scroll API only surfaces string-typed payload values, so a
689    /// parallel string field is necessary for enumeration. Points missing `entity_id_str`
690    /// (written before this field was added) are silently skipped — they will gain the field on
691    /// the next `merge_entity` or `store_entity_embedding` call.
692    ///
693    /// # Errors
694    ///
695    /// Returns an error if the underlying scroll operation fails.
696    #[tracing::instrument(name = "memory.embed_store.scroll_all_entity_ids", skip_all)]
697    pub async fn scroll_all_entity_ids(
698        &self,
699        collection: &str,
700    ) -> Result<Vec<(String, i64)>, MemoryError> {
701        let rows = self
702            .ops
703            .scroll_all_with_point_ids(collection, "entity_id_str")
704            .await?;
705        let mut out = Vec::with_capacity(rows.len());
706        for (point_id, fields) in rows {
707            let Some(s) = fields.get("entity_id_str") else {
708                continue;
709            };
710            if let Ok(id) = s.parse::<i64>() {
711                out.push((point_id, id));
712            } else {
713                tracing::debug!(point_id, value = %s, "entity_id_str unparseable, skipping");
714            }
715        }
716        Ok(out)
717    }
718
719    /// Delete a set of points from a named collection by their Qdrant point IDs.
720    ///
721    /// This is a thin wrapper over [`VectorStore::delete_by_ids`] for use by
722    /// the stale-embedding cleanup path in `community.rs`.
723    ///
724    /// # Errors
725    ///
726    /// Returns an error if the underlying delete operation fails.
727    #[tracing::instrument(name = "memory.embed_store.delete_from_collection", skip_all)]
728    pub async fn delete_from_collection(
729        &self,
730        collection: &str,
731        ids: Vec<String>,
732    ) -> Result<(), MemoryError> {
733        if ids.is_empty() {
734            return Ok(());
735        }
736        self.ops.delete_by_ids(collection, ids).await?;
737        Ok(())
738    }
739
740    /// Retrieve raw vectors for the given Qdrant point IDs from `collection`.
741    ///
742    /// Returns a map of `point_id → embedding`. Missing ids are silently dropped.
743    /// Returns an empty map when the backend does not support vector retrieval
744    /// (e.g. `DbVectorStore` / `InMemoryVectorStore` without an override).
745    ///
746    /// # Errors
747    ///
748    /// Returns an error if the underlying store returns a non-`Unsupported` error.
749    #[tracing::instrument(name = "memory.embed_store.get_vectors_from_collection", skip_all)]
750    pub async fn get_vectors_from_collection(
751        &self,
752        collection: &str,
753        point_ids: &[String],
754    ) -> Result<std::collections::HashMap<String, Vec<f32>>, MemoryError> {
755        if point_ids.is_empty() {
756            return Ok(std::collections::HashMap::new());
757        }
758        match self.ops.get_points(collection, point_ids.to_vec()).await {
759            Ok(points) => Ok(points.into_iter().map(|p| (p.id, p.vector)).collect()),
760            Err(crate::VectorStoreError::Unsupported(_)) => Ok(std::collections::HashMap::new()),
761            Err(e) => Err(MemoryError::VectorStore(e)),
762        }
763    }
764
765    /// Fetch raw vectors for the given message IDs from the `SQLite` vector store.
766    ///
767    /// Returns an empty map when using Qdrant backend (vectors not locally stored).
768    ///
769    /// # Errors
770    ///
771    /// Returns an error if the `SQLite` query fails.
772    #[tracing::instrument(name = "memory.embed_store.get_vectors", skip_all)]
773    pub async fn get_vectors(
774        &self,
775        ids: &[MessageId],
776    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
777        if ids.is_empty() {
778            return Ok(std::collections::HashMap::new());
779        }
780
781        let placeholders = zeph_db::placeholder_list(1, ids.len());
782        let query = format!(
783            "SELECT em.message_id, vp.vector \
784             FROM embeddings_metadata em \
785             JOIN vector_points vp ON vp.id = em.qdrant_point_id \
786             WHERE em.message_id IN ({placeholders}) AND em.chunk_index = 0"
787        );
788        let mut q = zeph_db::query_as::<_, (MessageId, Vec<u8>)>(sqlx::AssertSqlSafe(query));
789        for &id in ids {
790            q = q.bind(id);
791        }
792
793        let rows = q.fetch_all(&self.pool).await?;
794
795        let map = rows
796            .into_iter()
797            .filter_map(|(msg_id, blob)| {
798                if blob.len() % 4 != 0 {
799                    return None;
800                }
801                let vec: Vec<f32> = blob
802                    .chunks_exact(4)
803                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
804                    .collect();
805                Some((msg_id, vec))
806            })
807            .collect();
808
809        Ok(map)
810    }
811
812    /// Fetch embeddings for the given message IDs from the configured vector store.
813    ///
814    /// Resolves `message_id → qdrant_point_id` via `embeddings_metadata` (filtering to
815    /// `chunk_index = 0` so each message yields at most one vector), then retrieves the
816    /// vectors from the underlying [`VectorStore`].
817    ///
818    /// Returns a map from [`MessageId`] to embedding vector. Messages without an
819    /// `embeddings_metadata` row, or whose vector cannot be retrieved, are silently dropped.
820    /// When the backend returns [`crate::VectorStoreError::Unsupported`], an empty map is
821    /// returned without error (matches [`Self::get_vectors_from_collection`] semantics).
822    ///
823    /// # Errors
824    ///
825    /// Returns an error if the `SQLite` metadata query or vector store retrieval fails.
826    #[tracing::instrument(name = "memory.embed_store.get_vectors_for_messages", skip_all)]
827    pub async fn get_vectors_for_messages(
828        &self,
829        ids: &[MessageId],
830    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
831        if ids.is_empty() {
832            return Ok(std::collections::HashMap::new());
833        }
834
835        let placeholders = zeph_db::placeholder_list(1, ids.len());
836        let query = format!(
837            "SELECT message_id, qdrant_point_id \
838             FROM embeddings_metadata \
839             WHERE message_id IN ({placeholders}) AND chunk_index = 0"
840        );
841        let mut q = zeph_db::query_as::<_, (MessageId, String)>(sqlx::AssertSqlSafe(query));
842        for &id in ids {
843            q = q.bind(id);
844        }
845        let rows: Vec<(MessageId, String)> = q.fetch_all(&self.pool).await?;
846
847        if rows.is_empty() {
848            return Ok(std::collections::HashMap::new());
849        }
850
851        // Build reverse map: point_id → message_id for result translation.
852        let mut point_to_msg: std::collections::HashMap<String, MessageId> =
853            std::collections::HashMap::with_capacity(rows.len());
854        let point_ids: Vec<String> = rows
855            .into_iter()
856            .map(|(msg_id, point_id)| {
857                point_to_msg.insert(point_id.clone(), msg_id);
858                point_id
859            })
860            .collect();
861
862        let points = match self.ops.get_points(&self.collection, point_ids).await {
863            Ok(pts) => pts,
864            Err(crate::VectorStoreError::Unsupported(_)) => {
865                return Ok(std::collections::HashMap::new());
866            }
867            Err(e) => return Err(MemoryError::VectorStore(e)),
868        };
869
870        let result = points
871            .into_iter()
872            .filter_map(|p| {
873                let msg_id = point_to_msg.get(&p.id).copied()?;
874                Some((msg_id, p.vector))
875            })
876            .collect();
877
878        Ok(result)
879    }
880
881    /// Delete all Qdrant vectors associated with the given message IDs.
882    ///
883    /// Resolves `message_id → qdrant_point_id` via the `embeddings_metadata` table,
884    /// then calls the underlying vector store's `delete_by_ids`. The
885    /// `embeddings_metadata` rows are **not** removed here — the `SQLite` CASCADE on
886    /// `messages` handles that when the rows are hard-deleted later.
887    ///
888    /// Returns the number of Qdrant point IDs targeted for deletion (may be less than
889    /// `ids.len()` when some messages have no embeddings).
890    ///
891    /// # Errors
892    ///
893    /// Returns [`MemoryError`] if the `SQLite` query or the vector store delete fails.
894    #[tracing::instrument(name = "memory.embed_store.delete_by_message_ids", skip_all)]
895    pub async fn delete_by_message_ids(&self, ids: &[MessageId]) -> Result<usize, MemoryError> {
896        if ids.is_empty() {
897            return Ok(0);
898        }
899
900        let placeholders = zeph_db::placeholder_list(1, ids.len());
901        let query = format!(
902            "SELECT qdrant_point_id FROM embeddings_metadata WHERE message_id IN ({placeholders})"
903        );
904        let mut q = zeph_db::query_as::<_, (String,)>(sqlx::AssertSqlSafe(query));
905        for &id in ids {
906            q = q.bind(id);
907        }
908        let rows: Vec<(String,)> = q.fetch_all(&self.pool).await?;
909
910        let point_ids: Vec<String> = rows.into_iter().map(|(id,)| id).collect();
911        let count = point_ids.len();
912
913        if !point_ids.is_empty() {
914            self.ops.delete_by_ids(&self.collection, point_ids).await?;
915        }
916
917        Ok(count)
918    }
919
920    /// Check whether an embedding already exists for the given message ID.
921    ///
922    /// # Errors
923    ///
924    /// Returns an error if the `SQLite` query fails.
925    #[tracing::instrument(name = "memory.embed_store.has_embedding", skip_all)]
926    pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
927        let row: (i64,) = zeph_db::query_as(sql!(
928            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
929        ))
930        .bind(message_id)
931        .fetch_one(&self.pool)
932        .await?;
933
934        Ok(row.0 > 0)
935    }
936
937    /// Check whether a Qdrant embedding for `entity_name` is current by comparing the
938    /// Qdrant-side epoch against the epoch stored in `graph_entities`.
939    ///
940    /// Returns `true` if the Qdrant embedding is up-to-date or if the entity no longer
941    /// exists in `SQLite` (embedding should be cleaned up separately).
942    ///
943    /// # Errors
944    ///
945    /// Returns an error if the `SQLite` query fails.
946    #[tracing::instrument(name = "memory.embed_store.is_epoch_current", skip_all)]
947    pub async fn is_epoch_current(
948        &self,
949        entity_name: &str,
950        qdrant_epoch: u64,
951    ) -> Result<bool, MemoryError> {
952        let row: Option<(i64,)> = zeph_db::query_as(sql!(
953            "SELECT embedding_epoch FROM graph_entities WHERE name = ? LIMIT 1"
954        ))
955        .bind(entity_name)
956        .fetch_optional(&self.pool)
957        .await?;
958
959        match row {
960            None => Ok(true), // entity deleted; Qdrant point is orphaned, not stale per epoch
961            Some((db_epoch,)) => Ok(qdrant_epoch >= db_epoch.cast_unsigned()),
962        }
963    }
964}
965
966#[cfg(test)]
967mod tests {
968    use super::*;
969    use crate::db_vector_store::DbVectorStore;
970    use crate::in_memory_store::InMemoryVectorStore;
971    use crate::store::SqliteStore;
972
973    async fn setup() -> (SqliteStore, DbPool) {
974        let store = SqliteStore::new(":memory:").await.unwrap();
975        let pool = store.pool().clone();
976        (store, pool)
977    }
978
979    async fn setup_with_store() -> (EmbeddingStore, SqliteStore) {
980        let sqlite = SqliteStore::new(":memory:").await.unwrap();
981        let pool = sqlite.pool().clone();
982        let mem_store = Box::new(InMemoryVectorStore::new());
983        let embedding_store = EmbeddingStore::with_store(mem_store, pool);
984        // Create collection first
985        embedding_store.ensure_collection(4).await.unwrap();
986        (embedding_store, sqlite)
987    }
988
989    #[tokio::test]
990    async fn has_embedding_returns_false_when_none() {
991        let (_store, pool) = setup().await;
992
993        let row: (i64,) = zeph_db::query_as(sql!(
994            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
995        ))
996        .bind(999_i64)
997        .fetch_one(&pool)
998        .await
999        .unwrap();
1000
1001        assert_eq!(row.0, 0);
1002    }
1003
1004    #[tokio::test]
1005    async fn insert_and_query_embeddings_metadata() {
1006        let (sqlite, pool) = setup().await;
1007        let cid = sqlite.create_conversation().await.unwrap();
1008        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1009
1010        let point_id = uuid::Uuid::new_v4().to_string();
1011        zeph_db::query(sql!(
1012            "INSERT INTO embeddings_metadata \
1013             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1014             VALUES (?, ?, ?, ?, ?)"
1015        ))
1016        .bind(msg_id)
1017        .bind(0_i64)
1018        .bind(&point_id)
1019        .bind(768_i64)
1020        .bind("qwen3-embedding")
1021        .execute(&pool)
1022        .await
1023        .unwrap();
1024
1025        let row: (i64,) = zeph_db::query_as(sql!(
1026            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
1027        ))
1028        .bind(msg_id)
1029        .fetch_one(&pool)
1030        .await
1031        .unwrap();
1032        assert_eq!(row.0, 1);
1033    }
1034
1035    #[tokio::test]
1036    async fn embedding_store_search_empty_returns_empty() {
1037        let (store, _sqlite) = setup_with_store().await;
1038        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 10, None).await.unwrap();
1039        assert!(results.is_empty());
1040    }
1041
1042    #[tokio::test]
1043    async fn embedding_store_store_and_search() {
1044        let (store, sqlite) = setup_with_store().await;
1045        let cid = sqlite.create_conversation().await.unwrap();
1046        let msg_id = sqlite
1047            .save_message(cid, "user", "test message")
1048            .await
1049            .unwrap();
1050
1051        store
1052            .store(
1053                msg_id,
1054                cid,
1055                "user",
1056                vec![1.0, 0.0, 0.0, 0.0],
1057                MessageKind::Regular,
1058                "test-model",
1059                0,
1060                None,
1061            )
1062            .await
1063            .unwrap();
1064
1065        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5, None).await.unwrap();
1066        assert_eq!(results.len(), 1);
1067        assert_eq!(results[0].message_id, msg_id);
1068        assert_eq!(results[0].conversation_id, cid);
1069        assert!((results[0].score - 1.0).abs() < 0.001);
1070    }
1071
1072    /// Issue #6553: `search` must clamp an oversized `limit` to `MAX_SEARCH_LIMIT` internally,
1073    /// rather than relying on every caller to clamp before calling, and must log a one-shot
1074    /// warning so a config-driven candidate pool silently shrunk by the clamp is observable
1075    /// (critic finding S1).
1076    #[tokio::test]
1077    #[tracing_test::traced_test]
1078    async fn embedding_store_search_clamps_oversized_limit() {
1079        let (store, sqlite) = setup_with_store().await;
1080        let cid = sqlite.create_conversation().await.unwrap();
1081
1082        for i in 0..(crate::MAX_SEARCH_LIMIT + 10) {
1083            let msg_id = sqlite
1084                .save_message(cid, "user", &format!("message {i}"))
1085                .await
1086                .unwrap();
1087            store
1088                .store(
1089                    msg_id,
1090                    cid,
1091                    "user",
1092                    vec![1.0, 0.0, 0.0, 0.0],
1093                    MessageKind::Regular,
1094                    "test-model",
1095                    0,
1096                    None,
1097                )
1098                .await
1099                .unwrap();
1100        }
1101
1102        let results = store
1103            .search(&[1.0, 0.0, 0.0, 0.0], usize::MAX, None)
1104            .await
1105            .unwrap();
1106        assert_eq!(results.len(), crate::MAX_SEARCH_LIMIT);
1107        assert!(
1108            logs_contain("requested search limit exceeds MAX_SEARCH_LIMIT"),
1109            "expected a one-shot warn when the clamp actually reduces the requested limit"
1110        );
1111    }
1112
1113    /// #5742 regression: two independent databases (own `conversation_id` counters, each starting
1114    /// at 1) share one backing vector store. Before the fix, `search()`'s `conversation_id`-only
1115    /// filter would let a message stored under db-b leak into db-a's scoped search and vice versa.
1116    #[tokio::test]
1117    async fn embedding_store_search_scoped_excludes_other_db_instance_same_conversation_id() {
1118        let shared = SqliteStore::new(":memory:").await.unwrap();
1119        let shared_pool = shared.pool().clone();
1120
1121        let sqlite_a = SqliteStore::new(":memory:").await.unwrap();
1122        let store_a = EmbeddingStore::with_store(
1123            Box::new(DbVectorStore::new(shared_pool.clone())),
1124            sqlite_a.pool().clone(),
1125        )
1126        .with_db_instance_id("db-a");
1127        store_a.ensure_collection(4).await.unwrap();
1128
1129        let sqlite_b = SqliteStore::new(":memory:").await.unwrap();
1130        let store_b = EmbeddingStore::with_store(
1131            Box::new(DbVectorStore::new(shared_pool.clone())),
1132            sqlite_b.pool().clone(),
1133        )
1134        .with_db_instance_id("db-b");
1135
1136        let cid_a = sqlite_a.create_conversation().await.unwrap();
1137        let cid_b = sqlite_b.create_conversation().await.unwrap();
1138        assert_eq!(
1139            cid_a, cid_b,
1140            "both databases must independently start at conversation_id=1"
1141        );
1142
1143        let msg_a = sqlite_a
1144            .save_message(cid_a, "user", "message in db a")
1145            .await
1146            .unwrap();
1147        let msg_b = sqlite_b
1148            .save_message(cid_b, "user", "message in db b")
1149            .await
1150            .unwrap();
1151
1152        store_a
1153            .store(
1154                msg_a,
1155                cid_a,
1156                "user",
1157                vec![1.0, 0.0, 0.0, 0.0],
1158                MessageKind::Regular,
1159                "test-model",
1160                0,
1161                None,
1162            )
1163            .await
1164            .unwrap();
1165        store_b
1166            .store(
1167                msg_b,
1168                cid_b,
1169                "user",
1170                vec![1.0, 0.0, 0.0, 0.0],
1171                MessageKind::Regular,
1172                "test-model",
1173                0,
1174                None,
1175            )
1176            .await
1177            .unwrap();
1178
1179        let filter_a = Some(SearchFilter {
1180            conversation_id: Some(cid_a),
1181            role: None,
1182            category: None,
1183        });
1184        let results_a = store_a
1185            .search(&[1.0, 0.0, 0.0, 0.0], 5, filter_a)
1186            .await
1187            .unwrap();
1188        assert_eq!(
1189            results_a.len(),
1190            1,
1191            "db-a's scoped search must not see db-b's message even though both use conversation_id=1"
1192        );
1193
1194        let filter_b = Some(SearchFilter {
1195            conversation_id: Some(cid_b),
1196            role: None,
1197            category: None,
1198        });
1199        let results_b = store_b
1200            .search(&[1.0, 0.0, 0.0, 0.0], 5, filter_b)
1201            .await
1202            .unwrap();
1203        assert_eq!(
1204            results_b.len(),
1205            1,
1206            "db-b's scoped search must not see db-a's message even though both use conversation_id=1"
1207        );
1208    }
1209
1210    /// `store_with_category` must write a `category` payload field when given `Some` (#5486
1211    /// shared `store_impl` helper must preserve this variant's distinguishing behavior).
1212    #[tokio::test]
1213    async fn embedding_store_store_with_category_sets_payload_field() {
1214        let (store, sqlite) = setup_with_store().await;
1215        let cid = sqlite.create_conversation().await.unwrap();
1216        let msg_id = sqlite.save_message(cid, "user", "cat test").await.unwrap();
1217
1218        store
1219            .store_with_category(
1220                msg_id,
1221                cid,
1222                "user",
1223                vec![1.0, 0.0, 0.0, 0.0],
1224                MessageKind::Regular,
1225                "test-model",
1226                0,
1227                Some("preference"),
1228                None,
1229            )
1230            .await
1231            .unwrap();
1232
1233        let results = store
1234            .search_collection(COLLECTION_NAME, &[1.0, 0.0, 0.0, 0.0], 1, None)
1235            .await
1236            .unwrap();
1237        assert_eq!(results.len(), 1);
1238        assert_eq!(
1239            results[0].payload.get("category").and_then(|v| v.as_str()),
1240            Some("preference")
1241        );
1242    }
1243
1244    /// `store_with_category(None)` must omit the `category` field entirely (no false positives
1245    /// on category filters for pre-existing memories).
1246    #[tokio::test]
1247    async fn embedding_store_store_with_category_none_omits_payload_field() {
1248        let (store, sqlite) = setup_with_store().await;
1249        let cid = sqlite.create_conversation().await.unwrap();
1250        let msg_id = sqlite.save_message(cid, "user", "no cat").await.unwrap();
1251
1252        store
1253            .store_with_category(
1254                msg_id,
1255                cid,
1256                "user",
1257                vec![0.0, 1.0, 0.0, 0.0],
1258                MessageKind::Regular,
1259                "m",
1260                0,
1261                None,
1262                None,
1263            )
1264            .await
1265            .unwrap();
1266
1267        let results = store
1268            .search_collection(COLLECTION_NAME, &[0.0, 1.0, 0.0, 0.0], 1, None)
1269            .await
1270            .unwrap();
1271        assert_eq!(results.len(), 1);
1272        assert!(!results[0].payload.contains_key("category"));
1273    }
1274
1275    /// Issue #6553: `search_collection` must clamp an oversized `limit` to `MAX_SEARCH_LIMIT`
1276    /// internally, rather than relying on every caller to clamp before calling, and must log a
1277    /// one-shot warning so a config-driven candidate pool silently shrunk by the clamp is
1278    /// observable (critic finding S1).
1279    #[tokio::test]
1280    #[tracing_test::traced_test]
1281    async fn embedding_store_search_collection_clamps_oversized_limit() {
1282        let (store, _sqlite) = setup_with_store().await;
1283
1284        for _ in 0..(crate::MAX_SEARCH_LIMIT + 10) {
1285            store
1286                .store_to_collection(
1287                    COLLECTION_NAME,
1288                    serde_json::json!({}),
1289                    vec![1.0, 0.0, 0.0, 0.0],
1290                )
1291                .await
1292                .unwrap();
1293        }
1294
1295        let results = store
1296            .search_collection(COLLECTION_NAME, &[1.0, 0.0, 0.0, 0.0], usize::MAX, None)
1297            .await
1298            .unwrap();
1299        assert_eq!(results.len(), crate::MAX_SEARCH_LIMIT);
1300        assert!(
1301            logs_contain("requested search limit exceeds MAX_SEARCH_LIMIT"),
1302            "expected a one-shot warn when the clamp actually reduces the requested limit"
1303        );
1304    }
1305
1306    /// `store_with_tool_context` must write `tool_name`, `exit_code`, and `timestamp` payload
1307    /// fields (#5486 shared `store_impl` helper must preserve this variant's fields).
1308    #[tokio::test]
1309    async fn embedding_store_store_with_tool_context_sets_payload_fields() {
1310        let (store, sqlite) = setup_with_store().await;
1311        let cid = sqlite.create_conversation().await.unwrap();
1312        let msg_id = sqlite
1313            .save_message(cid, "assistant", "ran a tool")
1314            .await
1315            .unwrap();
1316
1317        store
1318            .store_with_tool_context(
1319                msg_id,
1320                cid,
1321                "assistant",
1322                vec![0.0, 0.0, 1.0, 0.0],
1323                MessageKind::Regular,
1324                "m",
1325                0,
1326                "shell",
1327                Some(0),
1328                Some("2026-07-02T00:00:00Z"),
1329                None,
1330            )
1331            .await
1332            .unwrap();
1333
1334        let results = store
1335            .search_collection(COLLECTION_NAME, &[0.0, 0.0, 1.0, 0.0], 1, None)
1336            .await
1337            .unwrap();
1338        assert_eq!(results.len(), 1);
1339        let payload = &results[0].payload;
1340        assert_eq!(
1341            payload.get("tool_name").and_then(|v| v.as_str()),
1342            Some("shell")
1343        );
1344        assert_eq!(
1345            payload.get("exit_code").and_then(serde_json::Value::as_i64),
1346            Some(0)
1347        );
1348        assert_eq!(
1349            payload.get("timestamp").and_then(|v| v.as_str()),
1350            Some("2026-07-02T00:00:00Z")
1351        );
1352    }
1353
1354    #[tokio::test]
1355    async fn embedding_store_has_embedding_false_for_unknown() {
1356        let (store, sqlite) = setup_with_store().await;
1357        let cid = sqlite.create_conversation().await.unwrap();
1358        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1359        assert!(!store.has_embedding(msg_id).await.unwrap());
1360    }
1361
1362    #[tokio::test]
1363    async fn embedding_store_has_embedding_true_after_store() {
1364        let (store, sqlite) = setup_with_store().await;
1365        let cid = sqlite.create_conversation().await.unwrap();
1366        let msg_id = sqlite.save_message(cid, "user", "hello").await.unwrap();
1367
1368        store
1369            .store(
1370                msg_id,
1371                cid,
1372                "user",
1373                vec![0.0, 1.0, 0.0, 0.0],
1374                MessageKind::Regular,
1375                "test-model",
1376                0,
1377                None,
1378            )
1379            .await
1380            .unwrap();
1381
1382        assert!(store.has_embedding(msg_id).await.unwrap());
1383    }
1384
1385    #[tokio::test]
1386    async fn embedding_store_search_with_conversation_filter() {
1387        let (store, sqlite) = setup_with_store().await;
1388        let cid1 = sqlite.create_conversation().await.unwrap();
1389        let cid2 = sqlite.create_conversation().await.unwrap();
1390        let msg1 = sqlite.save_message(cid1, "user", "msg1").await.unwrap();
1391        let msg2 = sqlite.save_message(cid2, "user", "msg2").await.unwrap();
1392
1393        store
1394            .store(
1395                msg1,
1396                cid1,
1397                "user",
1398                vec![1.0, 0.0, 0.0, 0.0],
1399                MessageKind::Regular,
1400                "m",
1401                0,
1402                None,
1403            )
1404            .await
1405            .unwrap();
1406        store
1407            .store(
1408                msg2,
1409                cid2,
1410                "user",
1411                vec![1.0, 0.0, 0.0, 0.0],
1412                MessageKind::Regular,
1413                "m",
1414                0,
1415                None,
1416            )
1417            .await
1418            .unwrap();
1419
1420        let results = store
1421            .search(
1422                &[1.0, 0.0, 0.0, 0.0],
1423                10,
1424                Some(SearchFilter {
1425                    conversation_id: Some(cid1),
1426                    role: None,
1427                    category: None,
1428                }),
1429            )
1430            .await
1431            .unwrap();
1432        assert_eq!(results.len(), 1);
1433        assert_eq!(results[0].conversation_id, cid1);
1434    }
1435
1436    #[tokio::test]
1437    async fn unique_constraint_on_message_chunk_and_model() {
1438        let (sqlite, pool) = setup().await;
1439        let cid = sqlite.create_conversation().await.unwrap();
1440        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1441
1442        let point_id1 = uuid::Uuid::new_v4().to_string();
1443        zeph_db::query(sql!(
1444            "INSERT INTO embeddings_metadata \
1445             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1446             VALUES (?, ?, ?, ?, ?)"
1447        ))
1448        .bind(msg_id)
1449        .bind(0_i64)
1450        .bind(&point_id1)
1451        .bind(768_i64)
1452        .bind("qwen3-embedding")
1453        .execute(&pool)
1454        .await
1455        .unwrap();
1456
1457        // Same (message_id, chunk_index, model) — must fail.
1458        let point_id2 = uuid::Uuid::new_v4().to_string();
1459        let result = zeph_db::query(sql!(
1460            "INSERT INTO embeddings_metadata \
1461             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1462             VALUES (?, ?, ?, ?, ?)"
1463        ))
1464        .bind(msg_id)
1465        .bind(0_i64)
1466        .bind(&point_id2)
1467        .bind(768_i64)
1468        .bind("qwen3-embedding")
1469        .execute(&pool)
1470        .await;
1471        assert!(result.is_err());
1472
1473        // Different chunk_index — must succeed.
1474        let point_id3 = uuid::Uuid::new_v4().to_string();
1475        zeph_db::query(sql!(
1476            "INSERT INTO embeddings_metadata \
1477             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1478             VALUES (?, ?, ?, ?, ?)"
1479        ))
1480        .bind(msg_id)
1481        .bind(1_i64)
1482        .bind(&point_id3)
1483        .bind(768_i64)
1484        .bind("qwen3-embedding")
1485        .execute(&pool)
1486        .await
1487        .unwrap();
1488    }
1489
1490    #[tokio::test]
1491    async fn get_vectors_for_messages_returns_correct_vectors() {
1492        let (store, sqlite) = setup_with_store().await;
1493        let cid = sqlite.create_conversation().await.unwrap();
1494        let msg1 = sqlite.save_message(cid, "user", "hello").await.unwrap();
1495        let msg2 = sqlite.save_message(cid, "user", "world").await.unwrap();
1496
1497        store
1498            .store(
1499                msg1,
1500                cid,
1501                "user",
1502                vec![1.0, 0.0, 0.0, 0.0],
1503                MessageKind::Regular,
1504                "m",
1505                0,
1506                None,
1507            )
1508            .await
1509            .unwrap();
1510        store
1511            .store(
1512                msg2,
1513                cid,
1514                "user",
1515                vec![0.0, 1.0, 0.0, 0.0],
1516                MessageKind::Regular,
1517                "m",
1518                0,
1519                None,
1520            )
1521            .await
1522            .unwrap();
1523
1524        let result = store.get_vectors_for_messages(&[msg1, msg2]).await.unwrap();
1525        assert_eq!(result.len(), 2);
1526        let v1 = result.get(&msg1).unwrap();
1527        let v2 = result.get(&msg2).unwrap();
1528        assert!((v1[0] - 1.0).abs() < f32::EPSILON);
1529        assert!((v2[1] - 1.0).abs() < f32::EPSILON);
1530    }
1531
1532    #[tokio::test]
1533    async fn get_vectors_for_messages_missing_id_is_dropped() {
1534        let (store, sqlite) = setup_with_store().await;
1535        let cid = sqlite.create_conversation().await.unwrap();
1536        let msg1 = sqlite.save_message(cid, "user", "present").await.unwrap();
1537        let msg_absent = MessageId(99_999);
1538
1539        store
1540            .store(
1541                msg1,
1542                cid,
1543                "user",
1544                vec![1.0, 0.0, 0.0, 0.0],
1545                MessageKind::Regular,
1546                "m",
1547                0,
1548                None,
1549            )
1550            .await
1551            .unwrap();
1552
1553        let result = store
1554            .get_vectors_for_messages(&[msg1, msg_absent])
1555            .await
1556            .unwrap();
1557        assert_eq!(result.len(), 1);
1558        assert!(result.contains_key(&msg1));
1559        assert!(!result.contains_key(&msg_absent));
1560    }
1561
1562    #[tokio::test]
1563    async fn get_vectors_for_messages_empty_input() {
1564        let (store, _sqlite) = setup_with_store().await;
1565        let result = store.get_vectors_for_messages(&[]).await.unwrap();
1566        assert!(result.is_empty());
1567    }
1568
1569    #[tokio::test]
1570    async fn get_vectors_for_messages_chunk_index_0_only() {
1571        // Store chunk_index=0 and chunk_index=1; only chunk_index=0 should be returned.
1572        let (store, sqlite) = setup_with_store().await;
1573        let cid = sqlite.create_conversation().await.unwrap();
1574        let msg = sqlite.save_message(cid, "user", "chunked").await.unwrap();
1575
1576        store
1577            .store(
1578                msg,
1579                cid,
1580                "user",
1581                vec![1.0, 0.0, 0.0, 0.0],
1582                MessageKind::Regular,
1583                "m",
1584                0,
1585                None,
1586            )
1587            .await
1588            .unwrap();
1589        store
1590            .store(
1591                msg,
1592                cid,
1593                "user",
1594                vec![0.0, 0.0, 1.0, 0.0],
1595                MessageKind::Regular,
1596                "m",
1597                1,
1598                None,
1599            )
1600            .await
1601            .unwrap();
1602
1603        let result = store.get_vectors_for_messages(&[msg]).await.unwrap();
1604        assert_eq!(result.len(), 1);
1605        // Must be the chunk_index=0 vector
1606        let v = result.get(&msg).unwrap();
1607        assert!(
1608            (v[0] - 1.0).abs() < f32::EPSILON,
1609            "expected chunk_index=0 vector"
1610        );
1611    }
1612
1613    /// `delete_by_message_ids` resolves `message_id → qdrant_point_id` via
1614    /// `embeddings_metadata` and deletes the matching vectors.
1615    ///
1616    /// Verifies: (a) the correct point id is targeted, (b) `embeddings_metadata`
1617    /// rows are NOT removed (CASCADE handles that on hard-delete later), and (c) the
1618    /// method returns the number of point IDs found.
1619    #[tokio::test]
1620    async fn embedding_store_delete_by_message_ids_resolves_via_metadata() {
1621        let (store, sqlite) = setup_with_store().await;
1622        let cid = sqlite.create_conversation().await.unwrap();
1623        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1624
1625        // Store a vector so embeddings_metadata gets a row.
1626        store
1627            .store(
1628                msg_id,
1629                cid,
1630                "user",
1631                vec![1.0, 0.0, 0.0, 0.0],
1632                MessageKind::Regular,
1633                "test-model",
1634                0,
1635                None,
1636            )
1637            .await
1638            .unwrap();
1639
1640        // Confirm the metadata row exists before deletion.
1641        assert!(store.has_embedding(msg_id).await.unwrap());
1642
1643        // Delete by message id — must succeed and return 1 (one point id resolved).
1644        let deleted = store.delete_by_message_ids(&[msg_id]).await.unwrap();
1645        assert_eq!(deleted, 1, "one point id should have been targeted");
1646
1647        // embeddings_metadata rows must still be present (CASCADE removes them later).
1648        let pool = sqlite.pool().clone();
1649        let row: (i64,) = zeph_db::query_as(sql!(
1650            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
1651        ))
1652        .bind(msg_id)
1653        .fetch_one(&pool)
1654        .await
1655        .unwrap();
1656        assert_eq!(
1657            row.0, 1,
1658            "embeddings_metadata row must survive delete_by_message_ids"
1659        );
1660    }
1661
1662    /// `delete_by_message_ids` is a no-op when the slice is empty.
1663    #[tokio::test]
1664    async fn embedding_store_delete_by_message_ids_empty_slice_is_noop() {
1665        let (store, _sqlite) = setup_with_store().await;
1666        let deleted = store.delete_by_message_ids(&[]).await.unwrap();
1667        assert_eq!(deleted, 0);
1668    }
1669}
zeph_memory/embedding_store.rs

zeph_memory/
embedding_store.rs
