Deploy Poimen Memory K8s cluster with ArgoCD tracking (M2.2, M3.5-M3.7)

crates/mem-store/src/pg_repo.rs

@@ -0,0 +1,210 @@
+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)
+}