chronos_core/vector_store.py

"""
KAAL — Vector Store
==========================
pgvector semantic search layer backed by Neon PostgreSQL.
Replaces ChromaDB while keeping the same API surface.

Dual-retrieval pipeline:
  • pgvector  → broad semantic recall  (cosine similarity, fuzzy)
  • PostgreSQL → precise temporal / entity filtering (deterministic)

Embeddings: all-MiniLM-L6-v2 (384 dims) via sentence-transformers.
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime
from typing import Optional

from .models import EventRecord

logger = logging.getLogger("chronos.vector_store")


class VectorStore:
    """
    pgvector-backed semantic search over Chronos events.
    Shares the same asyncpg pool as the MemoryStore for efficiency.
    """

    EMBEDDING_DIM = 384  # all-MiniLM-L6-v2

    def __init__(self, pool=None):
        # Pool is injected by api/main.py after MemoryStore initializes
        self._pool = pool
        self._model = None

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def initialize(self, pool=None) -> None:
        """Create the vectors table and HNSW index if not present."""
        if pool:
            self._pool = pool

        if not self._pool:
            raise RuntimeError("VectorStore requires an asyncpg pool.")

        async with self._pool.acquire() as conn:
            await conn.execute(f"""
                CREATE TABLE IF NOT EXISTS event_vectors (
                    event_id    TEXT PRIMARY KEY REFERENCES events(id) ON DELETE CASCADE,
                    source_id   TEXT NOT NULL,
                    owner_id    TEXT NOT NULL,
                    scope       TEXT NOT NULL DEFAULT 'default',
                    embedding   vector({self.EMBEDDING_DIM}) NOT NULL,
                    embed_text  TEXT NOT NULL,
                    timestamp   TIMESTAMPTZ NOT NULL
                );
            """)
            # Migration: add scope column if upgrading from an older schema
            await conn.execute("""
                DO $$ BEGIN
                    ALTER TABLE event_vectors ADD COLUMN IF NOT EXISTS scope TEXT NOT NULL DEFAULT 'default';
                EXCEPTION WHEN others THEN NULL;
                END $$;
                
                CREATE INDEX IF NOT EXISTS idx_vectors_source
                    ON event_vectors(source_id);
                CREATE INDEX IF NOT EXISTS idx_vectors_owner
                    ON event_vectors(owner_id);
                CREATE INDEX IF NOT EXISTS idx_vectors_scope
                    ON event_vectors(scope);
            """)

        # Load embedding model in background thread without blocking port binding
        import asyncio
        asyncio.create_task(asyncio.to_thread(self._load_model))
        logger.info(f"Vector store initialized (pgvector {self.EMBEDDING_DIM}d)")

    def _load_model(self) -> None:
        """Load the sentence-transformer model (cached after first call)."""
        if self._model is not None:
            return
        from sentence_transformers import SentenceTransformer
        self._model = SentenceTransformer("all-MiniLM-L6-v2")
        logger.info("Embedding model loaded: all-MiniLM-L6-v2")

    def _embed(self, text: str) -> list[float]:
        """Embed text using sentence-transformers."""
        if self._model is None:
            self._load_model()
        return self._model.encode(text, normalize_embeddings=True).tolist()

    # ------------------------------------------------------------------
    # Insert
    # ------------------------------------------------------------------

    async def add_event(self, event: EventRecord) -> None:
        """Embed and store a single event vector."""
        import asyncio
        embed_text = self._build_embed_text(event)
        embedding = await asyncio.to_thread(self._embed, embed_text)
        owner_id = event.metadata_json.get("owner_id", event.source_id)

        async with self._pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO event_vectors
                    (event_id, source_id, owner_id, scope, embedding, embed_text, timestamp)
                VALUES ($1, $2, $3, $4, $5::vector, $6, $7)
                ON CONFLICT (event_id) DO UPDATE SET
                    embedding = EXCLUDED.embedding,
                    embed_text = EXCLUDED.embed_text,
                    scope = EXCLUDED.scope
                """,
                event.id, event.source_id, owner_id, event.scope,
                f"[{','.join(str(x) for x in embedding)}]",
                embed_text, event.timestamp,
            )

    async def add_events_batch(self, events: list[EventRecord]) -> None:
        """Embed and store multiple events — embeddings computed in parallel."""
        if not events:
            return
        import asyncio

        embed_texts = [self._build_embed_text(e) for e in events]
        # Encode all at once (sentence-transformers batches efficiently)
        embeddings = await asyncio.to_thread(
            lambda: self._model.encode(embed_texts, normalize_embeddings=True, batch_size=32).tolist()
        )

        rows = [
            (
                e.id, e.source_id,
                e.metadata_json.get("owner_id", e.source_id),
                e.scope,
                f"[{','.join(str(x) for x in emb)}]",
                txt, e.timestamp,
            )
            for e, emb, txt in zip(events, embeddings, embed_texts)
        ]

        async with self._pool.acquire() as conn:
            await conn.executemany(
                """
                INSERT INTO event_vectors
                    (event_id, source_id, owner_id, scope, embedding, embed_text, timestamp)
                VALUES ($1, $2, $3, $4, $5::vector, $6, $7)
                ON CONFLICT (event_id) DO UPDATE SET
                    embedding = EXCLUDED.embedding,
                    embed_text = EXCLUDED.embed_text,
                    scope = EXCLUDED.scope
                """,
                rows,
            )
        logger.info(f"Batch added {len(events)} event vectors to pgvector")

    # ------------------------------------------------------------------
    # Search
    # ------------------------------------------------------------------

    async def semantic_search(
        self,
        query: str,
        n_results: int = 20,
        source_ids: Optional[list[str]] = None,
        owner_id: Optional[str] = None,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        scope: Optional[str] = None,
        similarity_threshold: Optional[float] = None,
    ) -> list[dict]:
        """
        Cosine similarity search over embedded events.

        similarity_threshold: cosine distance cutoff (lower = stricter).
          None  → read from SMRITI_SIMILARITY_THRESHOLD env var (default 0.15).
          0.10  → ≥90% cosine similarity (very strict)
          0.15  → ≥85% cosine similarity (default)
          0.30  → ≥70% cosine similarity (lenient)
        """
        from chronos_core.config import SIMILARITY_THRESHOLD
        threshold = similarity_threshold if similarity_threshold is not None else SIMILARITY_THRESHOLD
        import asyncio

        query_embedding = await asyncio.to_thread(self._embed, query)
        vec_str = f"[{','.join(str(x) for x in query_embedding)}]"

        conditions, params = [], [vec_str]
        i = 2

        # Tenant isolation (highest priority)
        if owner_id:
            conditions.append(f"ev.owner_id = ${i}"); params.append(owner_id); i += 1
        elif source_ids:
            conditions.append(f"ev.source_id = ANY(${i})"); params.append(source_ids); i += 1

        # Hard scope isolation
        if scope:
            conditions.append(f"ev.scope = ${i}"); params.append(scope); i += 1

        # Time range filters
        if start_time:
            conditions.append(f"ev.timestamp >= ${i}"); params.append(start_time); i += 1
        if end_time:
            conditions.append(f"ev.timestamp <= ${i}"); params.append(end_time); i += 1

        # Configurable similarity threshold (cosine distance, not similarity)
        params.append(threshold)  # resolved: never None
        threshold_param = i; i += 1

        params.append(n_results)
        limit_param = i

        where = f"WHERE {' AND '.join(conditions)}" if conditions else ""

        query_sql = f"""
            SELECT ev.event_id,
                   (ev.embedding <=> $1::vector) AS distance,
                   ev.source_id, ev.owner_id, ev.scope, ev.embed_text, ev.timestamp
            FROM event_vectors ev
            JOIN events e ON ev.event_id = e.id AND e.valid_to IS NULL
            {where}
              AND (ev.embedding <=> $1::vector) <= ${threshold_param}
            ORDER BY ev.embedding <=> $1::vector
            LIMIT ${limit_param}
        """

        async with self._pool.acquire() as conn:
            rows = await conn.fetch(query_sql, *params)

        return [
            {
                "id": r["event_id"],
                "distance": float(r["distance"]),
                "metadata": {
                    "source_id": r["source_id"],
                    "owner_id": r["owner_id"],
                    "scope": r["scope"],
                    "timestamp": r["timestamp"].isoformat(),
                },
                "document": r["embed_text"],
            }
            for r in rows
        ]

    async def count(self) -> int:
        """Get total number of stored embeddings."""
        async with self._pool.acquire() as conn:
            return await conn.fetchval("SELECT COUNT(*) FROM event_vectors")

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _build_embed_text(event: EventRecord) -> str:
        """Build rich text for embedding: SVO + raw text."""
        parts = [f"{event.subject} {event.verb} {event.object}"]
        if event.raw_text:
            parts.append(event.raw_text)
        return " | ".join(parts)