Newer
Older
navi-1 / navi / memory / _ddl.py
"""DDL for memory tables — conditional on pgvector/pg_trgm availability."""

from navi.config import settings


def _build_ddl(pgvector_available: bool) -> list[str]:
    """Return DDL statements depending on whether pgvector is installed."""
    embedding_col = f"embedding vector({settings.embedding_dimensions})," if pgvector_available else ""
    embedding_idx = (
        "CREATE INDEX IF NOT EXISTS idx_memory_facts_embedding ON memory_facts USING hnsw (embedding vector_cosine_ops)"
        if pgvector_available else ""
    )
    chunk_embedding_col = (
        f"embedding vector({settings.embedding_dimensions})," if pgvector_available else ""
    )
    stmts = [
        # Migration: add user_id to existing tables (nullable for legacy)
        "ALTER TABLE memory_facts ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES navi_users(id) ON DELETE CASCADE",
        "ALTER TABLE memory_summary ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES navi_users(id) ON DELETE CASCADE",
        # Merge losers: a fact hidden behind a canonical one written under a different
        # spelling of the same key. Nullable, filtered with IS NULL on read.
        "ALTER TABLE memory_facts ADD COLUMN IF NOT EXISTS merged_into TEXT",
        # Migrate unique constraint from (category, key) to (user_id, category, key)
        "DO $$ BEGIN IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'memory_facts_category_key_key') THEN ALTER TABLE memory_facts DROP CONSTRAINT memory_facts_category_key_key; END IF; END $$;",
        "DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'memory_facts_user_cat_key') THEN ALTER TABLE memory_facts ADD CONSTRAINT memory_facts_user_cat_key UNIQUE (user_id, category, key); END IF; END $$;",
        # Ensure memory_summary has unique constraint on (id, user_id) for ON CONFLICT
        "DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'memory_summary_id_user_key') THEN ALTER TABLE memory_summary ADD CONSTRAINT memory_summary_id_user_key UNIQUE (id, user_id); END IF; END $$;",

        """CREATE TABLE IF NOT EXISTS memory_facts (
            id          TEXT PRIMARY KEY,
            user_id     TEXT REFERENCES navi_users(id) ON DELETE CASCADE,
            category    TEXT NOT NULL,
            key         TEXT NOT NULL,
            value       TEXT NOT NULL,
            created_at  TIMESTAMPTZ NOT NULL,
            updated_at  TIMESTAMPTZ NOT NULL,
            source_session_id TEXT,
            %s
            source      TEXT NOT NULL DEFAULT 'conversation',
            confidence  SMALLINT NOT NULL DEFAULT 70,
            expires_at  TIMESTAMPTZ,
            last_verified_at TIMESTAMPTZ,
            source_context TEXT,
            merged_into TEXT,
            UNIQUE(user_id, category, key)
        )""" % embedding_col,
        "CREATE INDEX IF NOT EXISTS idx_memory_facts_expires ON memory_facts (expires_at) WHERE expires_at IS NOT NULL",
        "CREATE INDEX IF NOT EXISTS idx_memory_facts_source_cat ON memory_facts (source, category)",
        # GIN trigram indexes — only if pg_trgm extension is already installed.
        # CREATE EXTENSION requires superuser/CREATE privilege, so we skip it here.
        """DO $$
        BEGIN
            IF EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm') THEN
                CREATE INDEX IF NOT EXISTS idx_memory_facts_cat_trgm ON memory_facts USING gin (category gin_trgm_ops);
                CREATE INDEX IF NOT EXISTS idx_memory_facts_key_trgm ON memory_facts USING gin (key gin_trgm_ops);
                CREATE INDEX IF NOT EXISTS idx_memory_facts_value_trgm ON memory_facts USING gin (value gin_trgm_ops);
            END IF;
        END $$;""",
        """CREATE TABLE IF NOT EXISTS memory_summary (
            id          INTEGER PRIMARY KEY DEFAULT 1,
            user_id     TEXT REFERENCES navi_users(id) ON DELETE CASCADE,
            content     TEXT NOT NULL,
            generated_at TIMESTAMPTZ NOT NULL,
            UNIQUE(id, user_id)
        )""",
        """CREATE TABLE IF NOT EXISTS session_memory_state (
            session_id  TEXT PRIMARY KEY,
            extracted_at TIMESTAMPTZ NOT NULL
        )""",
        # Raw history dropped by compression: append-only, never updated. No FK to
        # `sessions` on purpose — this DDL can run before the sessions table exists
        # on a fresh database, and a failed CREATE here would silently disable the
        # whole store for the process. PgSessionStore.delete clears the rows.
        f"""CREATE TABLE IF NOT EXISTS memory_chunks (
            id           TEXT PRIMARY KEY,
            user_id      TEXT,
            session_id   TEXT NOT NULL,
            start_seq    INTEGER NOT NULL DEFAULT -1,
            end_seq      INTEGER NOT NULL DEFAULT -1,
            started_at   TIMESTAMPTZ,
            ended_at     TIMESTAMPTZ,
            content      TEXT NOT NULL,
            breadcrumbs  TEXT[] NOT NULL DEFAULT '{{}}',
            content_hash TEXT NOT NULL,
            created_at   TIMESTAMPTZ NOT NULL,
            {chunk_embedding_col}
            UNIQUE(session_id, content_hash)
        )""",
        "CREATE INDEX IF NOT EXISTS idx_memory_chunks_scope "
        "ON memory_chunks (user_id, session_id, ended_at DESC)",
        # `to_tsvector('simple', content)` has no stemming and no stopwords, which is
        # what a mixed-language corpus needs; the expression here must match the
        # query's exactly or the planner cannot use this index.
        "CREATE INDEX IF NOT EXISTS idx_memory_chunks_fts "
        "ON memory_chunks USING gin (to_tsvector('simple', content))",
        "CREATE INDEX IF NOT EXISTS idx_memory_chunks_breadcrumbs "
        "ON memory_chunks USING gin (breadcrumbs)",
    ]
    if pgvector_available:
        # An older database may hold a `memory_chunks` created before pgvector was
        # available: CREATE TABLE IF NOT EXISTS leaves it without the column, and
        # every INSERT then dies on "column embedding does not exist". Add it before
        # the HNSW index so the table heals itself on the first run that can use
        # vectors.
        stmts.append(
            "ALTER TABLE memory_chunks ADD COLUMN IF NOT EXISTS "
            "embedding vector(%d)" % settings.embedding_dimensions
        )
    if pgvector_available:
        stmts.append(
            "CREATE INDEX IF NOT EXISTS idx_memory_chunks_embedding "
            "ON memory_chunks USING hnsw (embedding vector_cosine_ops)"
        )
    if embedding_idx:
        stmts.insert(2, embedding_idx)
    return stmts