kiromi-ai-memory 0.2.2

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Plan 18 phase C — `VectorIndex` plugin trait.
//!
//! The default implementation [`super::sqlite_vec_index::SqliteVecIndex`]
//! talks to the `vec0` virtual table created idempotently on every
//! `Memory::open()` via
//! [`crate::metadata::MetadataStore::create_indices_if_missing`]. Future swaps
//! (Vectorlite, an external HNSW service, etc.) implement this trait and plug
//! into the engine via the builder.

use async_trait::async_trait;

use crate::attribute::AttributeValue;
use crate::error::Result;
use crate::memory::{MemoryId, MemoryKind};
use crate::partition::PartitionPath;
use crate::summarizer::SummaryStyle;
use crate::summary::SummaryId;

/// Where a `knn` call should look. The vec0-backed default impl turns each
/// variant into a `partition_path` filter on the virtual table.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub enum VectorScope {
    /// Every live row in the tenant.
    Tenant,
    /// Exactly one partition (no recursion).
    Partition(PartitionPath),
    /// Every partition under the supplied prefix (recursive).
    /// The prefix is matched as `<prefix>` or `<prefix>/%`.
    PartitionPrefix(String),
}

/// Per-attribute filter applied during KNN. The default impl turns this into
/// a JOIN against `memory_attribute`.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct VectorFilter {
    /// Attribute key.
    pub key: String,
    /// Attribute value the row must match exactly.
    pub value: AttributeValue,
}

/// Distance metric the index uses. `vec0` defaults to cosine for
/// `vec_distance_cosine`; other backends may differ.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DistanceMetric {
    /// 1 - cos(a, b). Smaller = more similar.
    CosineDistance,
    /// L2 (Euclidean). Smaller = more similar.
    L2Distance,
}

/// Self-declared capabilities for a `VectorIndex` impl.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct VectorIndexCapabilities {
    /// Whether the impl supports filtered KNN (attribute predicates).
    pub knn_filtered: bool,
    /// Maximum supported vector dimension.
    pub max_dimensions: usize,
    /// What metric the impl returns scores in.
    pub distance_metric: DistanceMetric,
}

impl Default for VectorIndexCapabilities {
    fn default() -> Self {
        Self {
            knn_filtered: false,
            max_dimensions: 4096,
            distance_metric: DistanceMetric::CosineDistance,
        }
    }
}

/// Plugin trait for the vector half of the engine's index surface.
///
/// All ops are async; the default impl shares the `SqlitePool` with
/// `SqliteMetadata` so an `upsert_memory` call from inside an `append`
/// transaction stays in the same sqlite connection. (See Plan 18's
/// "single-transaction append" property.)
///
/// # Atomicity
///
/// The `upsert_*` / `delete_*` methods on this trait are **not, by
/// themselves, atomic with the catalog row** they conceptually belong to.
/// A direct call (e.g. from a custom reindex tool) can leave the catalog
/// row and the vector index momentarily out of sync if the process crashes
/// between calls.
///
/// For the engine's authoritative writers — `MetadataStore::append_memory`,
/// `MetadataStore::insert_summary`, `MetadataStore::regenerate_memory_embedding`
/// — the index inserts are folded into the same SQL transaction as the
/// catalog writes (when the default [`super::SqliteVecIndex`] impl is used),
/// so the per-mutation single-transaction property holds.
///
/// Third-party impls that want the same property must run their own
/// upsert work inside whatever transactional context the engine offers
/// at the call site.
#[async_trait]
pub trait VectorIndex: Send + Sync + std::fmt::Debug + 'static {
    /// Insert or replace a memory's vector row.
    ///
    /// **Atomicity:** see the trait-level [Atomicity] note —
    /// invoking this method directly does **not** guarantee atomicity
    /// with the catalog row. Use the [`crate::metadata::MetadataStore`]
    /// helpers (`append_memory`, `regenerate_memory_embedding`) for
    /// transactional writes.
    ///
    /// [Atomicity]: VectorIndex#atomicity
    async fn upsert_memory(
        &self,
        id: &MemoryId,
        partition_path: &PartitionPath,
        kind: Option<&MemoryKind>,
        embedding: &[f32],
    ) -> Result<()>;

    /// Insert or replace a summary's vector row.
    async fn upsert_summary(
        &self,
        id: &SummaryId,
        parent_path: &str,
        style: &SummaryStyle,
        embedding: &[f32],
    ) -> Result<()>;

    /// Delete a memory's vector row. Idempotent.
    async fn delete_memory(&self, id: &MemoryId) -> Result<()>;

    /// Delete a summary's vector row. Idempotent.
    async fn delete_summary(&self, id: &SummaryId) -> Result<()>;

    /// k-nearest memories under `scope`, optionally filtered by an attribute.
    /// Returns `(memory_id, distance)` pairs ordered ascending.
    async fn knn_memory(
        &self,
        query: &[f32],
        k: u32,
        scope: VectorScope,
        filter: Option<&VectorFilter>,
    ) -> Result<Vec<(MemoryId, f32)>>;

    /// k-nearest summaries whose `parent_path` is under the supplied prefix.
    async fn knn_summary(
        &self,
        query: &[f32],
        k: u32,
        parent_path_prefix: &str,
    ) -> Result<Vec<(SummaryId, f32)>>;

    /// Stable identifier for tracing (`"sqlite-vec:vec0"`, …).
    fn id(&self) -> &str;

    /// Self-declared capabilities.
    fn capabilities(&self) -> VectorIndexCapabilities;
}