poimen-memory/crates/mem-store/src/pg_repo.rs

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

/// Vector kind (text or symptom).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum VectorKind {
    Text,
    Symptom,
}

impl std::fmt::Display for VectorKind {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            VectorKind::Text => write!(f, "text"),
            VectorKind::Symptom => write!(f, "symptom"),
        }
    }
}

/// Level (L0, L1, L2).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum Level {
    L0,
    L1,
    L2,
}

/// Memory node (idempotent upsert key: sha256).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryNode {
    pub sha256: String,
    pub level: Level,
    pub project: String,
    pub text: String,
    pub tokens: u32,
}

/// Scored search result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScoredNode {
    pub node: MemoryNode,
    pub distance: f32,
    pub matched_kind: VectorKind,
}

/// PostgreSQL repository (in-memory mock for now).
pub struct PgRepo {
    // Nodes by sha256
    nodes: BTreeMap<String, MemoryNode>,
    // Vectors by (sha256, kind)
    vectors: BTreeMap<(String, VectorKind), Vec<f32>>,
    // Parents edges: child_sha -> vec of parent_shas
    edges: BTreeMap<String, Vec<String>>,
}

impl PgRepo {
    /// Create new repo (mock, no real DB).
    pub fn new() -> Self {
        Self {
            nodes: BTreeMap::new(),
            vectors: BTreeMap::new(),
            edges: BTreeMap::new(),
        }
    }

    /// Upsert node (idempotent).
    pub fn upsert_node(&mut self, node: &MemoryNode) -> Result<()> {
        self.nodes.insert(node.sha256.clone(), node.clone());
        Ok(())
    }

    /// Upsert many nodes (batching embedding calls).
    pub fn upsert_many(&mut self, nodes: &[MemoryNode]) -> Result<()> {
        for node in nodes {
            self.upsert_node(node)?;
        }
        Ok(())
    }

    /// Upsert vector for node.
    pub fn upsert_vector(&mut self, sha: &str, kind: VectorKind, embedding: &[f32]) -> Result<()> {
        if !self.nodes.contains_key(sha) {
            return Err(anyhow::anyhow!("Node {} not found", sha));
        }
        self.vectors.insert((sha.to_string(), kind), embedding.to_vec());
        Ok(())
    }

    /// Insert edges (requires both endpoints exist).
    pub fn insert_edges(&mut self, child: &str, parents: &[String]) -> Result<()> {
        if !self.nodes.contains_key(child) {
            return Err(anyhow::anyhow!("Child node {} not found", child));
        }
        for parent in parents {
            if !self.nodes.contains_key(parent) {
                return Err(anyhow::anyhow!("Parent node {} not found", parent));
            }
        }
        self.edges.insert(child.to_string(), parents.to_vec());
        Ok(())
    }

    /// Search by cosine distance.
    pub fn search(
        &self,
        q: &[f32],
        kind: VectorKind,
        levels: &[Level],
    ) -> Result<Vec<ScoredNode>> {
        let mut results = Vec::new();

        for ((sha, vkind), embedding) in &self.vectors {
            if *vkind != kind {
                continue;
            }

            if let Some(node) = self.nodes.get(sha) {
                if !levels.contains(&node.level) {
                    continue;
                }

                if let Some(dist) = cosine_distance(q, embedding) {
                    results.push(ScoredNode {
                        node: node.clone(),
                        distance: dist,
                        matched_kind: kind,
                    });
                }
            }
        }

        // Sort by distance (ascending)
        results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        Ok(results)
    }

    /// Parents of node.
    pub fn parents_of(&self, sha: &str) -> Result<Vec<MemoryNode>> {
        let parent_shas = self.edges.get(sha).cloned().unwrap_or_default();
        let parents: Vec<_> = parent_shas
            .iter()
            .filter_map(|p_sha| self.nodes.get(p_sha).cloned())
            .collect();
        Ok(parents)
    }

    /// Clear all nodes for project.
    pub fn clear_project(&mut self, project: &str) -> Result<()> {
        let nodes_to_remove: Vec<String> = self
            .nodes
            .iter()
            .filter(|(_, n)| n.project == project)
            .map(|(sha, _)| sha.clone())
            .collect();

        // Remove vectors
        self.vectors.retain(|(sha, _), _| !nodes_to_remove.contains(sha));

        // Remove edges
        self.edges.retain(|child, _| !nodes_to_remove.contains(child));

        // Remove nodes
        self.nodes.retain(|sha, _| !nodes_to_remove.contains(sha));

        Ok(())
    }

    /// Get all nodes.
    pub fn all_nodes(&self) -> Vec<&MemoryNode> {
        self.nodes.values().collect()
    }

    /// Verify: count upserted nodes.
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Verify: count edges.
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }
}

/// Cosine distance (1 - cosine_similarity).
fn cosine_distance(a: &[f32], b: &[f32]) -> Option<f32> {
    if a.len() != b.len() || a.is_empty() {
        return None;
    }

    let mut dot = 0.0;
    let mut norm_a = 0.0;
    let mut norm_b = 0.0;

    for (x, y) in a.iter().zip(b.iter()) {
        dot += x * y;
        norm_a += x * x;
        norm_b += y * y;
    }

    let norm_a = norm_a.sqrt();
    let norm_b = norm_b.sqrt();

    if norm_a == 0.0 || norm_b == 0.0 {
        return None;
    }

    let similarity = dot / (norm_a * norm_b);
    Some(1.0 - similarity)
}