zeph_memory/

embedding_store.rs

// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Qdrant-backed embedding store for message vector search.
//!
//! [`EmbeddingStore`] owns a [`VectorStore`] implementation (Qdrant in production,
//! [`crate::db_vector_store::DbVectorStore`] in tests) and exposes typed `embed` /
//! `search` / `delete` operations used by [`crate::semantic::SemanticMemory`].
//!
//! Message vectors are stored in the `zeph_conversations` Qdrant collection with a
//! payload that includes `message_id`, `conversation_id`, `role`, and `category`.

pub use qdrant_client::qdrant::Filter;
use zeph_db::DbPool;
#[allow(unused_imports)]
use zeph_db::sql;

use crate::db_vector_store::DbVectorStore;
use crate::error::MemoryError;
use crate::qdrant_ops::QdrantOps;
use crate::types::{ConversationId, MessageId};
use crate::vector_store::{FieldCondition, FieldValue, VectorFilter, VectorPoint, VectorStore};

/// Distinguishes regular messages from summaries when storing embeddings.
///
/// The kind is encoded in the Qdrant payload so search filters can restrict
/// results to one category.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum MessageKind {
    /// A normal conversation message.
    Regular,
    /// A compression summary generated by the summarization subsystem.
    Summary,
}

impl MessageKind {
    #[must_use]
    pub fn is_summary(self) -> bool {
        matches!(self, Self::Summary)
    }
}

const COLLECTION_NAME: &str = "zeph_conversations";

/// Ensure a Qdrant collection exists with cosine distance vectors.
///
/// Idempotent: no-op if the collection already exists.
///
/// # Errors
///
/// Returns an error if Qdrant cannot be reached or collection creation fails.
pub async fn ensure_qdrant_collection(
    ops: &QdrantOps,
    collection: &str,
    vector_size: u64,
) -> Result<(), Box<qdrant_client::QdrantError>> {
    ops.ensure_collection(collection, vector_size).await
}

/// Typed wrapper over a [`VectorStore`] backend for conversation message embeddings.
///
/// Constructed via [`EmbeddingStore::new`] (Qdrant URL + optional API key) or
/// [`EmbeddingStore::with_store`] (custom backend for testing).
pub struct EmbeddingStore {
    ops: Box<dyn VectorStore>,
    collection: String,
    pool: DbPool,
}

impl std::fmt::Debug for EmbeddingStore {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EmbeddingStore")
            .field("collection", &self.collection)
            .finish_non_exhaustive()
    }
}

/// Optional filters applied to a vector similarity search.
#[derive(Debug)]
pub struct SearchFilter {
    /// Restrict results to a single conversation. `None` searches across all conversations.
    pub conversation_id: Option<ConversationId>,
    /// Restrict by message role (`"user"` / `"assistant"`). `None` returns all roles.
    pub role: Option<String>,
    /// Restrict by category payload field (category-aware memory, #2428).
    /// When `Some`, Qdrant search is restricted to vectors with a matching `category` payload.
    pub category: Option<String>,
}

/// A single result returned by [`EmbeddingStore::search`].
#[derive(Debug)]
pub struct SearchResult {
    /// Database row ID of the matching message.
    pub message_id: MessageId,
    /// Conversation the message belongs to.
    pub conversation_id: ConversationId,
    /// Cosine similarity score in `[0, 1]`.
    pub score: f32,
}

impl EmbeddingStore {
    /// Create a new `EmbeddingStore` connected to the given Qdrant URL with optional API key.
    ///
    /// `api_key` is forwarded to [`QdrantOps::new`]. The `pool` is used for `SQLite` metadata
    /// operations on the `embeddings_metadata` table (which must already exist via sqlx
    /// migrations).
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant client cannot be created.
    pub fn new(url: &str, api_key: Option<&str>, pool: DbPool) -> Result<Self, MemoryError> {
        let ops = QdrantOps::new(url, api_key).map_err(MemoryError::Qdrant)?;

        Ok(Self {
            ops: Box::new(ops),
            collection: COLLECTION_NAME.into(),
            pool,
        })
    }

    /// Create a new `EmbeddingStore` backed by `SQLite` for vector storage.
    ///
    /// Uses the same pool for both vector data and metadata. No external Qdrant required.
    #[must_use]
    pub fn new_sqlite(pool: DbPool) -> Self {
        let ops = DbVectorStore::new(pool.clone());
        Self {
            ops: Box::new(ops),
            collection: COLLECTION_NAME.into(),
            pool,
        }
    }

    /// Create an `EmbeddingStore` backed by an arbitrary [`VectorStore`] implementation.
    ///
    /// Intended for testing: inject a pre-configured or mock store without requiring
    /// an external Qdrant instance.
    #[must_use]
    pub fn with_store(store: Box<dyn VectorStore>, pool: DbPool) -> Self {
        Self {
            ops: store,
            collection: COLLECTION_NAME.into(),
            pool,
        }
    }

    /// Return `true` if the backing store is reachable and healthy.
    pub async fn health_check(&self) -> bool {
        self.ops.health_check().await.unwrap_or(false)
    }

    /// Ensure the collection exists in Qdrant with the given vector size.
    ///
    /// Idempotent: no-op if the collection already exists.
    ///
    /// # Errors
    ///
    /// Returns an error if Qdrant cannot be reached or collection creation fails.
    pub async fn ensure_collection(&self, vector_size: u64) -> Result<(), MemoryError> {
        self.ops
            .ensure_collection(&self.collection, vector_size)
            .await?;
        // Create keyword indexes for the fields used in filtered recall so Qdrant can satisfy
        // filter conditions in O(log n) instead of scanning all payload documents.
        self.ops
            .create_keyword_indexes(&self.collection, &["category", "conversation_id", "role"])
            .await?;
        Ok(())
    }

    /// Store a vector in Qdrant with additional tool execution metadata as payload fields.
    ///
    /// Metadata fields (`tool_name`, `exit_code`, `timestamp`) are stored as Qdrant payload
    /// alongside the standard fields. This allows filtering and scoring by tool context
    /// without corrupting the embedding vector with text prefixes.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
    pub async fn store_with_tool_context(
        &self,
        message_id: MessageId,
        conversation_id: ConversationId,
        role: &str,
        vector: Vec<f32>,
        kind: MessageKind,
        model: &str,
        chunk_index: u32,
        tool_name: &str,
        exit_code: Option<i32>,
        timestamp: Option<&str>,
    ) -> Result<String, MemoryError> {
        let point_id = uuid::Uuid::new_v4().to_string();
        let dimensions = i64::try_from(vector.len())?;

        let mut payload = std::collections::HashMap::from([
            ("message_id".to_owned(), serde_json::json!(message_id.0)),
            (
                "conversation_id".to_owned(),
                serde_json::json!(conversation_id.0),
            ),
            ("role".to_owned(), serde_json::json!(role)),
            (
                "is_summary".to_owned(),
                serde_json::json!(kind.is_summary()),
            ),
            ("tool_name".to_owned(), serde_json::json!(tool_name)),
        ]);
        if let Some(code) = exit_code {
            payload.insert("exit_code".to_owned(), serde_json::json!(code));
        }
        if let Some(ts) = timestamp {
            payload.insert("timestamp".to_owned(), serde_json::json!(ts));
        }

        let point = VectorPoint {
            id: point_id.clone(),
            vector,
            payload,
        };

        self.ops.upsert(&self.collection, vec![point]).await?;

        let chunk_index_i64 = i64::from(chunk_index);
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?) \
             ON CONFLICT(message_id, chunk_index, model) DO UPDATE SET \
             qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions"
        ))
        .bind(message_id)
        .bind(chunk_index_i64)
        .bind(&point_id)
        .bind(dimensions)
        .bind(model)
        .execute(&self.pool)
        .await?;

        Ok(point_id)
    }

    /// Store a vector in Qdrant and persist metadata to `SQLite`.
    ///
    /// `chunk_index` is 0 for single-vector messages and increases for each chunk
    /// when a long message is split into multiple embeddings.
    ///
    /// Returns the UUID of the newly created Qdrant point.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
    pub async fn store(
        &self,
        message_id: MessageId,
        conversation_id: ConversationId,
        role: &str,
        vector: Vec<f32>,
        kind: MessageKind,
        model: &str,
        chunk_index: u32,
    ) -> Result<String, MemoryError> {
        let point_id = uuid::Uuid::new_v4().to_string();
        let dimensions = i64::try_from(vector.len())?;

        let payload = std::collections::HashMap::from([
            ("message_id".to_owned(), serde_json::json!(message_id.0)),
            (
                "conversation_id".to_owned(),
                serde_json::json!(conversation_id.0),
            ),
            ("role".to_owned(), serde_json::json!(role)),
            (
                "is_summary".to_owned(),
                serde_json::json!(kind.is_summary()),
            ),
        ]);

        let point = VectorPoint {
            id: point_id.clone(),
            vector,
            payload,
        };

        self.ops.upsert(&self.collection, vec![point]).await?;

        let chunk_index_i64 = i64::from(chunk_index);
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?) \
             ON CONFLICT(message_id, chunk_index, model) DO UPDATE SET \
             qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions"
        ))
        .bind(message_id)
        .bind(chunk_index_i64)
        .bind(&point_id)
        .bind(dimensions)
        .bind(model)
        .execute(&self.pool)
        .await?;

        Ok(point_id)
    }

    /// Store a vector with an optional category tag in the Qdrant payload.
    ///
    /// Identical to [`Self::store`] but adds a `category` field to the payload when provided.
    /// Used by category-aware memory (#2428) to enable category-filtered recall.
    ///
    /// Note: when `category` is `None` no `category` field is written to the Qdrant payload.
    /// Memories stored before category-aware recall was enabled therefore won't match a
    /// category filter — this is intentional (no silent false-positives), but a backfill
    /// pass is needed if retrospective categorization is desired.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant upsert or `SQLite` insert fails.
    #[allow(clippy::too_many_arguments)] // function with many required inputs; a *Params struct would be more verbose without simplifying the call site
    pub async fn store_with_category(
        &self,
        message_id: MessageId,
        conversation_id: ConversationId,
        role: &str,
        vector: Vec<f32>,
        kind: MessageKind,
        model: &str,
        chunk_index: u32,
        category: Option<&str>,
    ) -> Result<String, MemoryError> {
        let point_id = uuid::Uuid::new_v4().to_string();
        let dimensions = i64::try_from(vector.len())?;

        let mut payload = std::collections::HashMap::from([
            ("message_id".to_owned(), serde_json::json!(message_id.0)),
            (
                "conversation_id".to_owned(),
                serde_json::json!(conversation_id.0),
            ),
            ("role".to_owned(), serde_json::json!(role)),
            (
                "is_summary".to_owned(),
                serde_json::json!(kind.is_summary()),
            ),
        ]);
        if let Some(cat) = category {
            payload.insert("category".to_owned(), serde_json::json!(cat));
        }

        let point = VectorPoint {
            id: point_id.clone(),
            vector,
            payload,
        };

        self.ops.upsert(&self.collection, vec![point]).await?;

        let chunk_index_i64 = i64::from(chunk_index);
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?) \
             ON CONFLICT(message_id, chunk_index, model) DO UPDATE SET \
             qdrant_point_id = excluded.qdrant_point_id, dimensions = excluded.dimensions"
        ))
        .bind(message_id)
        .bind(chunk_index_i64)
        .bind(&point_id)
        .bind(dimensions)
        .bind(model)
        .execute(&self.pool)
        .await?;

        Ok(point_id)
    }

    /// Search for similar vectors in Qdrant, returning up to `limit` results.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant search fails.
    pub async fn search(
        &self,
        query_vector: &[f32],
        limit: usize,
        filter: Option<SearchFilter>,
    ) -> Result<Vec<SearchResult>, MemoryError> {
        let limit_u64 = u64::try_from(limit)?;

        let vector_filter = filter.as_ref().and_then(|f| {
            let mut must = Vec::new();
            if let Some(cid) = f.conversation_id {
                must.push(FieldCondition {
                    field: "conversation_id".into(),
                    value: FieldValue::Integer(cid.0),
                });
            }
            if let Some(ref role) = f.role {
                must.push(FieldCondition {
                    field: "role".into(),
                    value: FieldValue::Text(role.clone()),
                });
            }
            if let Some(ref category) = f.category {
                must.push(FieldCondition {
                    field: "category".into(),
                    value: FieldValue::Text(category.clone()),
                });
            }
            if must.is_empty() {
                None
            } else {
                Some(VectorFilter {
                    must,
                    must_not: vec![],
                })
            }
        });

        let results = self
            .ops
            .search(
                &self.collection,
                query_vector.to_vec(),
                limit_u64,
                vector_filter,
            )
            .await?;

        // Deduplicate by message_id, keeping the chunk with the highest score.
        // A single message may produce multiple Qdrant points (one per chunk).
        let mut best: std::collections::HashMap<MessageId, SearchResult> =
            std::collections::HashMap::new();
        for point in results {
            let Some(message_id) = point
                .payload
                .get("message_id")
                .and_then(serde_json::Value::as_i64)
            else {
                continue;
            };
            let Some(conversation_id) = point
                .payload
                .get("conversation_id")
                .and_then(serde_json::Value::as_i64)
            else {
                continue;
            };
            let message_id = MessageId(message_id);
            let entry = best.entry(message_id).or_insert(SearchResult {
                message_id,
                conversation_id: ConversationId(conversation_id),
                score: f32::NEG_INFINITY,
            });
            if point.score > entry.score {
                entry.score = point.score;
            }
        }

        let mut search_results: Vec<SearchResult> = best.into_values().collect();
        search_results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        search_results.truncate(limit);

        Ok(search_results)
    }

    /// Check whether a named collection exists in the vector store.
    ///
    /// # Errors
    ///
    /// Returns an error if the store backend cannot be reached.
    pub async fn collection_exists(&self, name: &str) -> Result<bool, MemoryError> {
        self.ops.collection_exists(name).await.map_err(Into::into)
    }

    /// Ensure a named collection exists in Qdrant with the given vector size.
    ///
    /// # Errors
    ///
    /// Returns an error if Qdrant cannot be reached or collection creation fails.
    pub async fn ensure_named_collection(
        &self,
        name: &str,
        vector_size: u64,
    ) -> Result<(), MemoryError> {
        self.ops.ensure_collection(name, vector_size).await?;
        Ok(())
    }

    /// Store a vector in a named Qdrant collection with arbitrary payload.
    ///
    /// Returns the UUID of the newly created point.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant upsert fails.
    pub async fn store_to_collection(
        &self,
        collection: &str,
        payload: serde_json::Value,
        vector: Vec<f32>,
    ) -> Result<String, MemoryError> {
        let point_id = uuid::Uuid::new_v4().to_string();
        let payload_map: std::collections::HashMap<String, serde_json::Value> =
            serde_json::from_value(payload)?;
        let point = VectorPoint {
            id: point_id.clone(),
            vector,
            payload: payload_map,
        };
        self.ops.upsert(collection, vec![point]).await?;
        Ok(point_id)
    }

    /// Upsert a vector into a named collection, reusing an existing point ID.
    ///
    /// Use this when updating an existing entity to avoid orphaned Qdrant points.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant upsert fails.
    pub async fn upsert_to_collection(
        &self,
        collection: &str,
        point_id: &str,
        payload: serde_json::Value,
        vector: Vec<f32>,
    ) -> Result<(), MemoryError> {
        let payload_map: std::collections::HashMap<String, serde_json::Value> =
            serde_json::from_value(payload)?;
        let point = VectorPoint {
            id: point_id.to_owned(),
            vector,
            payload: payload_map,
        };
        self.ops.upsert(collection, vec![point]).await?;
        Ok(())
    }

    /// Search a named Qdrant collection, returning scored points with payloads.
    ///
    /// # Errors
    ///
    /// Returns an error if the Qdrant search fails.
    pub async fn search_collection(
        &self,
        collection: &str,
        query_vector: &[f32],
        limit: usize,
        filter: Option<VectorFilter>,
    ) -> Result<Vec<crate::ScoredVectorPoint>, MemoryError> {
        let limit_u64 = u64::try_from(limit)?;
        let results = self
            .ops
            .search(collection, query_vector.to_vec(), limit_u64, filter)
            .await?;
        Ok(results)
    }

    /// Enumerate `(point_id, entity_id)` pairs for all points in `collection` that carry
    /// an `entity_id_str` payload field.
    ///
    /// `entity_id_str` is a string mirror of the i64 `entity_id` written alongside the numeric
    /// field at embedding time. The scroll API only surfaces string-typed payload values, so a
    /// parallel string field is necessary for enumeration. Points missing `entity_id_str`
    /// (written before this field was added) are silently skipped — they will gain the field on
    /// the next `merge_entity` or `store_entity_embedding` call.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying scroll operation fails.
    pub async fn scroll_all_entity_ids(
        &self,
        collection: &str,
    ) -> Result<Vec<(String, i64)>, MemoryError> {
        let rows = self
            .ops
            .scroll_all_with_point_ids(collection, "entity_id_str")
            .await?;
        let mut out = Vec::with_capacity(rows.len());
        for (point_id, fields) in rows {
            let Some(s) = fields.get("entity_id_str") else {
                continue;
            };
            if let Ok(id) = s.parse::<i64>() {
                out.push((point_id, id));
            } else {
                tracing::debug!(point_id, value = %s, "entity_id_str unparseable, skipping");
            }
        }
        Ok(out)
    }

    /// Delete a set of points from a named collection by their Qdrant point IDs.
    ///
    /// This is a thin wrapper over [`VectorStore::delete_by_ids`] for use by
    /// the stale-embedding cleanup path in `community.rs`.
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying delete operation fails.
    pub async fn delete_from_collection(
        &self,
        collection: &str,
        ids: Vec<String>,
    ) -> Result<(), MemoryError> {
        if ids.is_empty() {
            return Ok(());
        }
        self.ops.delete_by_ids(collection, ids).await?;
        Ok(())
    }

    /// Retrieve raw vectors for the given Qdrant point IDs from `collection`.
    ///
    /// Returns a map of `point_id → embedding`. Missing ids are silently dropped.
    /// Returns an empty map when the backend does not support vector retrieval
    /// (e.g. `DbVectorStore` / `InMemoryVectorStore` without an override).
    ///
    /// # Errors
    ///
    /// Returns an error if the underlying store returns a non-`Unsupported` error.
    pub async fn get_vectors_from_collection(
        &self,
        collection: &str,
        point_ids: &[String],
    ) -> Result<std::collections::HashMap<String, Vec<f32>>, MemoryError> {
        if point_ids.is_empty() {
            return Ok(std::collections::HashMap::new());
        }
        match self.ops.get_points(collection, point_ids.to_vec()).await {
            Ok(points) => Ok(points.into_iter().map(|p| (p.id, p.vector)).collect()),
            Err(crate::VectorStoreError::Unsupported(_)) => Ok(std::collections::HashMap::new()),
            Err(e) => Err(MemoryError::VectorStore(e)),
        }
    }

    /// Fetch raw vectors for the given message IDs from the `SQLite` vector store.
    ///
    /// Returns an empty map when using Qdrant backend (vectors not locally stored).
    ///
    /// # Errors
    ///
    /// Returns an error if the `SQLite` query fails.
    pub async fn get_vectors(
        &self,
        ids: &[MessageId],
    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
        if ids.is_empty() {
            return Ok(std::collections::HashMap::new());
        }

        let placeholders = zeph_db::placeholder_list(1, ids.len());
        let query = format!(
            "SELECT em.message_id, vp.vector \
             FROM embeddings_metadata em \
             JOIN vector_points vp ON vp.id = em.qdrant_point_id \
             WHERE em.message_id IN ({placeholders}) AND em.chunk_index = 0"
        );
        let mut q = zeph_db::query_as::<_, (MessageId, Vec<u8>)>(&query);
        for &id in ids {
            q = q.bind(id);
        }

        let rows = q.fetch_all(&self.pool).await?;

        let map = rows
            .into_iter()
            .filter_map(|(msg_id, blob)| {
                if blob.len() % 4 != 0 {
                    return None;
                }
                let vec: Vec<f32> = blob
                    .chunks_exact(4)
                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                    .collect();
                Some((msg_id, vec))
            })
            .collect();

        Ok(map)
    }

    /// Fetch embeddings for the given message IDs from the configured vector store.
    ///
    /// Resolves `message_id → qdrant_point_id` via `embeddings_metadata` (filtering to
    /// `chunk_index = 0` so each message yields at most one vector), then retrieves the
    /// vectors from the underlying [`VectorStore`].
    ///
    /// Returns a map from [`MessageId`] to embedding vector. Messages without an
    /// `embeddings_metadata` row, or whose vector cannot be retrieved, are silently dropped.
    /// When the backend returns [`crate::VectorStoreError::Unsupported`], an empty map is
    /// returned without error (matches [`Self::get_vectors_from_collection`] semantics).
    ///
    /// # Errors
    ///
    /// Returns an error if the `SQLite` metadata query or vector store retrieval fails.
    pub async fn get_vectors_for_messages(
        &self,
        ids: &[MessageId],
    ) -> Result<std::collections::HashMap<MessageId, Vec<f32>>, MemoryError> {
        if ids.is_empty() {
            return Ok(std::collections::HashMap::new());
        }

        let placeholders = zeph_db::placeholder_list(1, ids.len());
        let query = format!(
            "SELECT message_id, qdrant_point_id \
             FROM embeddings_metadata \
             WHERE message_id IN ({placeholders}) AND chunk_index = 0"
        );
        let mut q = zeph_db::query_as::<_, (MessageId, String)>(&query);
        for &id in ids {
            q = q.bind(id);
        }
        let rows: Vec<(MessageId, String)> = q.fetch_all(&self.pool).await?;

        if rows.is_empty() {
            return Ok(std::collections::HashMap::new());
        }

        // Build reverse map: point_id → message_id for result translation.
        let mut point_to_msg: std::collections::HashMap<String, MessageId> =
            std::collections::HashMap::with_capacity(rows.len());
        let point_ids: Vec<String> = rows
            .into_iter()
            .map(|(msg_id, point_id)| {
                point_to_msg.insert(point_id.clone(), msg_id);
                point_id
            })
            .collect();

        let points = match self.ops.get_points(&self.collection, point_ids).await {
            Ok(pts) => pts,
            Err(crate::VectorStoreError::Unsupported(_)) => {
                return Ok(std::collections::HashMap::new());
            }
            Err(e) => return Err(MemoryError::VectorStore(e)),
        };

        let result = points
            .into_iter()
            .filter_map(|p| {
                let msg_id = point_to_msg.get(&p.id).copied()?;
                Some((msg_id, p.vector))
            })
            .collect();

        Ok(result)
    }

    /// Delete all Qdrant vectors associated with the given message IDs.
    ///
    /// Resolves `message_id → qdrant_point_id` via the `embeddings_metadata` table,
    /// then calls the underlying vector store's `delete_by_ids`. The
    /// `embeddings_metadata` rows are **not** removed here — the `SQLite` CASCADE on
    /// `messages` handles that when the rows are hard-deleted later.
    ///
    /// Returns the number of Qdrant point IDs targeted for deletion (may be less than
    /// `ids.len()` when some messages have no embeddings).
    ///
    /// # Errors
    ///
    /// Returns [`MemoryError`] if the `SQLite` query or the vector store delete fails.
    pub async fn delete_by_message_ids(&self, ids: &[MessageId]) -> Result<usize, MemoryError> {
        if ids.is_empty() {
            return Ok(0);
        }

        let placeholders = zeph_db::placeholder_list(1, ids.len());
        let query = format!(
            "SELECT qdrant_point_id FROM embeddings_metadata WHERE message_id IN ({placeholders})"
        );
        let mut q = zeph_db::query_as::<_, (String,)>(&query);
        for &id in ids {
            q = q.bind(id);
        }
        let rows: Vec<(String,)> = q.fetch_all(&self.pool).await?;

        let point_ids: Vec<String> = rows.into_iter().map(|(id,)| id).collect();
        let count = point_ids.len();

        if !point_ids.is_empty() {
            self.ops.delete_by_ids(&self.collection, point_ids).await?;
        }

        Ok(count)
    }

    /// Check whether an embedding already exists for the given message ID.
    ///
    /// # Errors
    ///
    /// Returns an error if the `SQLite` query fails.
    pub async fn has_embedding(&self, message_id: MessageId) -> Result<bool, MemoryError> {
        let row: (i64,) = zeph_db::query_as(sql!(
            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
        ))
        .bind(message_id)
        .fetch_one(&self.pool)
        .await?;

        Ok(row.0 > 0)
    }

    /// Check whether a Qdrant embedding for `entity_name` is current by comparing the
    /// Qdrant-side epoch against the epoch stored in `graph_entities`.
    ///
    /// Returns `true` if the Qdrant embedding is up-to-date or if the entity no longer
    /// exists in `SQLite` (embedding should be cleaned up separately).
    ///
    /// # Errors
    ///
    /// Returns an error if the `SQLite` query fails.
    pub async fn is_epoch_current(
        &self,
        entity_name: &str,
        qdrant_epoch: u64,
    ) -> Result<bool, MemoryError> {
        let row: Option<(i64,)> = zeph_db::query_as(sql!(
            "SELECT embedding_epoch FROM graph_entities WHERE name = ? LIMIT 1"
        ))
        .bind(entity_name)
        .fetch_optional(&self.pool)
        .await?;

        match row {
            None => Ok(true), // entity deleted; Qdrant point is orphaned, not stale per epoch
            Some((db_epoch,)) => Ok(qdrant_epoch >= db_epoch.cast_unsigned()),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::in_memory_store::InMemoryVectorStore;
    use crate::store::SqliteStore;

    async fn setup() -> (SqliteStore, DbPool) {
        let store = SqliteStore::new(":memory:").await.unwrap();
        let pool = store.pool().clone();
        (store, pool)
    }

    async fn setup_with_store() -> (EmbeddingStore, SqliteStore) {
        let sqlite = SqliteStore::new(":memory:").await.unwrap();
        let pool = sqlite.pool().clone();
        let mem_store = Box::new(InMemoryVectorStore::new());
        let embedding_store = EmbeddingStore::with_store(mem_store, pool);
        // Create collection first
        embedding_store.ensure_collection(4).await.unwrap();
        (embedding_store, sqlite)
    }

    #[tokio::test]
    async fn has_embedding_returns_false_when_none() {
        let (_store, pool) = setup().await;

        let row: (i64,) = zeph_db::query_as(sql!(
            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
        ))
        .bind(999_i64)
        .fetch_one(&pool)
        .await
        .unwrap();

        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn insert_and_query_embeddings_metadata() {
        let (sqlite, pool) = setup().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();

        let point_id = uuid::Uuid::new_v4().to_string();
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?)"
        ))
        .bind(msg_id)
        .bind(0_i64)
        .bind(&point_id)
        .bind(768_i64)
        .bind("qwen3-embedding")
        .execute(&pool)
        .await
        .unwrap();

        let row: (i64,) = zeph_db::query_as(sql!(
            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
        ))
        .bind(msg_id)
        .fetch_one(&pool)
        .await
        .unwrap();
        assert_eq!(row.0, 1);
    }

    #[tokio::test]
    async fn embedding_store_search_empty_returns_empty() {
        let (store, _sqlite) = setup_with_store().await;
        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 10, None).await.unwrap();
        assert!(results.is_empty());
    }

    #[tokio::test]
    async fn embedding_store_store_and_search() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite
            .save_message(cid, "user", "test message")
            .await
            .unwrap();

        store
            .store(
                msg_id,
                cid,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "test-model",
                0,
            )
            .await
            .unwrap();

        let results = store.search(&[1.0, 0.0, 0.0, 0.0], 5, None).await.unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].message_id, msg_id);
        assert_eq!(results[0].conversation_id, cid);
        assert!((results[0].score - 1.0).abs() < 0.001);
    }

    #[tokio::test]
    async fn embedding_store_has_embedding_false_for_unknown() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();
        assert!(!store.has_embedding(msg_id).await.unwrap());
    }

    #[tokio::test]
    async fn embedding_store_has_embedding_true_after_store() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite.save_message(cid, "user", "hello").await.unwrap();

        store
            .store(
                msg_id,
                cid,
                "user",
                vec![0.0, 1.0, 0.0, 0.0],
                MessageKind::Regular,
                "test-model",
                0,
            )
            .await
            .unwrap();

        assert!(store.has_embedding(msg_id).await.unwrap());
    }

    #[tokio::test]
    async fn embedding_store_search_with_conversation_filter() {
        let (store, sqlite) = setup_with_store().await;
        let cid1 = sqlite.create_conversation().await.unwrap();
        let cid2 = sqlite.create_conversation().await.unwrap();
        let msg1 = sqlite.save_message(cid1, "user", "msg1").await.unwrap();
        let msg2 = sqlite.save_message(cid2, "user", "msg2").await.unwrap();

        store
            .store(
                msg1,
                cid1,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();
        store
            .store(
                msg2,
                cid2,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();

        let results = store
            .search(
                &[1.0, 0.0, 0.0, 0.0],
                10,
                Some(SearchFilter {
                    conversation_id: Some(cid1),
                    role: None,
                    category: None,
                }),
            )
            .await
            .unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].conversation_id, cid1);
    }

    #[tokio::test]
    async fn unique_constraint_on_message_chunk_and_model() {
        let (sqlite, pool) = setup().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();

        let point_id1 = uuid::Uuid::new_v4().to_string();
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?)"
        ))
        .bind(msg_id)
        .bind(0_i64)
        .bind(&point_id1)
        .bind(768_i64)
        .bind("qwen3-embedding")
        .execute(&pool)
        .await
        .unwrap();

        // Same (message_id, chunk_index, model) — must fail.
        let point_id2 = uuid::Uuid::new_v4().to_string();
        let result = zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?)"
        ))
        .bind(msg_id)
        .bind(0_i64)
        .bind(&point_id2)
        .bind(768_i64)
        .bind("qwen3-embedding")
        .execute(&pool)
        .await;
        assert!(result.is_err());

        // Different chunk_index — must succeed.
        let point_id3 = uuid::Uuid::new_v4().to_string();
        zeph_db::query(sql!(
            "INSERT INTO embeddings_metadata \
             (message_id, chunk_index, qdrant_point_id, dimensions, model) \
             VALUES (?, ?, ?, ?, ?)"
        ))
        .bind(msg_id)
        .bind(1_i64)
        .bind(&point_id3)
        .bind(768_i64)
        .bind("qwen3-embedding")
        .execute(&pool)
        .await
        .unwrap();
    }

    #[tokio::test]
    async fn get_vectors_for_messages_returns_correct_vectors() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg1 = sqlite.save_message(cid, "user", "hello").await.unwrap();
        let msg2 = sqlite.save_message(cid, "user", "world").await.unwrap();

        store
            .store(
                msg1,
                cid,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();
        store
            .store(
                msg2,
                cid,
                "user",
                vec![0.0, 1.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();

        let result = store.get_vectors_for_messages(&[msg1, msg2]).await.unwrap();
        assert_eq!(result.len(), 2);
        let v1 = result.get(&msg1).unwrap();
        let v2 = result.get(&msg2).unwrap();
        assert!((v1[0] - 1.0).abs() < f32::EPSILON);
        assert!((v2[1] - 1.0).abs() < f32::EPSILON);
    }

    #[tokio::test]
    async fn get_vectors_for_messages_missing_id_is_dropped() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg1 = sqlite.save_message(cid, "user", "present").await.unwrap();
        let msg_absent = MessageId(99_999);

        store
            .store(
                msg1,
                cid,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();

        let result = store
            .get_vectors_for_messages(&[msg1, msg_absent])
            .await
            .unwrap();
        assert_eq!(result.len(), 1);
        assert!(result.contains_key(&msg1));
        assert!(!result.contains_key(&msg_absent));
    }

    #[tokio::test]
    async fn get_vectors_for_messages_empty_input() {
        let (store, _sqlite) = setup_with_store().await;
        let result = store.get_vectors_for_messages(&[]).await.unwrap();
        assert!(result.is_empty());
    }

    #[tokio::test]
    async fn get_vectors_for_messages_chunk_index_0_only() {
        // Store chunk_index=0 and chunk_index=1; only chunk_index=0 should be returned.
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg = sqlite.save_message(cid, "user", "chunked").await.unwrap();

        store
            .store(
                msg,
                cid,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "m",
                0,
            )
            .await
            .unwrap();
        store
            .store(
                msg,
                cid,
                "user",
                vec![0.0, 0.0, 1.0, 0.0],
                MessageKind::Regular,
                "m",
                1,
            )
            .await
            .unwrap();

        let result = store.get_vectors_for_messages(&[msg]).await.unwrap();
        assert_eq!(result.len(), 1);
        // Must be the chunk_index=0 vector
        let v = result.get(&msg).unwrap();
        assert!(
            (v[0] - 1.0).abs() < f32::EPSILON,
            "expected chunk_index=0 vector"
        );
    }

    /// `delete_by_message_ids` resolves `message_id → qdrant_point_id` via
    /// `embeddings_metadata` and deletes the matching vectors.
    ///
    /// Verifies: (a) the correct point id is targeted, (b) `embeddings_metadata`
    /// rows are NOT removed (CASCADE handles that on hard-delete later), and (c) the
    /// method returns the number of point IDs found.
    #[tokio::test]
    async fn embedding_store_delete_by_message_ids_resolves_via_metadata() {
        let (store, sqlite) = setup_with_store().await;
        let cid = sqlite.create_conversation().await.unwrap();
        let msg_id = sqlite.save_message(cid, "user", "test").await.unwrap();

        // Store a vector so embeddings_metadata gets a row.
        store
            .store(
                msg_id,
                cid,
                "user",
                vec![1.0, 0.0, 0.0, 0.0],
                MessageKind::Regular,
                "test-model",
                0,
            )
            .await
            .unwrap();

        // Confirm the metadata row exists before deletion.
        assert!(store.has_embedding(msg_id).await.unwrap());

        // Delete by message id — must succeed and return 1 (one point id resolved).
        let deleted = store.delete_by_message_ids(&[msg_id]).await.unwrap();
        assert_eq!(deleted, 1, "one point id should have been targeted");

        // embeddings_metadata rows must still be present (CASCADE removes them later).
        let pool = sqlite.pool().clone();
        let row: (i64,) = zeph_db::query_as(sql!(
            "SELECT COUNT(*) FROM embeddings_metadata WHERE message_id = ?"
        ))
        .bind(msg_id)
        .fetch_one(&pool)
        .await
        .unwrap();
        assert_eq!(
            row.0, 1,
            "embeddings_metadata row must survive delete_by_message_ids"
        );
    }

    /// `delete_by_message_ids` is a no-op when the slice is empty.
    #[tokio::test]
    async fn embedding_store_delete_by_message_ids_empty_slice_is_noop() {
        let (store, _sqlite) = setup_with_store().await;
        let deleted = store.delete_by_message_ids(&[]).await.unwrap();
        assert_eq!(deleted, 0);
    }
}