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    /// Ensure a named collection exists in Qdrant with the given vector size.
251    ///
252    /// # Errors
253    ///
254    /// Returns an error if Qdrant cannot be reached or collection creation fails.
255    pub async fn ensure_named_collection(
256        &self,
257        name: &str,
258        vector_size: u64,
259    ) -> Result<(), MemoryError> {
260        self.ops.ensure_collection(name, vector_size).await?;
261        Ok(())
262    }
263
264    /// Store a vector in a named Qdrant collection with arbitrary payload.
265    ///
266    /// Returns the UUID of the newly created point.
267    ///
268    /// # Errors
269    ///
270    /// Returns an error if the Qdrant upsert fails.
271    pub async fn store_to_collection(
272        &self,
273        collection: &str,
274        payload: serde_json::Value,
275        vector: Vec<f32>,
276    ) -> Result<String, MemoryError> {
277        let point_id = uuid::Uuid::new_v4().to_string();
278        let payload_map: std::collections::HashMap<String, serde_json::Value> =
279            serde_json::from_value(payload)?;
280        let point = VectorPoint {
281            id: point_id.clone(),
282            vector,
283            payload: payload_map,
284        };
285        self.ops.upsert(collection, vec![point]).await?;
286        Ok(point_id)
287    }
288
289    /// Search a named Qdrant collection, returning scored points with payloads.
290    ///
291    /// # Errors
292    ///
293    /// Returns an error if the Qdrant search fails.
294    pub async fn search_collection(
295        &self,
296        collection: &str,
297        query_vector: &[f32],
298        limit: usize,
299        filter: Option<VectorFilter>,
300    ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
301        let limit_u64 = u64::try_from(limit)?;
302        let results = self
303            .ops
304            .search(collection, query_vector.to_vec(), limit_u64, filter)
305            .await?;
306        Ok(results)
307    }
308
309    /// Fetch raw vectors for the given message IDs from the `SQLite` vector store.
310    ///
311    /// Returns an empty map when using Qdrant backend (vectors not locally stored).
312    ///
313    /// # Errors
314    ///
315    /// Returns an error if the `SQLite` query fails.
316    pub async fn get_vectors(
317        &self,
318        ids: &[MessageId],
319    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
320        if ids.is_empty() {
321            return Ok(std::collections::HashMap::new());
322        }
323
324        let placeholders: String = ids.iter().map(|_| "?").collect::<Vec<_>>().join(",");
325        let query = format!(
326            "SELECT em.message_id, vp.vector \
327             FROM embeddings_metadata em \
328             JOIN vector_points vp ON vp.id = em.qdrant_point_id \
329             WHERE em.message_id IN ({placeholders})"
330        );
331        let mut q = sqlx::query_as::<_, (MessageId, Vec<u8>)>(&query);
332        for &id in ids {
333            q = q.bind(id);
334        }
335
336        let rows = q.fetch_all(&self.pool).await.unwrap_or_default();
337
338        let map = rows
339            .into_iter()
340            .filter_map(|(msg_id, blob)| {
341                if blob.len() % 4 != 0 {
342                    return None;
343                }
344                let vec: Vec<f32> = blob
345                    .chunks_exact(4)
346                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
347                    .collect();
348                Some((msg_id, vec))
349            })
350            .collect();
351
352        Ok(map)
353    }
354
355    /// Check whether an embedding already exists for the given message ID.
356    ///
357    /// # Errors
358    ///
359    /// Returns an error if the `SQLite` query fails.
360    pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
361        let row: (i64,) =
362            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
363                .bind(message_id)
364                .fetch_one(&self.pool)
365                .await?;
366
367        Ok(row.0 > 0)
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374    use crate::in_memory_store::InMemoryVectorStore;
375    use crate::sqlite::SqliteStore;
376
377    async fn setup() -> (SqliteStore, SqlitePool) {
378        let store = SqliteStore::new(":memory:").await.unwrap();
379        let pool = store.pool().clone();
380        (store, pool)
381    }
382
383    async fn setup_with_store() -> (EmbeddingStore, SqliteStore) {
384        let sqlite = SqliteStore::new(":memory:").await.unwrap();
385        let pool = sqlite.pool().clone();
386        let mem_store = Box::new(InMemoryVectorStore::new());
387        let embedding_store = EmbeddingStore::with_store(mem_store, pool);
388        // Create collection first
389        embedding_store.ensure_collection(4).await.unwrap();
390        (embedding_store, sqlite)
391    }
392
393    #[tokio::test]
394    async fn has_embedding_returns_false_when_none() {
395        let (_store, pool) = setup().await;
396
397        let row: (i64,) =
398            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
399                .bind(999_i64)
400                .fetch_one(&pool)
401                .await
402                .unwrap();
403
404        assert_eq!(row.0, 0);
405    }
406
407    #[tokio::test]
408    async fn insert_and_query_embeddings_metadata() {
409        let (sqlite, pool) = setup().await;
410        let cid = sqlite.create_conversation().await.unwrap();
411        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
412
413        let point_id = uuid::Uuid::new_v4().to_string();
414        sqlx::query(
415            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
416             VALUES (?, ?, ?, ?)",
417        )
418        .bind(msg_id)
419        .bind(&point_id)
420        .bind(768_i64)
421        .bind("qwen3-embedding")
422        .execute(&pool)
423        .await
424        .unwrap();
425
426        let row: (i64,) =
427            sqlx::query_as("SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?")
428                .bind(msg_id)
429                .fetch_one(&pool)
430                .await
431                .unwrap();
432        assert_eq!(row.0, 1);
433    }
434
435    #[tokio::test]
436    async fn embedding_store_search_empty_returns_empty() {
437        let (store, _sqlite) = setup_with_store().await;
438        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 10, None).await.unwrap();
439        assert!(results.is_empty());
440    }
441
442    #[tokio::test]
443    async fn embedding_store_store_and_search() {
444        let (store, sqlite) = setup_with_store().await;
445        let cid = sqlite.create_conversation().await.unwrap();
446        let msg_id = sqlite
447            .save_message(cid, "user", "test message")
448            .await
449            .unwrap();
450
451        store
452            .store(
453                msg_id,
454                cid,
455                "user",
456                vec![1.0, 0.0, 0.0, 0.0],
457                MessageKind::Regular,
458                "test-model",
459            )
460            .await
461            .unwrap();
462
463        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5, None).await.unwrap();
464        assert_eq!(results.len(), 1);
465        assert_eq!(results[0].message_id, msg_id);
466        assert_eq!(results[0].conversation_id, cid);
467        assert!((results[0].score - 1.0).abs() < 0.001);
468    }
469
470    #[tokio::test]
471    async fn embedding_store_has_embedding_false_for_unknown() {
472        let (store, sqlite) = setup_with_store().await;
473        let cid = sqlite.create_conversation().await.unwrap();
474        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
475        assert!(!store.has_embedding(msg_id).await.unwrap());
476    }
477
478    #[tokio::test]
479    async fn embedding_store_has_embedding_true_after_store() {
480        let (store, sqlite) = setup_with_store().await;
481        let cid = sqlite.create_conversation().await.unwrap();
482        let msg_id = sqlite.save_message(cid, "user", "hello").await.unwrap();
483
484        store
485            .store(
486                msg_id,
487                cid,
488                "user",
489                vec![0.0, 1.0, 0.0, 0.0],
490                MessageKind::Regular,
491                "test-model",
492            )
493            .await
494            .unwrap();
495
496        assert!(store.has_embedding(msg_id).await.unwrap());
497    }
498
499    #[tokio::test]
500    async fn embedding_store_search_with_conversation_filter() {
501        let (store, sqlite) = setup_with_store().await;
502        let cid1 = sqlite.create_conversation().await.unwrap();
503        let cid2 = sqlite.create_conversation().await.unwrap();
504        let msg1 = sqlite.save_message(cid1, "user", "msg1").await.unwrap();
505        let msg2 = sqlite.save_message(cid2, "user", "msg2").await.unwrap();
506
507        store
508            .store(
509                msg1,
510                cid1,
511                "user",
512                vec![1.0, 0.0, 0.0, 0.0],
513                MessageKind::Regular,
514                "m",
515            )
516            .await
517            .unwrap();
518        store
519            .store(
520                msg2,
521                cid2,
522                "user",
523                vec![1.0, 0.0, 0.0, 0.0],
524                MessageKind::Regular,
525                "m",
526            )
527            .await
528            .unwrap();
529
530        let results = store
531            .search(
532                &[1.0, 0.0, 0.0, 0.0],
533                10,
534                Some(SearchFilter {
535                    conversation_id: Some(cid1),
536                    role: None,
537                }),
538            )
539            .await
540            .unwrap();
541        assert_eq!(results.len(), 1);
542        assert_eq!(results[0].conversation_id, cid1);
543    }
544
545    #[tokio::test]
546    async fn unique_constraint_on_message_and_model() {
547        let (sqlite, pool) = setup().await;
548        let cid = sqlite.create_conversation().await.unwrap();
549        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
550
551        let point_id1 = uuid::Uuid::new_v4().to_string();
552        sqlx::query(
553            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
554             VALUES (?, ?, ?, ?)",
555        )
556        .bind(msg_id)
557        .bind(&point_id1)
558        .bind(768_i64)
559        .bind("qwen3-embedding")
560        .execute(&pool)
561        .await
562        .unwrap();
563
564        let point_id2 = uuid::Uuid::new_v4().to_string();
565        let result = sqlx::query(
566            "INSERT INTO embeddings_metadata (message_id, qdrant_point_id, dimensions, model) \
567             VALUES (?, ?, ?, ?)",
568        )
569        .bind(msg_id)
570        .bind(&point_id2)
571        .bind(768_i64)
572        .bind("qwen3-embedding")
573        .execute(&pool)
574        .await;
575
576        assert!(result.is_err());
577    }
578}
zeph_memory/embedding_store.rs

zeph_memory/
embedding_store.rs
