145 lines
5.7 KiB
Python
145 lines
5.7 KiB
Python
"""Add memory tables and user encryption_key column.
|
|
|
|
Memory tables:
|
|
memory_core — per-user key/value preferences (encrypted)
|
|
memory_associative — semantic memory with pgvector embedding (encrypted)
|
|
memory_episodic — session summaries (encrypted)
|
|
memory_proactive — behavioral patterns (encrypted)
|
|
|
|
Also adds encryption_key column to users table.
|
|
|
|
Revision ID: 004
|
|
Revises: 003
|
|
Create Date: 2026-03-08
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
revision: str = "004"
|
|
down_revision: Union[str, None] = "003"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# ── Enable pgvector extension (idempotent) ────────────────────────────────
|
|
op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
|
|
|
# ── Add encryption_key to users ───────────────────────────────────────────
|
|
op.add_column(
|
|
"users",
|
|
sa.Column("encryption_key", sa.String(64), nullable=True),
|
|
)
|
|
|
|
# ── memory_core ───────────────────────────────────────────────────────────
|
|
op.create_table(
|
|
"memory_core",
|
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
|
sa.Column(
|
|
"user_id",
|
|
postgresql.UUID(as_uuid=False),
|
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("key", sa.String(255), nullable=False),
|
|
sa.Column("value_encrypted", sa.Text, nullable=False),
|
|
sa.Column(
|
|
"updated_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.func.now(),
|
|
),
|
|
)
|
|
op.create_index("ix_memory_core_user_id", "memory_core", ["user_id"])
|
|
|
|
# ── memory_associative ────────────────────────────────────────────────────
|
|
# The embedding column uses pgvector's vector(1536) type.
|
|
op.create_table(
|
|
"memory_associative",
|
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
|
sa.Column(
|
|
"user_id",
|
|
postgresql.UUID(as_uuid=False),
|
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("content_encrypted", sa.Text, nullable=False),
|
|
sa.Column("entity_type", sa.String(100), nullable=True),
|
|
sa.Column("entity_id", sa.String(255), nullable=True),
|
|
sa.Column(
|
|
"updated_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.func.now(),
|
|
),
|
|
)
|
|
# Add the pgvector column separately (not supported by generic sa types)
|
|
op.execute(
|
|
"ALTER TABLE memory_associative ADD COLUMN embedding vector(1536);"
|
|
)
|
|
op.create_index("ix_memory_associative_user_id", "memory_associative", ["user_id"])
|
|
# IVFFlat index for approximate nearest-neighbour search
|
|
op.execute(
|
|
"CREATE INDEX ix_memory_associative_embedding "
|
|
"ON memory_associative USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);"
|
|
)
|
|
|
|
# ── memory_episodic ───────────────────────────────────────────────────────
|
|
op.create_table(
|
|
"memory_episodic",
|
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
|
sa.Column(
|
|
"user_id",
|
|
postgresql.UUID(as_uuid=False),
|
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("summary_encrypted", sa.Text, nullable=False),
|
|
sa.Column("session_id", sa.String(255), nullable=False),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.func.now(),
|
|
),
|
|
)
|
|
op.create_index("ix_memory_episodic_user_id", "memory_episodic", ["user_id"])
|
|
op.create_index("ix_memory_episodic_session_id", "memory_episodic", ["session_id"])
|
|
|
|
# ── memory_proactive ──────────────────────────────────────────────────────
|
|
op.create_table(
|
|
"memory_proactive",
|
|
sa.Column("id", postgresql.UUID(as_uuid=False), primary_key=True),
|
|
sa.Column(
|
|
"user_id",
|
|
postgresql.UUID(as_uuid=False),
|
|
sa.ForeignKey("users.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
),
|
|
sa.Column("pattern_encrypted", sa.Text, nullable=False),
|
|
sa.Column("confidence", sa.Float, nullable=False, server_default="0.5"),
|
|
sa.Column("source", sa.String(50), nullable=False, server_default="inferred"),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.func.now(),
|
|
),
|
|
)
|
|
op.create_index("ix_memory_proactive_user_id", "memory_proactive", ["user_id"])
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.drop_table("memory_proactive")
|
|
op.drop_table("memory_episodic")
|
|
op.drop_index("ix_memory_associative_embedding", "memory_associative")
|
|
op.drop_table("memory_associative")
|
|
op.drop_table("memory_core")
|
|
op.drop_column("users", "encryption_key")
|