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

use anyhow::Result;
use pgvector::Vector;
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;

/// L0: Evidence chunk (raw source span)
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct ChunkL0 {
    pub id: Uuid,
    pub project: String,
    pub query_id: String,
    pub source: String,  // "pi", "claude", "transcript"
    pub content: String,
    pub tokens: i32,
}

/// L1: Per-query memory (1024 token bound)
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct MemoryL1 {
    pub id: Uuid,
    pub project: String,
    pub query_id: String,
    pub content: String,
    pub tokens: i32,
    #[sqlx(skip)]
    pub embedding: Option<Vec<f32>>,
    pub chunks_seen: i32,
    pub chunks_used: i32,
    pub run_id: String,
}

/// L2: Project synthesis (1024 token bound)
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct MemoryL2 {
    pub id: Uuid,
    pub project: String,
    pub content: String,
    pub tokens: i32,
    #[sqlx(skip)]
    pub embedding: Option<Vec<f32>>,
    pub l1_count: i32,
    pub run_id: String,
}

/// Reference corpus entry (documentation, skills, etc.)
#[derive(Debug, Clone, Serialize, Deserialize, sqlx::FromRow)]
pub struct RefCorpus {
    pub id: Uuid,
    pub project: String,
    pub name: String,
    pub content: String,
    #[sqlx(skip)]
    pub embedding: Option<Vec<f32>>,
}

/// Vector record for embedding storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VectorRecord {
    pub id: String,
    pub chunk_id: String,
    pub kind: String,  // "l1", "l2", "corpus"
    pub embedding: Vec<f32>,
    pub tokens: u32,
}

/// Scored search result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ScoredResult<T> {
    pub item: T,
    pub score: f32,
}

/// PostgreSQL vector store — backed by pgvector
pub struct VectorStore {
    pool: PgPool,
}

impl VectorStore {
    /// Create or get vector store from connection pool
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    /// Get access to the connection pool (for testing)
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Store L0 chunk
    pub async fn store_chunk_l0(&self, chunk: &ChunkL0) -> Result<()> {
        sqlx::query(
            "INSERT INTO chunks_l0 (id, project, query_id, source, content, tokens) 
             VALUES ($1, $2, $3, $4, $5, $6)
             ON CONFLICT (id) DO NOTHING",
        )
        .bind(chunk.id)
        .bind(&chunk.project)
        .bind(&chunk.query_id)
        .bind(&chunk.source)
        .bind(&chunk.content)
        .bind(chunk.tokens)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Store L1 memory with embedding
    pub async fn store_memory_l1(
        &self,
        mem: &MemoryL1,
        embedding: &Vector,
    ) -> Result<()> {
        sqlx::query(
            "INSERT INTO memories_l1 (id, project, query_id, content, tokens, embedding, chunks_seen, chunks_used, run_id)
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
             ON CONFLICT (project, query_id) DO UPDATE SET
               content = EXCLUDED.content,
               tokens = EXCLUDED.tokens,
               embedding = EXCLUDED.embedding,
               chunks_seen = EXCLUDED.chunks_seen,
               chunks_used = EXCLUDED.chunks_used,
               updated_at = CURRENT_TIMESTAMP,
               run_id = EXCLUDED.run_id",
        )
        .bind(mem.id)
        .bind(&mem.project)
        .bind(&mem.query_id)
        .bind(&mem.content)
        .bind(mem.tokens)
        .bind(embedding)
        .bind(mem.chunks_seen)
        .bind(mem.chunks_used)
        .bind(&mem.run_id)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Store L2 synthesis with embedding
    pub async fn store_memory_l2(
        &self,
        mem: &MemoryL2,
        embedding: &Vector,
    ) -> Result<()> {
        sqlx::query(
            "INSERT INTO memories_l2 (id, project, content, tokens, embedding, l1_count, run_id)
             VALUES ($1, $2, $3, $4, $5, $6, $7)
             ON CONFLICT (project) DO UPDATE SET
               content = EXCLUDED.content,
               tokens = EXCLUDED.tokens,
               embedding = EXCLUDED.embedding,
               l1_count = EXCLUDED.l1_count,
               updated_at = CURRENT_TIMESTAMP,
               run_id = EXCLUDED.run_id",
        )
        .bind(mem.id)
        .bind(&mem.project)
        .bind(&mem.content)
        .bind(mem.tokens)
        .bind(embedding)
        .bind(mem.l1_count)
        .bind(&mem.run_id)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Store reference corpus entry with embedding
    pub async fn store_corpus(
        &self,
        project: &str,
        name: &str,
        content: &str,
        embedding: &Vector,
    ) -> Result<()> {
        sqlx::query(
            "INSERT INTO reference_corpus (id, project, name, content, embedding)
             VALUES ($1, $2, $3, $4, $5)
             ON CONFLICT (project, name) DO UPDATE SET
               content = EXCLUDED.content,
               embedding = EXCLUDED.embedding",
        )
        .bind(Uuid::new_v4())
        .bind(project)
        .bind(name)
        .bind(content)
        .bind(embedding)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    /// Search L1 memories by embedding similarity
    pub async fn search_l1(
        &self,
        project: &str,
        embedding: &Vector,
        limit: i64,
    ) -> Result<Vec<ScoredResult<MemoryL1>>> {
        let rows = sqlx::query_as::<_, (Uuid, String, String, String, i32, i32, i32, String)>(
            "SELECT id, project, query_id, content, tokens, chunks_seen, chunks_used, run_id
             FROM memories_l1
             WHERE project = $1
             ORDER BY embedding <=> $2
             LIMIT $3",
        )
        .bind(project)
        .bind(embedding)
        .bind(limit)
        .fetch_all(&self.pool)
        .await?;

        Ok(rows
            .into_iter()
            .enumerate()
            .map(|(i, (id, proj, qid, content, tokens, seen, used, run))| {
                // Calculate similarity score (1 / (1 + distance))
                let distance = (i as f32) * 0.1; // Rough approximation from rank
                let score = 1.0 / (1.0 + distance);
                ScoredResult {
                    item: MemoryL1 {
                        id,
                        project: proj,
                        query_id: qid,
                        content,
                        tokens,
                        embedding: None,
                        chunks_seen: seen,
                        chunks_used: used,
                        run_id: run,
                    },
                    score,
                }
            })
            .collect())
    }

    /// Search L2 memories by embedding similarity
    pub async fn search_l2(
        &self,
        project: &str,
        embedding: &Vector,
    ) -> Result<Option<ScoredResult<MemoryL2>>> {
        let row = sqlx::query_as::<_, (Uuid, String, String, i32, i32, String)>(
            "SELECT id, project, content, tokens, l1_count, run_id
             FROM memories_l2
             WHERE project = $1
             ORDER BY embedding <=> $2
             LIMIT 1",
        )
        .bind(project)
        .bind(embedding)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row.map(|(id, proj, content, tokens, count, run)| ScoredResult {
            item: MemoryL2 {
                id,
                project: proj,
                content,
                tokens,
                embedding: None,
                l1_count: count,
                run_id: run,
            },
            score: 0.95, // Perfect match for same project
        }))
    }

    /// Search reference corpus by embedding similarity
    pub async fn search_corpus(
        &self,
        project: &str,
        embedding: &Vector,
        limit: i64,
    ) -> Result<Vec<ScoredResult<RefCorpus>>> {
        let rows = sqlx::query_as::<_, (Uuid, String, String, String)>(
            "SELECT id, project, name, content
             FROM reference_corpus
             WHERE project = $1
             ORDER BY embedding <=> $2
             LIMIT $3",
        )
        .bind(project)
        .bind(embedding)
        .bind(limit)
        .fetch_all(&self.pool)
        .await?;

        Ok(rows
            .into_iter()
            .enumerate()
            .map(|(i, (id, proj, name, content))| {
                let distance = (i as f32) * 0.1;
                let score = 1.0 / (1.0 + distance);
                ScoredResult {
                    item: RefCorpus {
                        id,
                        project: proj,
                        name,
                        content,
                        embedding: None,
                    },
                    score,
                }
            })
            .collect())
    }

    /// Get L1 memory by query_id
    pub async fn get_l1(&self, project: &str, query_id: &str) -> Result<Option<MemoryL1>> {
        let row = sqlx::query_as::<_, (Uuid, String, String, String, i32, i32, i32, String)>(
            "SELECT id, project, query_id, content, tokens, chunks_seen, chunks_used, run_id
             FROM memories_l1
             WHERE project = $1 AND query_id = $2",
        )
        .bind(project)
        .bind(query_id)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row.map(|(id, proj, qid, content, tokens, seen, used, run)| MemoryL1 {
            id,
            project: proj,
            query_id: qid,
            content,
            tokens,
            embedding: None,
            chunks_seen: seen,
            chunks_used: used,
            run_id: run,
        }))
    }

    /// Get L2 memory by project
    pub async fn get_l2(&self, project: &str) -> Result<Option<MemoryL2>> {
        let row = sqlx::query_as::<_, (Uuid, String, String, i32, i32, String)>(
            "SELECT id, project, content, tokens, l1_count, run_id
             FROM memories_l2
             WHERE project = $1",
        )
        .bind(project)
        .fetch_optional(&self.pool)
        .await?;

        Ok(row.map(|(id, proj, content, tokens, count, run)| MemoryL2 {
            id,
            project: proj,
            content,
            tokens,
            embedding: None,
            l1_count: count,
            run_id: run,
        }))
    }

    /// Get L0 chunks for a query (for provenance)
    pub async fn get_l0_chunks(&self, project: &str, query_id: &str) -> Result<Vec<ChunkL0>> {
        sqlx::query_as::<_, (Uuid, String, String, String, String, i32)>(
            "SELECT id, project, query_id, source, content, tokens
             FROM chunks_l0
             WHERE project = $1 AND query_id = $2
             ORDER BY created_at",
        )
        .bind(project)
        .bind(query_id)
        .fetch_all(&self.pool)
        .await?
        .into_iter()
        .map(|(id, proj, qid, src, content, tokens)| {
            Ok(ChunkL0 {
                id,
                project: proj,
                query_id: qid,
                source: src,
                content,
                tokens,
            })
        })
        .collect()
    }
}