zeph_memory/
embedding_store.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4pub use qdrant_client::qdrant::Filter;
5use sqlx::SqlitePool;
6
7use crate::error::MemoryError;
8use crate::qdrant_ops::QdrantOps;
9use crate::sqlite_vector_store::SqliteVectorStore;
10use crate::types::{ConversationId, MessageId};
11use crate::vector_store::{FieldCondition, FieldValue, VectorFilter, VectorPoint, VectorStore};
12
13/// Distinguishes regular messages from summaries when storing embeddings.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum MessageKind {
16    Regular,
17    Summary,
18}
19
20impl MessageKind {
21    #[must_use]
22    pub fn is_summary(self) -> bool {
23        matches!(self, Self::Summary)
24    }
25}
26
27const COLLECTION_NAME: &str = "zeph_conversations";
28
29/// Ensure a Qdrant collection exists with cosine distance vectors.
30///
31/// Idempotent: no-op if the collection already exists.
32///
33/// # Errors
34///
35/// Returns an error if Qdrant cannot be reached or collection creation fails.
36pub async fn ensure_qdrant_collection(
37    ops: &QdrantOps,
38    collection: &str,
39    vector_size: u64,
40) -> Result<(), Box<qdrant_client::QdrantError>> {
41    ops.ensure_collection(collection, vector_size).await
42}
43
44pub struct EmbeddingStore {
45    ops: Box<dyn VectorStore>,
46    collection: String,
47    pool: SqlitePool,
48}
49
50impl std::fmt::Debug for EmbeddingStore {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("EmbeddingStore")
53            .field("collection", &self.collection)
54            .finish_non_exhaustive()
55    }
56}
57
58#[derive(Debug)]
59pub struct SearchFilter {
60    pub conversation_id: Option<ConversationId>,
61    pub role: Option<String>,
62}
63
64#[derive(Debug)]
65pub struct SearchResult {
66    pub message_id: MessageId,
67    pub conversation_id: ConversationId,
68    pub score: f32,
69}
70
71impl EmbeddingStore {
72    /// Create a new `EmbeddingStore` connected to the given Qdrant URL.
73    ///
74    /// The `pool` is used for `SQLite` metadata operations on the `embeddings_metadata`
75    /// table (which must already exist via sqlx migrations).
76    ///
77    /// # Errors
78    ///
79    /// Returns an error if the Qdrant client cannot be created.
80    pub fn new(url: &str, pool: SqlitePool) -> Result<Self, MemoryError> {
81        let ops = QdrantOps::new(url).map_err(MemoryError::Qdrant)?;
82
83        Ok(Self {
84            ops: Box::new(ops),
85            collection: COLLECTION_NAME.into(),
86            pool,
87        })
88    }
89
90    /// Create a new `EmbeddingStore` backed by `SQLite` for vector storage.
91    ///
92    /// Uses the same pool for both vector data and metadata. No external Qdrant required.
93    #[must_use]
94    pub fn new_sqlite(pool: SqlitePool) -> Self {
95        let ops = SqliteVectorStore::new(pool.clone());
96        Self {
97            ops: Box::new(ops),
98            collection: COLLECTION_NAME.into(),
99            pool,
100        }
101    }
102
103    #[must_use]
104    pub fn with_store(store: Box<dyn VectorStore>, pool: SqlitePool) -> Self {
105        Self {
106            ops: store,
107            collection: COLLECTION_NAME.into(),
108            pool,
109        }
110    }
111
112    pub async fn health_check(&self) -> bool {
113        self.ops.health_check().await.unwrap_or(false)
114    }
115
116    /// Ensure the collection exists in Qdrant with the given vector size.
117    ///
118    /// Idempotent: no-op if the collection already exists.
119    ///
120    /// # Errors
121    ///
122    /// Returns an error if Qdrant cannot be reached or collection creation fails.
123    pub async fn ensure_collection(&self, vector_size: u64) -> Result<(), MemoryError> {
124        self.ops
125            .ensure_collection(&self.collection, vector_size)
126            .await?;
127        Ok(())
128    }
129
130    /// Store a vector in Qdrant and persist metadata to `SQLite`.
131    ///
132    /// Returns the UUID of the newly created Qdrant point.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
137    pub async fn store(
138        &self,
139        message_id: MessageId,
140        conversation_id: ConversationId,
141        role: &str,
142        vector: Vec<f32>,
143        kind: MessageKind,
144        model: &str,
145    ) -> Result<String, MemoryError> {
146        let point_id = uuid::Uuid::new_v4().to_string();
147        let dimensions = i64::try_from(vector.len())?;
148
149        let payload = std::collections::HashMap::from([
150            ("message_id".to_owned(), serde_json::json!(message_id.0)),
151            (
152                "conversation_id".to_owned(),
153                serde_json::json!(conversation_id.0),
154            ),
155            ("role".to_owned(), serde_json::json!(role)),
156            (
157                "is_summary".to_owned(),
158                serde_json::json!(kind.is_summary()),
159            ),
160        ]);
161
162        let point = VectorPoint {
163            id: point_id.clone(),
164            vector,
165            payload,
166        };
167
168        self.ops.upsert(&self.collection, vec![point]).await?;
169
170        sqlx::query(
171            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
172             VALUES (?, ?, ?, ?) \
173             ON CONFLICT(message_id, model) DO UPDATE SET \
174             qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions",
175        )
176        .bind(message_id)
177        .bind(&point_id)
178        .bind(dimensions)
179        .bind(model)
180        .execute(&self.pool)
181        .await?;
182
183        Ok(point_id)
184    }
185
186    /// Search for similar vectors in Qdrant, returning up to `limit` results.
187    ///
188    /// # Errors
189    ///
190    /// Returns an error if the Qdrant search fails.
191    pub async fn search(
192        &self,
193        query_vector: &[f32],
194        limit: usize,
195        filter: Option<SearchFilter>,
196    ) -> Result<Vec<SearchResult>, MemoryError> {
197        let limit_u64 = u64::try_from(limit)?;
198
199        let vector_filter = filter.as_ref().and_then(|f| {
200            let mut must = Vec::new();
201            if let Some(cid) = f.conversation_id {
202                must.push(FieldCondition {
203                    field: "conversation_id".into(),
204                    value: FieldValue::Integer(cid.0),
205                });
206            }
207            if let Some(ref role) = f.role {
208                must.push(FieldCondition {
209                    field: "role".into(),
210                    value: FieldValue::Text(role.clone()),
211                });
212            }
213            if must.is_empty() {
214                None
215            } else {
216                Some(VectorFilter {
217                    must,
218                    must_not: vec![],
219                })
220            }
221        });
222
223        let results = self
224            .ops
225            .search(
226                &self.collection,
227                query_vector.to_vec(),
228                limit_u64,
229                vector_filter,
230            )
231            .await?;
232
233        let search_results = results
234            .into_iter()
235            .filter_map(|point| {
236                let message_id = MessageId(point.payload.get("message_id")?.as_i64()?);
237                let conversation_id =
238                    ConversationId(point.payload.get("conversation_id")?.as_i64()?);
239                Some(SearchResult {
240                    message_id,
241                    conversation_id,
242                    score: point.score,
243                })
244            })
245            .collect();
246
247        Ok(search_results)
248    }
249
250    /// Check whether a named collection exists in the vector store.
251    ///
252    /// # Errors
253    ///
254    /// Returns an error if the store backend cannot be reached.
255    pub async fn collection_exists(&self, name: &str) -> Result<bool, MemoryError> {
256        self.ops.collection_exists(name).await.map_err(Into::into)
257    }
258
259    /// Ensure a named collection exists in Qdrant with the given vector size.
260    ///
261    /// # Errors
262    ///
263    /// Returns an error if Qdrant cannot be reached or collection creation fails.
264    pub async fn ensure_named_collection(
265        &self,
266        name: &str,
267        vector_size: u64,
268    ) -> Result<(), MemoryError> {
269        self.ops.ensure_collection(name, vector_size).await?;
270        Ok(())
271    }
272
273    /// Store a vector in a named Qdrant collection with arbitrary payload.
274    ///
275    /// Returns the UUID of the newly created point.
276    ///
277    /// # Errors
278    ///
279    /// Returns an error if the Qdrant upsert fails.
280    pub async fn store_to_collection(
281        &self,
282        collection: &str,
283        payload: serde_json::Value,
284        vector: Vec<f32>,
285    ) -> Result<String, MemoryError> {
286        let point_id = uuid::Uuid::new_v4().to_string();
287        let payload_map: std::collections::HashMap<String, serde_json::Value> =
288            serde_json::from_value(payload)?;
289        let point = VectorPoint {
290            id: point_id.clone(),
291            vector,
292            payload: payload_map,
293        };
294        self.ops.upsert(collection, vec![point]).await?;
295        Ok(point_id)
296    }
297
298    /// Search a named Qdrant collection, returning scored points with payloads.
299    ///
300    /// # Errors
301    ///
302    /// Returns an error if the Qdrant search fails.
303    pub async fn search_collection(
304        &self,
305        collection: &str,
306        query_vector: &[f32],
307        limit: usize,
308        filter: Option<VectorFilter>,
309    ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
310        let limit_u64 = u64::try_from(limit)?;
311        let results = self
312            .ops
313            .search(collection, query_vector.to_vec(), limit_u64, filter)
314            .await?;
315        Ok(results)
316    }
317
318    /// Fetch raw vectors for the given message IDs from the `SQLite` vector store.
319    ///
320    /// Returns an empty map when using Qdrant backend (vectors not locally stored).
321    ///
322    /// # Errors
323    ///
324    /// Returns an error if the `SQLite` query fails.
325    pub async fn get_vectors(
326        &self,
327        ids: &[MessageId],
328    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
329        if ids.is_empty() {
330            return Ok(std::collections::HashMap::new());
331        }
332
333        let placeholders: String = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
334        let query = format!(
335            "SELECT em.message_id, vp.vector \
336             FROM embeddings_metadata em \
337             JOIN vector_points vp ON vp.id = em.qdrant_point_id \
338             WHERE em.message_id IN ({placeholders})"
339        );
340        let mut q = sqlx::query_as::<_, (MessageId, Vec<u8>)>(&query);
341        for &id in ids {
342            q = q.bind(id);
343        }
344
345        let rows = q.fetch_all(&self.pool).await.unwrap_or_default();
346
347        let map = rows
348            .into_iter()
349            .filter_map(|(msg_id, blob)| {
350                if blob.len() % 4 != 0 {
351                    return None;
352                }
353                let vec: Vec<f32> = blob
354                    .chunks_exact(4)
355                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
356                    .collect();
357                Some((msg_id, vec))
358            })
359            .collect();
360
361        Ok(map)
362    }
363
364    /// Check whether an embedding already exists for the given message ID.
365    ///
366    /// # Errors
367    ///
368    /// Returns an error if the `SQLite` query fails.
369    pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
370        let row: (i64,) =
371            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
372                .bind(message_id)
373                .fetch_one(&self.pool)
374                .await?;
375
376        Ok(row.0 > 0)
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383    use crate::in_memory_store::InMemoryVectorStore;
384    use crate::sqlite::SqliteStore;
385
386    async fn setup() -> (SqliteStore, SqlitePool) {
387        let store = SqliteStore::new(":memory:").await.unwrap();
388        let pool = store.pool().clone();
389        (store, pool)
390    }
391
392    async fn setup_with_store() -> (EmbeddingStore, SqliteStore) {
393        let sqlite = SqliteStore::new(":memory:").await.unwrap();
394        let pool = sqlite.pool().clone();
395        let mem_store = Box::new(InMemoryVectorStore::new());
396        let embedding_store = EmbeddingStore::with_store(mem_store, pool);
397        // Create collection first
398        embedding_store.ensure_collection(4).await.unwrap();
399        (embedding_store, sqlite)
400    }
401
402    #[tokio::test]
403    async fn has_embedding_returns_false_when_none() {
404        let (_store, pool) = setup().await;
405
406        let row: (i64,) =
407            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
408                .bind(999_i64)
409                .fetch_one(&pool)
410                .await
411                .unwrap();
412
413        assert_eq!(row.0, 0);
414    }
415
416    #[tokio::test]
417    async fn insert_and_query_embeddings_metadata() {
418        let (sqlite, pool) = setup().await;
419        let cid = sqlite.create_conversation().await.unwrap();
420        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
421
422        let point_id = uuid::Uuid::new_v4().to_string();
423        sqlx::query(
424            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
425             VALUES (?, ?, ?, ?)",
426        )
427        .bind(msg_id)
428        .bind(&point_id)
429        .bind(768_i64)
430        .bind("qwen3-embedding")
431        .execute(&pool)
432        .await
433        .unwrap();
434
435        let row: (i64,) =
436            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
437                .bind(msg_id)
438                .fetch_one(&pool)
439                .await
440                .unwrap();
441        assert_eq!(row.0, 1);
442    }
443
444    #[tokio::test]
445    async fn embedding_store_search_empty_returns_empty() {
446        let (store, _sqlite) = setup_with_store().await;
447        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 10, None).await.unwrap();
448        assert!(results.is_empty());
449    }
450
451    #[tokio::test]
452    async fn embedding_store_store_and_search() {
453        let (store, sqlite) = setup_with_store().await;
454        let cid = sqlite.create_conversation().await.unwrap();
455        let msg_id = sqlite
456            .save_message(cid, "user", "test message")
457            .await
458            .unwrap();
459
460        store
461            .store(
462                msg_id,
463                cid,
464                "user",
465                vec![1.0, 0.0, 0.0, 0.0],
466                MessageKind::Regular,
467                "test-model",
468            )
469            .await
470            .unwrap();
471
472        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5, None).await.unwrap();
473        assert_eq!(results.len(), 1);
474        assert_eq!(results[0].message_id, msg_id);
475        assert_eq!(results[0].conversation_id, cid);
476        assert!((results[0].score - 1.0).abs() < 0.001);
477    }
478
479    #[tokio::test]
480    async fn embedding_store_has_embedding_false_for_unknown() {
481        let (store, sqlite) = setup_with_store().await;
482        let cid = sqlite.create_conversation().await.unwrap();
483        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
484        assert!(!store.has_embedding(msg_id).await.unwrap());
485    }
486
487    #[tokio::test]
488    async fn embedding_store_has_embedding_true_after_store() {
489        let (store, sqlite) = setup_with_store().await;
490        let cid = sqlite.create_conversation().await.unwrap();
491        let msg_id = sqlite.save_message(cid, "user", "hello").await.unwrap();
492
493        store
494            .store(
495                msg_id,
496                cid,
497                "user",
498                vec![0.0, 1.0, 0.0, 0.0],
499                MessageKind::Regular,
500                "test-model",
501            )
502            .await
503            .unwrap();
504
505        assert!(store.has_embedding(msg_id).await.unwrap());
506    }
507
508    #[tokio::test]
509    async fn embedding_store_search_with_conversation_filter() {
510        let (store, sqlite) = setup_with_store().await;
511        let cid1 = sqlite.create_conversation().await.unwrap();
512        let cid2 = sqlite.create_conversation().await.unwrap();
513        let msg1 = sqlite.save_message(cid1, "user", "msg1").await.unwrap();
514        let msg2 = sqlite.save_message(cid2, "user", "msg2").await.unwrap();
515
516        store
517            .store(
518                msg1,
519                cid1,
520                "user",
521                vec![1.0, 0.0, 0.0, 0.0],
522                MessageKind::Regular,
523                "m",
524            )
525            .await
526            .unwrap();
527        store
528            .store(
529                msg2,
530                cid2,
531                "user",
532                vec![1.0, 0.0, 0.0, 0.0],
533                MessageKind::Regular,
534                "m",
535            )
536            .await
537            .unwrap();
538
539        let results = store
540            .search(
541                &[1.0, 0.0, 0.0, 0.0],
542                10,
543                Some(SearchFilter {
544                    conversation_id: Some(cid1),
545                    role: None,
546                }),
547            )
548            .await
549            .unwrap();
550        assert_eq!(results.len(), 1);
551        assert_eq!(results[0].conversation_id, cid1);
552    }
553
554    #[tokio::test]
555    async fn unique_constraint_on_message_and_model() {
556        let (sqlite, pool) = setup().await;
557        let cid = sqlite.create_conversation().await.unwrap();
558        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
559
560        let point_id1 = uuid::Uuid::new_v4().to_string();
561        sqlx::query(
562            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
563             VALUES (?, ?, ?, ?)",
564        )
565        .bind(msg_id)
566        .bind(&point_id1)
567        .bind(768_i64)
568        .bind("qwen3-embedding")
569        .execute(&pool)
570        .await
571        .unwrap();
572
573        let point_id2 = uuid::Uuid::new_v4().to_string();
574        let result = sqlx::query(
575            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
576             VALUES (?, ?, ?, ?)",
577        )
578        .bind(msg_id)
579        .bind(&point_id2)
580        .bind(768_i64)
581        .bind("qwen3-embedding")
582        .execute(&pool)
583        .await;
584
585        assert!(result.is_err());
586    }
587}
zeph_memory/embedding_store.rs

zeph_memory/
embedding_store.rs
