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 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 #[tracing::instrument(name = "memory.embed_store.store", skip_all)]
308 #[allow(clippy::too_many_arguments)] 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 #[tracing::instrument(name = "memory.embed_store.store_with_category", skip_all)]
348 #[allow(clippy::too_many_arguments)] 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 #[allow(clippy::too_many_arguments)] 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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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), 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 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 #[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 #[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 #[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 #[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 #[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 #[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 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 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 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 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 #[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
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 assert!(store.has_embedding(msg_id).await.unwrap());
1642
1643 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 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 #[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}