IdeaSDK/crates/infrastructure/src/store/vector.rs

//! Étage-2 semantic recall (LOT C, §14.5.3): [`VectorMemoryRecall`] and the
//! [`AdaptiveMemoryRecall`] switch.
//!
//! [`VectorMemoryRecall`] composes an [`Embedder`] + a [`MemoryStore`] + a small
//! **derived vector store** under `.ideai/memory/.index/`, ranks the memory index
//! entries by cosine similarity to the query, and truncates to the token budget —
//! the **same** budget/emptiness semantics as [`super::NaiveMemoryRecall`]
//! (Liskov), only the relevance strategy differs.
//!
//! [`AdaptiveMemoryRecall`] composes both and routes between them by an
//! **objective, pure, I/O-free** decision ([`should_use_vector`]): tiny memories
//! or a `none` strategy go to the naïve recall; otherwise to the vector recall,
//! with an automatic **fallback to naïve** whenever the embedder is unavailable.
//! No path of this module ever fails hard on a missing/unavailable embedder.
//!
//! ## Derived vector store format
//!
//! `.ideai/memory/.index/vectors.json` holds, per slug, the embedding of that
//! note's index line (title + hook), tagged with the producing embedder id and
//! its dimension:
//!
//! ```json
//! {
//!   "version": 1,
//!   "embedderId": "hash-embedder",
//!   "dimension": 64,
//!   "vectors": { "my-note": [0.1, 0.0, ...] }
//! }
//! ```
//!
//! It is **derived data, fully reconstructible** from the `.md` source of truth,
//! and is **gitignored** (`.ideai/memory/.index/`). A stale or absent file is not
//! an error: missing vectors are recomputed on demand and the file is refreshed
//! best-effort. A change of embedder id/dimension invalidates the whole file.

use std::collections::HashMap;
use std::sync::Arc;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use domain::memory::MemoryIndexEntry;
use domain::ports::{
    Embedder, EmbedderError, FileSystem, FsError, MemoryError, MemoryQuery, MemoryRecall,
    MemoryStore, RemotePath,
};
use domain::profile::EmbedderStrategy;
use domain::project::ProjectPath;

use super::memory::{index_token_size, truncate_to_budget};

/// `.ideai/` directory name (mirrors [`super::memory`]).
const IDEAI_DIR: &str = ".ideai";
/// Memory sub-dir.
const MEMORY_DIR: &str = "memory";
/// Derived vector-store sub-dir (gitignored).
const INDEX_DIR: &str = ".index";
/// Derived vector-store file.
const VECTORS_FILE: &str = "vectors.json";
/// Schema version of the derived vector store.
const VECTORS_VERSION: u32 = 1;

/// On-disk shape of `.ideai/memory/.index/vectors.json`.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct VectorDoc {
    /// Schema version.
    version: u32,
    /// Id of the embedder that produced these vectors (a mismatch invalidates).
    embedder_id: String,
    /// Vector length (a mismatch invalidates).
    dimension: usize,
    /// Per-slug embedding of the note's index line.
    vectors: HashMap<String, Vec<f32>>,
}

impl VectorDoc {
    fn new(embedder_id: String, dimension: usize) -> Self {
        Self {
            version: VECTORS_VERSION,
            embedder_id,
            dimension,
            vectors: HashMap::new(),
        }
    }

    /// Whether this doc was produced by the given embedder (id + dimension).
    fn matches(&self, embedder_id: &str, dimension: usize) -> bool {
        self.version == VECTORS_VERSION
            && self.embedder_id == embedder_id
            && self.dimension == dimension
    }
}

/// Semantic [`MemoryRecall`] (étage 2): ranks the memory index by cosine
/// similarity of each note's index line to the query text, truncated to the token
/// budget (identical budget semantics to [`super::NaiveMemoryRecall`]).
///
/// Composes an [`Embedder`], a [`MemoryStore`] (the index source of truth), and a
/// [`FileSystem`] for the derived vector cache. **Best-effort**: any
/// [`EmbedderError`] surfaces as [`EmbedderError`] to the caller
/// ([`AdaptiveMemoryRecall`] turns it into a naïve fallback); a missing/empty
/// memory yields an empty list, never an error.
#[derive(Clone)]
pub struct VectorMemoryRecall {
    embedder: Arc<dyn Embedder>,
    store: Arc<dyn MemoryStore>,
    fs: Arc<dyn FileSystem>,
}

impl VectorMemoryRecall {
    /// Builds the vector recall from its composed ports.
    #[must_use]
    pub fn new(
        embedder: Arc<dyn Embedder>,
        store: Arc<dyn MemoryStore>,
        fs: Arc<dyn FileSystem>,
    ) -> Self {
        Self {
            embedder,
            store,
            fs,
        }
    }

    /// `<root>/.ideai/memory/.index`.
    fn index_dir(root: &ProjectPath) -> String {
        let base = root.as_str().trim_end_matches(['/', '\\']);
        format!("{base}/{IDEAI_DIR}/{MEMORY_DIR}/{INDEX_DIR}")
    }

    /// `<index-dir>/vectors.json`.
    fn vectors_path(root: &ProjectPath) -> RemotePath {
        RemotePath::new(format!("{}/{VECTORS_FILE}", Self::index_dir(root)))
    }

    /// The text embedded for an index entry (its index-line payload).
    fn entry_text(entry: &MemoryIndexEntry) -> String {
        format!("{} {}", entry.title, entry.hook)
    }

    /// Loads the derived vector doc, or a fresh empty one when absent, malformed,
    /// or produced by a different embedder/dimension (best-effort: never errors on
    /// a stale cache — it is simply rebuilt).
    async fn load_doc(&self, root: &ProjectPath) -> VectorDoc {
        let empty = || VectorDoc::new(self.embedder.id().to_owned(), self.embedder.dimension());
        match self.fs.read(&Self::vectors_path(root)).await {
            Ok(bytes) => match serde_json::from_slice::<VectorDoc>(&bytes) {
                Ok(doc) if doc.matches(self.embedder.id(), self.embedder.dimension()) => doc,
                _ => empty(),
            },
            Err(FsError::NotFound(_)) | Err(_) => empty(),
        }
    }

    /// Persists the derived vector doc best-effort (a write failure is swallowed:
    /// the cache is reconstructible, recall already has its result).
    async fn save_doc(&self, root: &ProjectPath, doc: &VectorDoc) {
        let dir = RemotePath::new(Self::index_dir(root));
        if self.fs.create_dir_all(&dir).await.is_err() {
            return;
        }
        if let Ok(bytes) = serde_json::to_vec_pretty(doc) {
            let _ = self.fs.write(&Self::vectors_path(root), &bytes).await;
        }
    }
}

#[async_trait]
impl MemoryRecall for VectorMemoryRecall {
    async fn recall(
        &self,
        root: &ProjectPath,
        query: &MemoryQuery,
    ) -> Result<Vec<MemoryIndexEntry>, MemoryError> {
        if query.token_budget == 0 {
            return Ok(Vec::new());
        }
        let entries = self.store.read_index(root).await?;
        if entries.is_empty() {
            return Ok(Vec::new());
        }

        // Load the derived cache and compute any missing note vectors.
        let mut doc = self.load_doc(root).await;
        let missing: Vec<&MemoryIndexEntry> = entries
            .iter()
            .filter(|e| !doc.vectors.contains_key(e.slug.as_str()))
            .collect();
        if !missing.is_empty() {
            let texts: Vec<String> = missing.iter().map(|e| Self::entry_text(e)).collect();
            let vectors = self.recall_embed(&texts).await?;
            for (entry, vector) in missing.iter().zip(vectors) {
                doc.vectors.insert(entry.slug.as_str().to_owned(), vector);
            }
            // Drop vectors of notes that no longer exist, then persist best-effort.
            let live: std::collections::HashSet<&str> =
                entries.iter().map(|e| e.slug.as_str()).collect();
            doc.vectors.retain(|slug, _| live.contains(slug.as_str()));
            self.save_doc(root, &doc).await;
        }

        // Embed the query and rank by cosine similarity (descending).
        let query_vec = self
            .recall_embed(std::slice::from_ref(&query.text))
            .await?
            .into_iter()
            .next()
            .unwrap_or_default();

        let mut ranked: Vec<(f32, MemoryIndexEntry)> = entries
            .into_iter()
            .map(|entry| {
                let score = doc
                    .vectors
                    .get(entry.slug.as_str())
                    .map_or(0.0, |v| cosine_similarity(&query_vec, v));
                (score, entry)
            })
            .collect();
        // Stable, deterministic order: score desc, then slug asc for ties.
        ranked.sort_by(|a, b| {
            b.0.partial_cmp(&a.0)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.1.slug.cmp(&b.1.slug))
        });

        let ordered = ranked.into_iter().map(|(_, e)| e);
        Ok(truncate_to_budget(ordered, query.token_budget))
    }
}

impl VectorMemoryRecall {
    /// Embeds `texts`, mapping an [`EmbedderError`] into a [`MemoryError`] only so
    /// the `?` plumbing compiles; in practice [`AdaptiveMemoryRecall`] guards this
    /// path and falls back to naïve before any such error reaches a use case.
    async fn recall_embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
        self.embedder.embed(texts).await.map_err(map_embedder_error)
    }
}

/// Maps an [`EmbedderError`] to a [`MemoryError`] for the (guarded) `?` path.
fn map_embedder_error(e: EmbedderError) -> MemoryError {
    match e {
        EmbedderError::Io(m) => MemoryError::Io(m),
        other => MemoryError::Serialization(other.to_string()),
    }
}

/// Cosine similarity of two equal-length vectors; `0.0` when a length mismatches
/// or either vector is zero (defensive — never panics).
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    if a.len() != b.len() {
        return 0.0;
    }
    let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    if na == 0.0 || nb == 0.0 {
        0.0
    } else {
        dot / (na * nb)
    }
}

// ---------------------------------------------------------------------------
// AdaptiveMemoryRecall — the étage-1/étage-2 switch.
// ---------------------------------------------------------------------------

/// **Pure, I/O-free** routing decision for [`AdaptiveMemoryRecall`].
///
/// Returns `true` (use the vector recall) **iff** the strategy is not
/// [`EmbedderStrategy::None`] **and** the memory is larger than the budget
/// (i.e. there is something to rank/prune semantically). Otherwise the naïve
/// recall is sufficient and cheaper. This is the single objective rule, unit
/// testable without any store.
#[must_use]
pub fn should_use_vector(memory_size: usize, budget: usize, strategy: EmbedderStrategy) -> bool {
    if strategy == EmbedderStrategy::None {
        return false;
    }
    memory_size > budget
}

/// Adaptive [`MemoryRecall`]: routes between a naïve étage-1 recall and a vector
/// étage-2 recall by [`should_use_vector`], with an automatic **fallback to
/// naïve** whenever the vector path fails (embedder unavailable/unsupported).
///
/// Substitutable for either composed recall (Liskov): same emptiness/budget
/// guarantees, and — by construction — it **never fails hard** on an embedder
/// problem. With strategy [`EmbedderStrategy::None`] it is behaviourally identical
/// to the naïve recall (the default product posture, zero dependency).
#[derive(Clone)]
pub struct AdaptiveMemoryRecall {
    naive: Arc<dyn MemoryRecall>,
    vector: Arc<dyn MemoryRecall>,
    store: Arc<dyn MemoryStore>,
    strategy: EmbedderStrategy,
}

impl AdaptiveMemoryRecall {
    /// Builds the switch from the two recalls, the index store (for the pure size
    /// measure), and the active embedder strategy.
    #[must_use]
    pub fn new(
        naive: Arc<dyn MemoryRecall>,
        vector: Arc<dyn MemoryRecall>,
        store: Arc<dyn MemoryStore>,
        strategy: EmbedderStrategy,
    ) -> Self {
        Self {
            naive,
            vector,
            store,
            strategy,
        }
    }
}

#[async_trait]
impl MemoryRecall for AdaptiveMemoryRecall {
    async fn recall(
        &self,
        root: &ProjectPath,
        query: &MemoryQuery,
    ) -> Result<Vec<MemoryIndexEntry>, MemoryError> {
        // Budget-0 short-circuits before any I/O, homogeneously with the naïve and
        // vector recalls (a zero budget can hold no entry: nothing to read/measure).
        if query.token_budget == 0 {
            return Ok(Vec::new());
        }
        // Pure measure from the index (the only I/O is reading the index once;
        // the decision itself is pure via `should_use_vector`).
        let memory_size = index_token_size(&self.store.read_index(root).await?);
        if !should_use_vector(memory_size, query.token_budget, self.strategy) {
            return self.naive.recall(root, query).await;
        }
        // Vector path with naïve fallback — never fail hard on the embedder.
        match self.vector.recall(root, query).await {
            Ok(entries) => Ok(entries),
            Err(_) => self.naive.recall(root, query).await,
        }
    }
}