1pub 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29#[non_exhaustive]
30pub enum MessageKind {
31 Regular,
33 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#[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
62pub struct EmbeddingStore {
67 ops: Box<dyn VectorStore>,
68 collection: String,
69 pool: DbPool,
70 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#[derive(Debug)]
87pub struct SearchFilter {
88 pub conversation_id: Option<ConversationId>,
90 pub role: Option<String>,
92 pub category: Option<String>,
95}
96
97#[derive(Debug)]
99pub struct SearchResult {
100 pub message_id: MessageId,
102 pub conversation_id: ConversationId,
104 pub score: f32,
106}
107
108enum StoreExtra<'a> {
114 None,
116 Category(Option<&'a str>),
118 ToolContext {
120 tool_name: &'a str,
121 exit_code: Option<i32>,
122 timestamp: Option<&'a str>,
123 },
124}
125
126impl EmbeddingStore {
127 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 #[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 #[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 #[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 #[must_use]
189 pub fn db_instance_id(&self) -> &str {
190 &self.db_instance_id
191 }
192
193 #[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 #[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 self.ops
214 .create_keyword_indexes(&self.collection, &["category", "conversation_id", "role"])
215 .await?;
216 Ok(())
217 }
218
219 pub async fn ensure_collection_for_vector(&self, vector: &[f32]) -> Result<(), MemoryError> {
230 self.ensure_collection(vector.len() as u64).await
232 }
233
234 pub async fn ensure_named_collection_for_vector(
245 &self,
246 name: &str,
247 vector: &[f32],
248 ) -> Result<(), MemoryError> {
249 self.ensure_named_collection(name, vector.len() as u64)
251 .await
252 }
253
254 #[tracing::instrument(name = "memory.embed_store.store_with_tool_context", skip_all)]
264 #[allow(clippy::too_many_arguments)] 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 ) -> Result<String, MemoryError> {
278 self.store_impl(
279 message_id,
280 conversation_id,
281 role,
282 vector,
283 kind,
284 model,
285 chunk_index,
286 StoreExtra::ToolContext {
287 tool_name,
288 exit_code,
289 timestamp,
290 },
291 )
292 .await
293 }
294
295 #[tracing::instrument(name = "memory.embed_store.store", skip_all)]
306 #[allow(clippy::too_many_arguments)] pub async fn store(
308 &self,
309 message_id: MessageId,
310 conversation_id: ConversationId,
311 role: &str,
312 vector: Vec<f32>,
313 kind: MessageKind,
314 model: &str,
315 chunk_index: u32,
316 ) -> Result<String, MemoryError> {
317 self.store_impl(
318 message_id,
319 conversation_id,
320 role,
321 vector,
322 kind,
323 model,
324 chunk_index,
325 StoreExtra::None,
326 )
327 .await
328 }
329
330 #[tracing::instrument(name = "memory.embed_store.store_with_category", skip_all)]
344 #[allow(clippy::too_many_arguments)] pub async fn store_with_category(
346 &self,
347 message_id: MessageId,
348 conversation_id: ConversationId,
349 role: &str,
350 vector: Vec<f32>,
351 kind: MessageKind,
352 model: &str,
353 chunk_index: u32,
354 category: Option<&str>,
355 ) -> Result<String, MemoryError> {
356 self.store_impl(
357 message_id,
358 conversation_id,
359 role,
360 vector,
361 kind,
362 model,
363 chunk_index,
364 StoreExtra::Category(category),
365 )
366 .await
367 }
368
369 #[allow(clippy::too_many_arguments)] async fn store_impl(
374 &self,
375 message_id: MessageId,
376 conversation_id: ConversationId,
377 role: &str,
378 vector: Vec<f32>,
379 kind: MessageKind,
380 model: &str,
381 chunk_index: u32,
382 extra: StoreExtra<'_>,
383 ) -> Result<String, MemoryError> {
384 let point_id = uuid::Uuid::new_v4().to_string();
385 let dimensions = i64::try_from(vector.len())?;
386
387 let mut payload = std::collections::HashMap::from([
388 ("message_id".to_owned(), serde_json::json!(message_id.0)),
389 (
390 "conversation_id".to_owned(),
391 serde_json::json!(conversation_id.0),
392 ),
393 (
394 "db_instance_id".to_owned(),
395 serde_json::json!(self.db_instance_id),
396 ),
397 ("role".to_owned(), serde_json::json!(role)),
398 (
399 "is_summary".to_owned(),
400 serde_json::json!(kind.is_summary()),
401 ),
402 ]);
403 match extra {
404 StoreExtra::None => {}
405 StoreExtra::Category(category) => {
406 if let Some(cat) = category {
407 payload.insert("category".to_owned(), serde_json::json!(cat));
408 }
409 }
410 StoreExtra::ToolContext {
411 tool_name,
412 exit_code,
413 timestamp,
414 } => {
415 payload.insert("tool_name".to_owned(), serde_json::json!(tool_name));
416 if let Some(code) = exit_code {
417 payload.insert("exit_code".to_owned(), serde_json::json!(code));
418 }
419 if let Some(ts) = timestamp {
420 payload.insert("timestamp".to_owned(), serde_json::json!(ts));
421 }
422 }
423 }
424
425 let point = VectorPoint {
426 id: point_id.clone(),
427 vector,
428 payload,
429 };
430
431 self.ops.upsert(&self.collection, vec![point]).await?;
432
433 let chunk_index_i64 = i64::from(chunk_index);
434 zeph_db::query(sql!(
435 "INSERT INTO embeddings_metadata \
436 (message_id, chunk_index, qdrant_point_id, dimensions, model) \
437 VALUES (?, ?, ?, ?, ?) \
438 ON CONFLICT(message_id, chunk_index, model) DO UPDATE SET \
439 qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions"
440 ))
441 .bind(message_id)
442 .bind(chunk_index_i64)
443 .bind(&point_id)
444 .bind(dimensions)
445 .bind(model)
446 .execute(&self.pool)
447 .await?;
448
449 Ok(point_id)
450 }
451
452 #[tracing::instrument(name = "memory.embed_store.search", skip_all)]
458 pub async fn search(
459 &self,
460 query_vector: &[f32],
461 limit: usize,
462 filter: Option<SearchFilter>,
463 ) -> Result<Vec<SearchResult>, MemoryError> {
464 let limit_u64 = u64::try_from(limit)?;
465
466 let vector_filter = filter.as_ref().and_then(|f| {
467 let mut must = Vec::new();
468 if let Some(cid) = f.conversation_id {
469 must.push(FieldCondition {
470 field: "conversation_id".into(),
471 value: FieldValue::Integer(cid.0),
472 });
473 must.push(FieldCondition {
474 field: "db_instance_id".into(),
475 value: FieldValue::Text(self.db_instance_id.clone()),
476 });
477 }
478 if let Some(ref role) = f.role {
479 must.push(FieldCondition {
480 field: "role".into(),
481 value: FieldValue::Text(role.clone()),
482 });
483 }
484 if let Some(ref category) = f.category {
485 must.push(FieldCondition {
486 field: "category".into(),
487 value: FieldValue::Text(category.clone()),
488 });
489 }
490 if must.is_empty() {
491 None
492 } else {
493 Some(VectorFilter {
494 must,
495 must_not: vec![],
496 })
497 }
498 });
499
500 let results = self
501 .ops
502 .search(
503 &self.collection,
504 query_vector.to_vec(),
505 limit_u64,
506 vector_filter,
507 )
508 .await?;
509
510 let mut best: std::collections::HashMap<MessageId, SearchResult> =
513 std::collections::HashMap::new();
514 for point in results {
515 let Some(message_id) = point
516 .payload
517 .get("message_id")
518 .and_then(serde_json::Value::as_i64)
519 else {
520 continue;
521 };
522 let Some(conversation_id) = point
523 .payload
524 .get("conversation_id")
525 .and_then(serde_json::Value::as_i64)
526 else {
527 continue;
528 };
529 let message_id = MessageId(message_id);
530 let entry = best.entry(message_id).or_insert(SearchResult {
531 message_id,
532 conversation_id: ConversationId(conversation_id),
533 score: f32::NEG_INFINITY,
534 });
535 if point.score > entry.score {
536 entry.score = point.score;
537 }
538 }
539
540 let mut search_results: Vec<SearchResult> = best.into_values().collect();
541 search_results.sort_by(|a, b| {
542 b.score
543 .partial_cmp(&a.score)
544 .unwrap_or(std::cmp::Ordering::Equal)
545 });
546 search_results.truncate(limit);
547
548 Ok(search_results)
549 }
550
551 #[tracing::instrument(name = "memory.embed_store.collection_exists", skip_all)]
557 pub async fn collection_exists(&self, name: &str) -> Result<bool, MemoryError> {
558 self.ops.collection_exists(name).await.map_err(Into::into)
559 }
560
561 #[tracing::instrument(name = "memory.embed_store.ensure_named_collection", skip_all)]
567 pub async fn ensure_named_collection(
568 &self,
569 name: &str,
570 vector_size: u64,
571 ) -> Result<(), MemoryError> {
572 self.ops.ensure_collection(name, vector_size).await?;
573 Ok(())
574 }
575
576 #[tracing::instrument(name = "memory.embed_store.store_to_collection", skip_all)]
584 pub async fn store_to_collection(
585 &self,
586 collection: &str,
587 payload: serde_json::Value,
588 vector: Vec<f32>,
589 ) -> Result<String, MemoryError> {
590 let point_id = uuid::Uuid::new_v4().to_string();
591 let payload_map: std::collections::HashMap<String, serde_json::Value> =
592 serde_json::from_value(payload)?;
593 let point = VectorPoint {
594 id: point_id.clone(),
595 vector,
596 payload: payload_map,
597 };
598 self.ops.upsert(collection, vec![point]).await?;
599 Ok(point_id)
600 }
601
602 #[tracing::instrument(name = "memory.embed_store.upsert_to_collection", skip_all)]
610 pub async fn upsert_to_collection(
611 &self,
612 collection: &str,
613 point_id: &str,
614 payload: serde_json::Value,
615 vector: Vec<f32>,
616 ) -> Result<(), MemoryError> {
617 let payload_map: std::collections::HashMap<String, serde_json::Value> =
618 serde_json::from_value(payload)?;
619 let point = VectorPoint {
620 id: point_id.to_owned(),
621 vector,
622 payload: payload_map,
623 };
624 self.ops.upsert(collection, vec![point]).await?;
625 Ok(())
626 }
627
628 #[tracing::instrument(name = "memory.embed_store.search_collection", skip_all)]
634 pub async fn search_collection(
635 &self,
636 collection: &str,
637 query_vector: &[f32],
638 limit: usize,
639 filter: Option<VectorFilter>,
640 ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
641 let limit_u64 = u64::try_from(limit)?;
642 let results = self
643 .ops
644 .search(collection, query_vector.to_vec(), limit_u64, filter)
645 .await?;
646 Ok(results)
647 }
648
649 #[tracing::instrument(name = "memory.embed_store.scroll_all_entity_ids", skip_all)]
662 pub async fn scroll_all_entity_ids(
663 &self,
664 collection: &str,
665 ) -> Result<Vec<(String, i64)>, MemoryError> {
666 let rows = self
667 .ops
668 .scroll_all_with_point_ids(collection, "entity_id_str")
669 .await?;
670 let mut out = Vec::with_capacity(rows.len());
671 for (point_id, fields) in rows {
672 let Some(s) = fields.get("entity_id_str") else {
673 continue;
674 };
675 if let Ok(id) = s.parse::<i64>() {
676 out.push((point_id, id));
677 } else {
678 tracing::debug!(point_id, value = %s, "entity_id_str unparseable, skipping");
679 }
680 }
681 Ok(out)
682 }
683
684 #[tracing::instrument(name = "memory.embed_store.delete_from_collection", skip_all)]
693 pub async fn delete_from_collection(
694 &self,
695 collection: &str,
696 ids: Vec<String>,
697 ) -> Result<(), MemoryError> {
698 if ids.is_empty() {
699 return Ok(());
700 }
701 self.ops.delete_by_ids(collection, ids).await?;
702 Ok(())
703 }
704
705 #[tracing::instrument(name = "memory.embed_store.get_vectors_from_collection", skip_all)]
715 pub async fn get_vectors_from_collection(
716 &self,
717 collection: &str,
718 point_ids: &[String],
719 ) -> Result<std::collections::HashMap<String, Vec<f32>>, MemoryError> {
720 if point_ids.is_empty() {
721 return Ok(std::collections::HashMap::new());
722 }
723 match self.ops.get_points(collection, point_ids.to_vec()).await {
724 Ok(points) => Ok(points.into_iter().map(|p| (p.id, p.vector)).collect()),
725 Err(crate::VectorStoreError::Unsupported(_)) => Ok(std::collections::HashMap::new()),
726 Err(e) => Err(MemoryError::VectorStore(e)),
727 }
728 }
729
730 #[tracing::instrument(name = "memory.embed_store.get_vectors", skip_all)]
738 pub async fn get_vectors(
739 &self,
740 ids: &[MessageId],
741 ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
742 if ids.is_empty() {
743 return Ok(std::collections::HashMap::new());
744 }
745
746 let placeholders = zeph_db::placeholder_list(1, ids.len());
747 let query = format!(
748 "SELECT em.message_id, vp.vector \
749 FROM embeddings_metadata em \
750 JOIN vector_points vp ON vp.id = em.qdrant_point_id \
751 WHERE em.message_id IN ({placeholders}) AND em.chunk_index = 0"
752 );
753 let mut q = zeph_db::query_as::<_, (MessageId, Vec<u8>)>(sqlx::AssertSqlSafe(query));
754 for &id in ids {
755 q = q.bind(id);
756 }
757
758 let rows = q.fetch_all(&self.pool).await?;
759
760 let map = rows
761 .into_iter()
762 .filter_map(|(msg_id, blob)| {
763 if blob.len() % 4 != 0 {
764 return None;
765 }
766 let vec: Vec<f32> = blob
767 .chunks_exact(4)
768 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
769 .collect();
770 Some((msg_id, vec))
771 })
772 .collect();
773
774 Ok(map)
775 }
776
777 #[tracing::instrument(name = "memory.embed_store.get_vectors_for_messages", skip_all)]
792 pub async fn get_vectors_for_messages(
793 &self,
794 ids: &[MessageId],
795 ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
796 if ids.is_empty() {
797 return Ok(std::collections::HashMap::new());
798 }
799
800 let placeholders = zeph_db::placeholder_list(1, ids.len());
801 let query = format!(
802 "SELECT message_id, qdrant_point_id \
803 FROM embeddings_metadata \
804 WHERE message_id IN ({placeholders}) AND chunk_index = 0"
805 );
806 let mut q = zeph_db::query_as::<_, (MessageId, String)>(sqlx::AssertSqlSafe(query));
807 for &id in ids {
808 q = q.bind(id);
809 }
810 let rows: Vec<(MessageId, String)> = q.fetch_all(&self.pool).await?;
811
812 if rows.is_empty() {
813 return Ok(std::collections::HashMap::new());
814 }
815
816 let mut point_to_msg: std::collections::HashMap<String, MessageId> =
818 std::collections::HashMap::with_capacity(rows.len());
819 let point_ids: Vec<String> = rows
820 .into_iter()
821 .map(|(msg_id, point_id)| {
822 point_to_msg.insert(point_id.clone(), msg_id);
823 point_id
824 })
825 .collect();
826
827 let points = match self.ops.get_points(&self.collection, point_ids).await {
828 Ok(pts) => pts,
829 Err(crate::VectorStoreError::Unsupported(_)) => {
830 return Ok(std::collections::HashMap::new());
831 }
832 Err(e) => return Err(MemoryError::VectorStore(e)),
833 };
834
835 let result = points
836 .into_iter()
837 .filter_map(|p| {
838 let msg_id = point_to_msg.get(&p.id).copied()?;
839 Some((msg_id, p.vector))
840 })
841 .collect();
842
843 Ok(result)
844 }
845
846 #[tracing::instrument(name = "memory.embed_store.delete_by_message_ids", skip_all)]
860 pub async fn delete_by_message_ids(&self, ids: &[MessageId]) -> Result<usize, MemoryError> {
861 if ids.is_empty() {
862 return Ok(0);
863 }
864
865 let placeholders = zeph_db::placeholder_list(1, ids.len());
866 let query = format!(
867 "SELECT qdrant_point_id FROM embeddings_metadata WHERE message_id IN ({placeholders})"
868 );
869 let mut q = zeph_db::query_as::<_, (String,)>(sqlx::AssertSqlSafe(query));
870 for &id in ids {
871 q = q.bind(id);
872 }
873 let rows: Vec<(String,)> = q.fetch_all(&self.pool).await?;
874
875 let point_ids: Vec<String> = rows.into_iter().map(|(id,)| id).collect();
876 let count = point_ids.len();
877
878 if !point_ids.is_empty() {
879 self.ops.delete_by_ids(&self.collection, point_ids).await?;
880 }
881
882 Ok(count)
883 }
884
885 #[tracing::instrument(name = "memory.embed_store.has_embedding", skip_all)]
891 pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
892 let row: (i64,) = zeph_db::query_as(sql!(
893 "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
894 ))
895 .bind(message_id)
896 .fetch_one(&self.pool)
897 .await?;
898
899 Ok(row.0 > 0)
900 }
901
902 #[tracing::instrument(name = "memory.embed_store.is_epoch_current", skip_all)]
912 pub async fn is_epoch_current(
913 &self,
914 entity_name: &str,
915 qdrant_epoch: u64,
916 ) -> Result<bool, MemoryError> {
917 let row: Option<(i64,)> = zeph_db::query_as(sql!(
918 "SELECT embedding_epoch FROM graph_entities WHERE name = ? LIMIT 1"
919 ))
920 .bind(entity_name)
921 .fetch_optional(&self.pool)
922 .await?;
923
924 match row {
925 None => Ok(true), Some((db_epoch,)) => Ok(qdrant_epoch >= db_epoch.cast_unsigned()),
927 }
928 }
929}
930
931#[cfg(test)]
932mod tests {
933 use super::*;
934 use crate::db_vector_store::DbVectorStore;
935 use crate::in_memory_store::InMemoryVectorStore;
936 use crate::store::SqliteStore;
937
938 async fn setup() -> (SqliteStore, DbPool) {
939 let store = SqliteStore::new(":memory:").await.unwrap();
940 let pool = store.pool().clone();
941 (store, pool)
942 }
943
944 async fn setup_with_store() -> (EmbeddingStore, SqliteStore) {
945 let sqlite = SqliteStore::new(":memory:").await.unwrap();
946 let pool = sqlite.pool().clone();
947 let mem_store = Box::new(InMemoryVectorStore::new());
948 let embedding_store = EmbeddingStore::with_store(mem_store, pool);
949 embedding_store.ensure_collection(4).await.unwrap();
951 (embedding_store, sqlite)
952 }
953
954 #[tokio::test]
955 async fn has_embedding_returns_false_when_none() {
956 let (_store, pool) = setup().await;
957
958 let row: (i64,) = zeph_db::query_as(sql!(
959 "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
960 ))
961 .bind(999_i64)
962 .fetch_one(&pool)
963 .await
964 .unwrap();
965
966 assert_eq!(row.0, 0);
967 }
968
969 #[tokio::test]
970 async fn insert_and_query_embeddings_metadata() {
971 let (sqlite, pool) = setup().await;
972 let cid = sqlite.create_conversation().await.unwrap();
973 let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
974
975 let point_id = uuid::Uuid::new_v4().to_string();
976 zeph_db::query(sql!(
977 "INSERT INTO embeddings_metadata \
978 (message_id, chunk_index, qdrant_point_id, dimensions, model) \
979 VALUES (?, ?, ?, ?, ?)"
980 ))
981 .bind(msg_id)
982 .bind(0_i64)
983 .bind(&point_id)
984 .bind(768_i64)
985 .bind("qwen3-embedding")
986 .execute(&pool)
987 .await
988 .unwrap();
989
990 let row: (i64,) = zeph_db::query_as(sql!(
991 "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
992 ))
993 .bind(msg_id)
994 .fetch_one(&pool)
995 .await
996 .unwrap();
997 assert_eq!(row.0, 1);
998 }
999
1000 #[tokio::test]
1001 async fn embedding_store_search_empty_returns_empty() {
1002 let (store, _sqlite) = setup_with_store().await;
1003 let results = store.search(&[1.0, 0.0, 0.0, 0.0], 10, None).await.unwrap();
1004 assert!(results.is_empty());
1005 }
1006
1007 #[tokio::test]
1008 async fn embedding_store_store_and_search() {
1009 let (store, sqlite) = setup_with_store().await;
1010 let cid = sqlite.create_conversation().await.unwrap();
1011 let msg_id = sqlite
1012 .save_message(cid, "user", "test message")
1013 .await
1014 .unwrap();
1015
1016 store
1017 .store(
1018 msg_id,
1019 cid,
1020 "user",
1021 vec![1.0, 0.0, 0.0, 0.0],
1022 MessageKind::Regular,
1023 "test-model",
1024 0,
1025 )
1026 .await
1027 .unwrap();
1028
1029 let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5, None).await.unwrap();
1030 assert_eq!(results.len(), 1);
1031 assert_eq!(results[0].message_id, msg_id);
1032 assert_eq!(results[0].conversation_id, cid);
1033 assert!((results[0].score - 1.0).abs() < 0.001);
1034 }
1035
1036 #[tokio::test]
1040 async fn embedding_store_search_scoped_excludes_other_db_instance_same_conversation_id() {
1041 let shared = SqliteStore::new(":memory:").await.unwrap();
1042 let shared_pool = shared.pool().clone();
1043
1044 let sqlite_a = SqliteStore::new(":memory:").await.unwrap();
1045 let store_a = EmbeddingStore::with_store(
1046 Box::new(DbVectorStore::new(shared_pool.clone())),
1047 sqlite_a.pool().clone(),
1048 )
1049 .with_db_instance_id("db-a");
1050 store_a.ensure_collection(4).await.unwrap();
1051
1052 let sqlite_b = SqliteStore::new(":memory:").await.unwrap();
1053 let store_b = EmbeddingStore::with_store(
1054 Box::new(DbVectorStore::new(shared_pool.clone())),
1055 sqlite_b.pool().clone(),
1056 )
1057 .with_db_instance_id("db-b");
1058
1059 let cid_a = sqlite_a.create_conversation().await.unwrap();
1060 let cid_b = sqlite_b.create_conversation().await.unwrap();
1061 assert_eq!(
1062 cid_a, cid_b,
1063 "both databases must independently start at conversation_id=1"
1064 );
1065
1066 let msg_a = sqlite_a
1067 .save_message(cid_a, "user", "message in db a")
1068 .await
1069 .unwrap();
1070 let msg_b = sqlite_b
1071 .save_message(cid_b, "user", "message in db b")
1072 .await
1073 .unwrap();
1074
1075 store_a
1076 .store(
1077 msg_a,
1078 cid_a,
1079 "user",
1080 vec![1.0, 0.0, 0.0, 0.0],
1081 MessageKind::Regular,
1082 "test-model",
1083 0,
1084 )
1085 .await
1086 .unwrap();
1087 store_b
1088 .store(
1089 msg_b,
1090 cid_b,
1091 "user",
1092 vec![1.0, 0.0, 0.0, 0.0],
1093 MessageKind::Regular,
1094 "test-model",
1095 0,
1096 )
1097 .await
1098 .unwrap();
1099
1100 let filter_a = Some(SearchFilter {
1101 conversation_id: Some(cid_a),
1102 role: None,
1103 category: None,
1104 });
1105 let results_a = store_a
1106 .search(&[1.0, 0.0, 0.0, 0.0], 5, filter_a)
1107 .await
1108 .unwrap();
1109 assert_eq!(
1110 results_a.len(),
1111 1,
1112 "db-a's scoped search must not see db-b's message even though both use conversation_id=1"
1113 );
1114
1115 let filter_b = Some(SearchFilter {
1116 conversation_id: Some(cid_b),
1117 role: None,
1118 category: None,
1119 });
1120 let results_b = store_b
1121 .search(&[1.0, 0.0, 0.0, 0.0], 5, filter_b)
1122 .await
1123 .unwrap();
1124 assert_eq!(
1125 results_b.len(),
1126 1,
1127 "db-b's scoped search must not see db-a's message even though both use conversation_id=1"
1128 );
1129 }
1130
1131 #[tokio::test]
1134 async fn embedding_store_store_with_category_sets_payload_field() {
1135 let (store, sqlite) = setup_with_store().await;
1136 let cid = sqlite.create_conversation().await.unwrap();
1137 let msg_id = sqlite.save_message(cid, "user", "cat test").await.unwrap();
1138
1139 store
1140 .store_with_category(
1141 msg_id,
1142 cid,
1143 "user",
1144 vec![1.0, 0.0, 0.0, 0.0],
1145 MessageKind::Regular,
1146 "test-model",
1147 0,
1148 Some("preference"),
1149 )
1150 .await
1151 .unwrap();
1152
1153 let results = store
1154 .search_collection(COLLECTION_NAME, &[1.0, 0.0, 0.0, 0.0], 1, None)
1155 .await
1156 .unwrap();
1157 assert_eq!(results.len(), 1);
1158 assert_eq!(
1159 results[0].payload.get("category").and_then(|v| v.as_str()),
1160 Some("preference")
1161 );
1162 }
1163
1164 #[tokio::test]
1167 async fn embedding_store_store_with_category_none_omits_payload_field() {
1168 let (store, sqlite) = setup_with_store().await;
1169 let cid = sqlite.create_conversation().await.unwrap();
1170 let msg_id = sqlite.save_message(cid, "user", "no cat").await.unwrap();
1171
1172 store
1173 .store_with_category(
1174 msg_id,
1175 cid,
1176 "user",
1177 vec![0.0, 1.0, 0.0, 0.0],
1178 MessageKind::Regular,
1179 "m",
1180 0,
1181 None,
1182 )
1183 .await
1184 .unwrap();
1185
1186 let results = store
1187 .search_collection(COLLECTION_NAME, &[0.0, 1.0, 0.0, 0.0], 1, None)
1188 .await
1189 .unwrap();
1190 assert_eq!(results.len(), 1);
1191 assert!(!results[0].payload.contains_key("category"));
1192 }
1193
1194 #[tokio::test]
1197 async fn embedding_store_store_with_tool_context_sets_payload_fields() {
1198 let (store, sqlite) = setup_with_store().await;
1199 let cid = sqlite.create_conversation().await.unwrap();
1200 let msg_id = sqlite
1201 .save_message(cid, "assistant", "ran a tool")
1202 .await
1203 .unwrap();
1204
1205 store
1206 .store_with_tool_context(
1207 msg_id,
1208 cid,
1209 "assistant",
1210 vec![0.0, 0.0, 1.0, 0.0],
1211 MessageKind::Regular,
1212 "m",
1213 0,
1214 "shell",
1215 Some(0),
1216 Some("2026-07-02T00:00:00Z"),
1217 )
1218 .await
1219 .unwrap();
1220
1221 let results = store
1222 .search_collection(COLLECTION_NAME, &[0.0, 0.0, 1.0, 0.0], 1, None)
1223 .await
1224 .unwrap();
1225 assert_eq!(results.len(), 1);
1226 let payload = &results[0].payload;
1227 assert_eq!(
1228 payload.get("tool_name").and_then(|v| v.as_str()),
1229 Some("shell")
1230 );
1231 assert_eq!(
1232 payload.get("exit_code").and_then(serde_json::Value::as_i64),
1233 Some(0)
1234 );
1235 assert_eq!(
1236 payload.get("timestamp").and_then(|v| v.as_str()),
1237 Some("2026-07-02T00:00:00Z")
1238 );
1239 }
1240
1241 #[tokio::test]
1242 async fn embedding_store_has_embedding_false_for_unknown() {
1243 let (store, sqlite) = setup_with_store().await;
1244 let cid = sqlite.create_conversation().await.unwrap();
1245 let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1246 assert!(!store.has_embedding(msg_id).await.unwrap());
1247 }
1248
1249 #[tokio::test]
1250 async fn embedding_store_has_embedding_true_after_store() {
1251 let (store, sqlite) = setup_with_store().await;
1252 let cid = sqlite.create_conversation().await.unwrap();
1253 let msg_id = sqlite.save_message(cid, "user", "hello").await.unwrap();
1254
1255 store
1256 .store(
1257 msg_id,
1258 cid,
1259 "user",
1260 vec![0.0, 1.0, 0.0, 0.0],
1261 MessageKind::Regular,
1262 "test-model",
1263 0,
1264 )
1265 .await
1266 .unwrap();
1267
1268 assert!(store.has_embedding(msg_id).await.unwrap());
1269 }
1270
1271 #[tokio::test]
1272 async fn embedding_store_search_with_conversation_filter() {
1273 let (store, sqlite) = setup_with_store().await;
1274 let cid1 = sqlite.create_conversation().await.unwrap();
1275 let cid2 = sqlite.create_conversation().await.unwrap();
1276 let msg1 = sqlite.save_message(cid1, "user", "msg1").await.unwrap();
1277 let msg2 = sqlite.save_message(cid2, "user", "msg2").await.unwrap();
1278
1279 store
1280 .store(
1281 msg1,
1282 cid1,
1283 "user",
1284 vec![1.0, 0.0, 0.0, 0.0],
1285 MessageKind::Regular,
1286 "m",
1287 0,
1288 )
1289 .await
1290 .unwrap();
1291 store
1292 .store(
1293 msg2,
1294 cid2,
1295 "user",
1296 vec![1.0, 0.0, 0.0, 0.0],
1297 MessageKind::Regular,
1298 "m",
1299 0,
1300 )
1301 .await
1302 .unwrap();
1303
1304 let results = store
1305 .search(
1306 &[1.0, 0.0, 0.0, 0.0],
1307 10,
1308 Some(SearchFilter {
1309 conversation_id: Some(cid1),
1310 role: None,
1311 category: None,
1312 }),
1313 )
1314 .await
1315 .unwrap();
1316 assert_eq!(results.len(), 1);
1317 assert_eq!(results[0].conversation_id, cid1);
1318 }
1319
1320 #[tokio::test]
1321 async fn unique_constraint_on_message_chunk_and_model() {
1322 let (sqlite, pool) = setup().await;
1323 let cid = sqlite.create_conversation().await.unwrap();
1324 let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1325
1326 let point_id1 = uuid::Uuid::new_v4().to_string();
1327 zeph_db::query(sql!(
1328 "INSERT INTO embeddings_metadata \
1329 (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1330 VALUES (?, ?, ?, ?, ?)"
1331 ))
1332 .bind(msg_id)
1333 .bind(0_i64)
1334 .bind(&point_id1)
1335 .bind(768_i64)
1336 .bind("qwen3-embedding")
1337 .execute(&pool)
1338 .await
1339 .unwrap();
1340
1341 let point_id2 = uuid::Uuid::new_v4().to_string();
1343 let result = zeph_db::query(sql!(
1344 "INSERT INTO embeddings_metadata \
1345 (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1346 VALUES (?, ?, ?, ?, ?)"
1347 ))
1348 .bind(msg_id)
1349 .bind(0_i64)
1350 .bind(&point_id2)
1351 .bind(768_i64)
1352 .bind("qwen3-embedding")
1353 .execute(&pool)
1354 .await;
1355 assert!(result.is_err());
1356
1357 let point_id3 = uuid::Uuid::new_v4().to_string();
1359 zeph_db::query(sql!(
1360 "INSERT INTO embeddings_metadata \
1361 (message_id, chunk_index, qdrant_point_id, dimensions, model) \
1362 VALUES (?, ?, ?, ?, ?)"
1363 ))
1364 .bind(msg_id)
1365 .bind(1_i64)
1366 .bind(&point_id3)
1367 .bind(768_i64)
1368 .bind("qwen3-embedding")
1369 .execute(&pool)
1370 .await
1371 .unwrap();
1372 }
1373
1374 #[tokio::test]
1375 async fn get_vectors_for_messages_returns_correct_vectors() {
1376 let (store, sqlite) = setup_with_store().await;
1377 let cid = sqlite.create_conversation().await.unwrap();
1378 let msg1 = sqlite.save_message(cid, "user", "hello").await.unwrap();
1379 let msg2 = sqlite.save_message(cid, "user", "world").await.unwrap();
1380
1381 store
1382 .store(
1383 msg1,
1384 cid,
1385 "user",
1386 vec![1.0, 0.0, 0.0, 0.0],
1387 MessageKind::Regular,
1388 "m",
1389 0,
1390 )
1391 .await
1392 .unwrap();
1393 store
1394 .store(
1395 msg2,
1396 cid,
1397 "user",
1398 vec![0.0, 1.0, 0.0, 0.0],
1399 MessageKind::Regular,
1400 "m",
1401 0,
1402 )
1403 .await
1404 .unwrap();
1405
1406 let result = store.get_vectors_for_messages(&[msg1, msg2]).await.unwrap();
1407 assert_eq!(result.len(), 2);
1408 let v1 = result.get(&msg1).unwrap();
1409 let v2 = result.get(&msg2).unwrap();
1410 assert!((v1[0] - 1.0).abs() < f32::EPSILON);
1411 assert!((v2[1] - 1.0).abs() < f32::EPSILON);
1412 }
1413
1414 #[tokio::test]
1415 async fn get_vectors_for_messages_missing_id_is_dropped() {
1416 let (store, sqlite) = setup_with_store().await;
1417 let cid = sqlite.create_conversation().await.unwrap();
1418 let msg1 = sqlite.save_message(cid, "user", "present").await.unwrap();
1419 let msg_absent = MessageId(99_999);
1420
1421 store
1422 .store(
1423 msg1,
1424 cid,
1425 "user",
1426 vec![1.0, 0.0, 0.0, 0.0],
1427 MessageKind::Regular,
1428 "m",
1429 0,
1430 )
1431 .await
1432 .unwrap();
1433
1434 let result = store
1435 .get_vectors_for_messages(&[msg1, msg_absent])
1436 .await
1437 .unwrap();
1438 assert_eq!(result.len(), 1);
1439 assert!(result.contains_key(&msg1));
1440 assert!(!result.contains_key(&msg_absent));
1441 }
1442
1443 #[tokio::test]
1444 async fn get_vectors_for_messages_empty_input() {
1445 let (store, _sqlite) = setup_with_store().await;
1446 let result = store.get_vectors_for_messages(&[]).await.unwrap();
1447 assert!(result.is_empty());
1448 }
1449
1450 #[tokio::test]
1451 async fn get_vectors_for_messages_chunk_index_0_only() {
1452 let (store, sqlite) = setup_with_store().await;
1454 let cid = sqlite.create_conversation().await.unwrap();
1455 let msg = sqlite.save_message(cid, "user", "chunked").await.unwrap();
1456
1457 store
1458 .store(
1459 msg,
1460 cid,
1461 "user",
1462 vec![1.0, 0.0, 0.0, 0.0],
1463 MessageKind::Regular,
1464 "m",
1465 0,
1466 )
1467 .await
1468 .unwrap();
1469 store
1470 .store(
1471 msg,
1472 cid,
1473 "user",
1474 vec![0.0, 0.0, 1.0, 0.0],
1475 MessageKind::Regular,
1476 "m",
1477 1,
1478 )
1479 .await
1480 .unwrap();
1481
1482 let result = store.get_vectors_for_messages(&[msg]).await.unwrap();
1483 assert_eq!(result.len(), 1);
1484 let v = result.get(&msg).unwrap();
1486 assert!(
1487 (v[0] - 1.0).abs() < f32::EPSILON,
1488 "expected chunk_index=0 vector"
1489 );
1490 }
1491
1492 #[tokio::test]
1499 async fn embedding_store_delete_by_message_ids_resolves_via_metadata() {
1500 let (store, sqlite) = setup_with_store().await;
1501 let cid = sqlite.create_conversation().await.unwrap();
1502 let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
1503
1504 store
1506 .store(
1507 msg_id,
1508 cid,
1509 "user",
1510 vec![1.0, 0.0, 0.0, 0.0],
1511 MessageKind::Regular,
1512 "test-model",
1513 0,
1514 )
1515 .await
1516 .unwrap();
1517
1518 assert!(store.has_embedding(msg_id).await.unwrap());
1520
1521 let deleted = store.delete_by_message_ids(&[msg_id]).await.unwrap();
1523 assert_eq!(deleted, 1, "one point id should have been targeted");
1524
1525 let pool = sqlite.pool().clone();
1527 let row: (i64,) = zeph_db::query_as(sql!(
1528 "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
1529 ))
1530 .bind(msg_id)
1531 .fetch_one(&pool)
1532 .await
1533 .unwrap();
1534 assert_eq!(
1535 row.0, 1,
1536 "embeddings_metadata row must survive delete_by_message_ids"
1537 );
1538 }
1539
1540 #[tokio::test]
1542 async fn embedding_store_delete_by_message_ids_empty_slice_is_noop() {
1543 let (store, _sqlite) = setup_with_store().await;
1544 let deleted = store.delete_by_message_ids(&[]).await.unwrap();
1545 assert_eq!(deleted, 0);
1546 }
1547}