use anyhow::{anyhow, Result}; use pgvector::Vector; use serde::{Deserialize, Serialize}; use sqlx::{PgPool, Row}; /// Memory level #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] pub enum Level { L0, L1, L2, R, } impl std::fmt::Display for Level { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Level::L0 => write!(f, "L0"), Level::L1 => write!(f, "L1"), Level::L2 => write!(f, "L2"), Level::R => write!(f, "R"), } } } impl std::str::FromStr for Level { type Err = anyhow::Error; fn from_str(s: &str) -> Result { match s { "L0" => Ok(Level::L0), "L1" => Ok(Level::L1), "L2" => Ok(Level::L2), "R" => Ok(Level::R), _ => Err(anyhow!("invalid level: {}", s)), } } } /// Vector kind (text embedding or symptom projection) #[derive(Debug, Clone, Copy, PartialEq, Eq, 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"), } } } /// Search scope (project-specific or federated) #[derive(Debug, Clone, PartialEq, Eq)] pub enum Scope { Project(String), // project-specific search AllProjects, // federated across all projects (for tool-failure lookups) } /// Memory node (content-addressable by sha256) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryNode { pub sha256: String, pub level: Level, pub project: String, pub query_id: Option, // NULL at L2, R pub run_id: String, pub t: i32, pub source: Option, // set at L0; URI at R pub text: String, } /// Scored search result with metadata #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ScoredNode { pub node: MemoryNode, pub distance: f32, // raw cosine distance (not similarity) pub matched_kind: VectorKind, } /// Failure signature hit (exact-match tier) #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SignatureHit { pub node_sha: String, pub tool: String, pub raw: String, pub seen_count: i32, } /// PostgreSQL repository for memory projection pub struct PgRepo { pool: PgPool, } impl PgRepo { /// Connect to Postgres and run migrations pub async fn connect(url: &str) -> Result { let pool = PgPool::connect(url).await?; sqlx::migrate!("./migrations") .run(&pool) .await?; Ok(Self { pool }) } /// Upsert a memory node (ON CONFLICT DO NOTHING — idempotent) pub async fn upsert_node(&self, node: &MemoryNode) -> Result<()> { let level_str = node.level.to_string(); sqlx::query( r#" INSERT INTO memory_node (sha256, level, project, query_id, run_id, t, source, text) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ON CONFLICT (sha256) DO NOTHING "#, ) .bind(&node.sha256) .bind(&level_str) .bind(&node.project) .bind(&node.query_id) .bind(&node.run_id) .bind(node.t) .bind(&node.source) .bind(&node.text) .execute(&self.pool) .await?; Ok(()) } /// Store an embedding for a node (text or symptom kind) pub async fn upsert_vector( &self, node_sha: &str, kind: VectorKind, embedding: &[f32], ) -> Result<()> { if embedding.len() != 768 { return Err(anyhow!( "invalid embedding dimension: expected 768, got {}", embedding.len() )); } let kind_str = kind.to_string(); let vec = Vector::from(embedding.to_vec()); sqlx::query( r#" INSERT INTO memory_vector (node_sha, kind, embedding) VALUES ($1, $2, $3) ON CONFLICT (node_sha, kind) DO UPDATE SET embedding = $3 "#, ) .bind(node_sha) .bind(&kind_str) .bind(&vec) .execute(&self.pool) .await?; Ok(()) } /// Insert edges (child → parent pointers) — fails if endpoints don't exist pub async fn insert_edges(&self, child_sha: &str, parent_shas: &[String]) -> Result<()> { // Batch insert edges; foreign key constraints ensure endpoints exist for parent_sha in parent_shas { sqlx::query( r#" INSERT INTO memory_edge (child_sha, parent_sha) VALUES ($1, $2) ON CONFLICT (child_sha, parent_sha) DO NOTHING "#, ) .bind(child_sha) .bind(parent_sha) .execute(&self.pool) .await?; } Ok(()) } /// Search by embedding (cosine distance, kind-filtered) /// - For text kind: uses partial index WHERE kind = 'text' /// - For symptom kind: uses partial index WHERE kind = 'symptom' pub async fn search( &self, query_embedding: &[f32], kind: VectorKind, levels: &[Level], scope: &Scope, k: usize, ) -> Result> { if query_embedding.len() != 768 { return Err(anyhow!("query embedding must be 768-dim")); } let q_vec = Vector::from(query_embedding.to_vec()); let kind_str = kind.to_string(); let level_strs: Vec = levels.iter().map(|l| l.to_string()).collect(); // Build the WHERE clause based on scope let (where_clause, project_param) = match scope { Scope::Project(proj) => ("AND n.project = $4".to_string(), Some(proj.clone())), Scope::AllProjects => ("".to_string(), None), }; let query_sql = format!( r#" SELECT n.sha256, n.level, n.project, n.query_id, n.run_id, n.t, n.source, n.text, v.embedding <=> $1 AS distance FROM memory_vector v JOIN memory_node n ON v.node_sha = n.sha256 WHERE v.kind = $2 AND n.level = ANY($3) {} AND (SELECT COUNT(*) FROM memory_supersede WHERE old_sha = n.sha256) = 0 ORDER BY v.embedding <=> $1 ASC LIMIT $5 "#, if where_clause.is_empty() { "" } else { &where_clause } ); let rows = if let Some(proj) = project_param { sqlx::query(&query_sql) .bind(&q_vec) .bind(&kind_str) .bind(&level_strs) .bind(&proj) .bind(k as i64) .fetch_all(&self.pool) .await? } else { sqlx::query(&query_sql) .bind(&q_vec) .bind(&kind_str) .bind(&level_strs) .bind(k as i64) .fetch_all(&self.pool) .await? }; let mut results = Vec::new(); for row in rows { let level_str: String = row.get("level"); let node = MemoryNode { sha256: row.get("sha256"), level: level_str.parse()?, project: row.get("project"), query_id: row.get("query_id"), run_id: row.get("run_id"), t: row.get("t"), source: row.get("source"), text: row.get("text"), }; let distance: f32 = row.get("distance"); results.push(ScoredNode { node, distance, matched_kind: kind, }); } Ok(results) } /// Exact-match lookup on failure signature pub async fn lookup_signature(&self, sig_sha: &str) -> Result> { let row = sqlx::query( r#" SELECT node_sha, tool, raw, seen_count FROM failure_signature WHERE sig_sha = $1 "#, ) .bind(sig_sha) .fetch_optional(&self.pool) .await?; Ok(row.map(|r| SignatureHit { node_sha: r.get("node_sha"), tool: r.get("tool"), raw: r.get("raw"), seen_count: r.get("seen_count"), })) } /// Traverse parent nodes via edges pub async fn parents_of(&self, sha: &str) -> Result> { let rows = sqlx::query( r#" SELECT n.sha256, n.level, n.project, n.query_id, n.run_id, n.t, n.source, n.text FROM memory_node n JOIN memory_edge e ON e.parent_sha = n.sha256 WHERE e.child_sha = $1 "#, ) .bind(sha) .fetch_all(&self.pool) .await?; let mut parents = Vec::new(); for row in rows { let level_str: String = row.get("level"); parents.push(MemoryNode { sha256: row.get("sha256"), level: level_str.parse()?, project: row.get("project"), query_id: row.get("query_id"), run_id: row.get("run_id"), t: row.get("t"), source: row.get("source"), text: row.get("text"), }); } Ok(parents) } /// Clear all nodes for a project (edges cascade delete) pub async fn clear_project(&self, project: &str) -> Result<()> { sqlx::query("DELETE FROM memory_node WHERE project = $1") .bind(project) .execute(&self.pool) .await?; Ok(()) } /// Get node count (for testing) pub async fn node_count(&self) -> Result { let row = sqlx::query("SELECT COUNT(*) as cnt FROM memory_node") .fetch_one(&self.pool) .await?; Ok(row.get("cnt")) } /// Get edge count (for testing) pub async fn edge_count(&self) -> Result { let row = sqlx::query("SELECT COUNT(*) as cnt FROM memory_edge") .fetch_one(&self.pool) .await?; Ok(row.get("cnt")) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_level_display() { assert_eq!(Level::L0.to_string(), "L0"); assert_eq!(Level::L1.to_string(), "L1"); assert_eq!(Level::L2.to_string(), "L2"); assert_eq!(Level::R.to_string(), "R"); } #[test] fn test_vector_kind_display() { assert_eq!(VectorKind::Text.to_string(), "text"); assert_eq!(VectorKind::Symptom.to_string(), "symptom"); } #[test] fn test_scope_variants() { let proj_scope = Scope::Project("test".to_string()); let all_scope = Scope::AllProjects; assert_ne!(proj_scope, all_scope); } }