275 lines
16 KiB
Python
275 lines
16 KiB
Python
"""initial_schema
|
|
|
|
Revision ID: 5a366d414dce
|
|
Revises:
|
|
Create Date: 2025-11-27 11:54:19.228030
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
from pgvector.sqlalchemy import Vector
|
|
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = '5a366d414dce'
|
|
down_revision: Union[str, Sequence[str], None] = None
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Upgrade schema - create all tables from scratch."""
|
|
|
|
# Enable required extensions
|
|
op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
|
|
op.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
|
|
# Create banks table
|
|
op.create_table(
|
|
'banks',
|
|
sa.Column('bank_id', sa.Text(), nullable=False),
|
|
sa.Column('name', sa.Text(), nullable=True),
|
|
sa.Column('personality', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
|
|
sa.Column('background', sa.Text(), nullable=True),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.PrimaryKeyConstraint('bank_id', name=op.f('pk_banks'))
|
|
)
|
|
|
|
# Create documents table
|
|
op.create_table(
|
|
'documents',
|
|
sa.Column('id', sa.Text(), nullable=False),
|
|
sa.Column('bank_id', sa.Text(), nullable=False),
|
|
sa.Column('original_text', sa.Text(), nullable=True),
|
|
sa.Column('content_hash', sa.Text(), nullable=True),
|
|
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.PrimaryKeyConstraint('id', 'bank_id', name=op.f('pk_documents'))
|
|
)
|
|
op.create_index('idx_documents_bank_id', 'documents', ['bank_id'])
|
|
op.create_index('idx_documents_content_hash', 'documents', ['content_hash'])
|
|
|
|
# Create async_operations table
|
|
op.create_table(
|
|
'async_operations',
|
|
sa.Column('operation_id', postgresql.UUID(as_uuid=True), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
|
sa.Column('bank_id', sa.Text(), nullable=False),
|
|
sa.Column('operation_type', sa.Text(), nullable=False),
|
|
sa.Column('status', sa.Text(), server_default='pending', nullable=False),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('completed_at', postgresql.TIMESTAMP(timezone=True), nullable=True),
|
|
sa.Column('error_message', sa.Text(), nullable=True),
|
|
sa.Column('result_metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
|
|
sa.PrimaryKeyConstraint('operation_id', name=op.f('pk_async_operations')),
|
|
sa.CheckConstraint("status IN ('pending', 'processing', 'completed', 'failed')", name='async_operations_status_check')
|
|
)
|
|
op.create_index('idx_async_operations_bank_id', 'async_operations', ['bank_id'])
|
|
op.create_index('idx_async_operations_status', 'async_operations', ['status'])
|
|
op.create_index('idx_async_operations_bank_status', 'async_operations', ['bank_id', 'status'])
|
|
|
|
# Create entities table
|
|
op.create_table(
|
|
'entities',
|
|
sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
|
sa.Column('canonical_name', sa.Text(), nullable=False),
|
|
sa.Column('bank_id', sa.Text(), nullable=False),
|
|
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
|
|
sa.Column('first_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('last_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('mention_count', sa.Integer(), server_default='1', nullable=False),
|
|
sa.PrimaryKeyConstraint('id', name=op.f('pk_entities'))
|
|
)
|
|
op.create_index('idx_entities_bank_id', 'entities', ['bank_id'])
|
|
op.create_index('idx_entities_canonical_name', 'entities', ['canonical_name'])
|
|
op.create_index('idx_entities_bank_name', 'entities', ['bank_id', 'canonical_name'])
|
|
# Create unique index on (bank_id, LOWER(canonical_name)) for entity resolution
|
|
op.execute('CREATE UNIQUE INDEX idx_entities_bank_lower_name ON entities (bank_id, LOWER(canonical_name))')
|
|
|
|
# Create memory_units table
|
|
op.create_table(
|
|
'memory_units',
|
|
sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
|
sa.Column('bank_id', sa.Text(), nullable=False),
|
|
sa.Column('document_id', sa.Text(), nullable=True),
|
|
sa.Column('text', sa.Text(), nullable=False),
|
|
sa.Column('embedding', Vector(384), nullable=True),
|
|
sa.Column('context', sa.Text(), nullable=True),
|
|
sa.Column('event_date', postgresql.TIMESTAMP(timezone=True), nullable=False),
|
|
sa.Column('occurred_start', postgresql.TIMESTAMP(timezone=True), nullable=True),
|
|
sa.Column('occurred_end', postgresql.TIMESTAMP(timezone=True), nullable=True),
|
|
sa.Column('mentioned_at', postgresql.TIMESTAMP(timezone=True), nullable=True),
|
|
sa.Column('fact_type', sa.Text(), server_default='world', nullable=False),
|
|
sa.Column('confidence_score', sa.Float(), nullable=True),
|
|
sa.Column('access_count', sa.Integer(), server_default='0', nullable=False),
|
|
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.ForeignKeyConstraint(['document_id', 'bank_id'], ['documents.id', 'documents.bank_id'], name='memory_units_document_fkey', ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('id', name=op.f('pk_memory_units')),
|
|
sa.CheckConstraint("fact_type IN ('world', 'bank', 'opinion', 'observation')", name='memory_units_fact_type_check'),
|
|
sa.CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)", name='memory_units_confidence_range_check'),
|
|
sa.CheckConstraint(
|
|
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
|
|
"(fact_type = 'observation') OR "
|
|
"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
|
|
name='confidence_score_fact_type_check'
|
|
)
|
|
)
|
|
|
|
# Add search_vector column for full-text search
|
|
op.execute("""
|
|
ALTER TABLE memory_units
|
|
ADD COLUMN search_vector tsvector
|
|
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED
|
|
""")
|
|
|
|
op.create_index('idx_memory_units_bank_id', 'memory_units', ['bank_id'])
|
|
op.create_index('idx_memory_units_document_id', 'memory_units', ['document_id'])
|
|
op.create_index('idx_memory_units_event_date', 'memory_units', [sa.text('event_date DESC')])
|
|
op.create_index('idx_memory_units_bank_date', 'memory_units', ['bank_id', sa.text('event_date DESC')])
|
|
op.create_index('idx_memory_units_access_count', 'memory_units', [sa.text('access_count DESC')])
|
|
op.create_index('idx_memory_units_fact_type', 'memory_units', ['fact_type'])
|
|
op.create_index('idx_memory_units_bank_fact_type', 'memory_units', ['bank_id', 'fact_type'])
|
|
op.create_index('idx_memory_units_bank_type_date', 'memory_units', ['bank_id', 'fact_type', sa.text('event_date DESC')])
|
|
op.create_index('idx_memory_units_opinion_confidence', 'memory_units', ['bank_id', sa.text('confidence_score DESC')], postgresql_where=sa.text("fact_type = 'opinion'"))
|
|
op.create_index('idx_memory_units_opinion_date', 'memory_units', ['bank_id', sa.text('event_date DESC')], postgresql_where=sa.text("fact_type = 'opinion'"))
|
|
op.create_index('idx_memory_units_observation_date', 'memory_units', ['bank_id', sa.text('event_date DESC')], postgresql_where=sa.text("fact_type = 'observation'"))
|
|
op.create_index('idx_memory_units_embedding', 'memory_units', ['embedding'], postgresql_using='hnsw', postgresql_ops={'embedding': 'vector_cosine_ops'})
|
|
|
|
# Create BM25 full-text search index on search_vector
|
|
op.execute("""
|
|
CREATE INDEX idx_memory_units_text_search ON memory_units
|
|
USING gin(search_vector)
|
|
""")
|
|
|
|
op.execute("""
|
|
CREATE MATERIALIZED VIEW memory_units_bm25 AS
|
|
SELECT
|
|
id,
|
|
bank_id,
|
|
text,
|
|
to_tsvector('english', text) AS text_vector,
|
|
log(1.0 + length(text)::float / (SELECT avg(length(text)) FROM memory_units)) AS doc_length_factor
|
|
FROM memory_units
|
|
""")
|
|
|
|
op.create_index('idx_memory_units_bm25_bank', 'memory_units_bm25', ['bank_id'])
|
|
op.create_index('idx_memory_units_bm25_text_vector', 'memory_units_bm25', ['text_vector'], postgresql_using='gin')
|
|
|
|
# Create entity_cooccurrences table
|
|
op.create_table(
|
|
'entity_cooccurrences',
|
|
sa.Column('entity_id_1', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column('entity_id_2', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column('cooccurrence_count', sa.Integer(), server_default='1', nullable=False),
|
|
sa.Column('last_cooccurred', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.ForeignKeyConstraint(['entity_id_1'], ['entities.id'], name=op.f('fk_entity_cooccurrences_entity_id_1_entities'), ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['entity_id_2'], ['entities.id'], name=op.f('fk_entity_cooccurrences_entity_id_2_entities'), ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('entity_id_1', 'entity_id_2', name=op.f('pk_entity_cooccurrences')),
|
|
sa.CheckConstraint('entity_id_1 < entity_id_2', name='entity_cooccurrence_order_check')
|
|
)
|
|
op.create_index('idx_entity_cooccurrences_entity1', 'entity_cooccurrences', ['entity_id_1'])
|
|
op.create_index('idx_entity_cooccurrences_entity2', 'entity_cooccurrences', ['entity_id_2'])
|
|
op.create_index('idx_entity_cooccurrences_count', 'entity_cooccurrences', [sa.text('cooccurrence_count DESC')])
|
|
|
|
# Create memory_links table
|
|
op.create_table(
|
|
'memory_links',
|
|
sa.Column('from_unit_id', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column('to_unit_id', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column('link_type', sa.Text(), nullable=False),
|
|
sa.Column('entity_id', postgresql.UUID(as_uuid=True), nullable=True),
|
|
sa.Column('weight', sa.Float(), server_default='1.0', nullable=False),
|
|
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
|
|
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], name=op.f('fk_memory_links_entity_id_entities'), ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['from_unit_id'], ['memory_units.id'], name=op.f('fk_memory_links_from_unit_id_memory_units'), ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['to_unit_id'], ['memory_units.id'], name=op.f('fk_memory_links_to_unit_id_memory_units'), ondelete='CASCADE'),
|
|
sa.CheckConstraint("link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')", name='memory_links_link_type_check'),
|
|
sa.CheckConstraint('weight >= 0.0 AND weight <= 1.0', name='memory_links_weight_check')
|
|
)
|
|
# Create unique constraint using COALESCE for nullable entity_id
|
|
op.execute("CREATE UNIQUE INDEX idx_memory_links_unique ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid))")
|
|
op.create_index('idx_memory_links_from_unit', 'memory_links', ['from_unit_id'])
|
|
op.create_index('idx_memory_links_to_unit', 'memory_links', ['to_unit_id'])
|
|
op.create_index('idx_memory_links_entity', 'memory_links', ['entity_id'])
|
|
op.create_index('idx_memory_links_link_type', 'memory_links', ['link_type'])
|
|
|
|
# Create unit_entities table
|
|
op.create_table(
|
|
'unit_entities',
|
|
sa.Column('unit_id', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column('entity_id', postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], name=op.f('fk_unit_entities_entity_id_entities'), ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['unit_id'], ['memory_units.id'], name=op.f('fk_unit_entities_unit_id_memory_units'), ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('unit_id', 'entity_id', name=op.f('pk_unit_entities'))
|
|
)
|
|
op.create_index('idx_unit_entities_unit', 'unit_entities', ['unit_id'])
|
|
op.create_index('idx_unit_entities_entity', 'unit_entities', ['entity_id'])
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Downgrade schema - drop all tables."""
|
|
|
|
# Drop tables in reverse dependency order
|
|
op.drop_index('idx_unit_entities_entity', table_name='unit_entities')
|
|
op.drop_index('idx_unit_entities_unit', table_name='unit_entities')
|
|
op.drop_table('unit_entities')
|
|
|
|
op.drop_index('idx_memory_links_link_type', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_entity', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_to_unit', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_from_unit', table_name='memory_links')
|
|
op.execute('DROP INDEX IF EXISTS idx_memory_links_unique')
|
|
op.drop_table('memory_links')
|
|
|
|
op.drop_index('idx_entity_cooccurrences_count', table_name='entity_cooccurrences')
|
|
op.drop_index('idx_entity_cooccurrences_entity2', table_name='entity_cooccurrences')
|
|
op.drop_index('idx_entity_cooccurrences_entity1', table_name='entity_cooccurrences')
|
|
op.drop_table('entity_cooccurrences')
|
|
|
|
# Drop BM25 materialized view and index
|
|
op.drop_index('idx_memory_units_bm25_text_vector', table_name='memory_units_bm25')
|
|
op.drop_index('idx_memory_units_bm25_bank', table_name='memory_units_bm25')
|
|
op.execute('DROP MATERIALIZED VIEW IF EXISTS memory_units_bm25')
|
|
|
|
op.drop_index('idx_memory_units_embedding', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_observation_date', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_opinion_date', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_opinion_confidence', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_bank_type_date', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_bank_fact_type', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_fact_type', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_access_count', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_bank_date', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_event_date', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_document_id', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_bank_id', table_name='memory_units')
|
|
op.execute('DROP INDEX IF EXISTS idx_memory_units_text_search')
|
|
op.drop_table('memory_units')
|
|
|
|
op.execute('DROP INDEX IF EXISTS idx_entities_bank_lower_name')
|
|
op.drop_index('idx_entities_bank_name', table_name='entities')
|
|
op.drop_index('idx_entities_canonical_name', table_name='entities')
|
|
op.drop_index('idx_entities_bank_id', table_name='entities')
|
|
op.drop_table('entities')
|
|
|
|
op.drop_index('idx_async_operations_bank_status', table_name='async_operations')
|
|
op.drop_index('idx_async_operations_status', table_name='async_operations')
|
|
op.drop_index('idx_async_operations_bank_id', table_name='async_operations')
|
|
op.drop_table('async_operations')
|
|
|
|
op.drop_index('idx_documents_content_hash', table_name='documents')
|
|
op.drop_index('idx_documents_bank_id', table_name='documents')
|
|
op.drop_table('documents')
|
|
|
|
op.drop_table('banks')
|
|
|
|
# Drop extensions (optional - comment out if you want to keep them)
|
|
# op.execute('DROP EXTENSION IF EXISTS vector')
|
|
# op.execute('DROP EXTENSION IF EXISTS "uuid-ossp"')
|