196 lines
13 KiB
Python
196 lines
13 KiB
Python
"""Initial schema
|
|
|
|
Revision ID: af0413383b3e
|
|
Revises:
|
|
Create Date: 2025-11-03 14:31:53.245542
|
|
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
import pgvector.sqlalchemy
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = 'af0413383b3e'
|
|
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 pgvector extension
|
|
op.execute('CREATE EXTENSION IF NOT EXISTS vector')
|
|
op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
|
|
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.create_table('documents',
|
|
sa.Column('id', sa.Text(), nullable=False),
|
|
sa.Column('agent_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', 'agent_id')
|
|
)
|
|
op.create_index('idx_documents_agent_id', 'documents', ['agent_id'], unique=False)
|
|
op.create_index('idx_documents_content_hash', 'documents', ['content_hash'], unique=False)
|
|
op.create_table('entities',
|
|
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
|
sa.Column('canonical_name', sa.Text(), nullable=False),
|
|
sa.Column('entity_type', sa.Text(), nullable=False),
|
|
sa.Column('agent_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')
|
|
)
|
|
op.create_index('idx_entities_agent_id', 'entities', ['agent_id'], unique=False)
|
|
op.create_index('idx_entities_agent_name_type', 'entities', ['agent_id', 'canonical_name', 'entity_type'],
|
|
unique=False)
|
|
op.create_index('idx_entities_canonical_name', 'entities', ['canonical_name'], unique=False)
|
|
op.create_index('idx_entities_type', 'entities', ['entity_type'], unique=False)
|
|
op.create_table('entity_cooccurrences',
|
|
sa.Column('entity_id_1', sa.UUID(), nullable=False),
|
|
sa.Column('entity_id_2', sa.UUID(), 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.CheckConstraint('entity_id_1 < entity_id_2', name='entity_cooccurrence_order_check'),
|
|
sa.ForeignKeyConstraint(['entity_id_1'], ['entities.id'], ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['entity_id_2'], ['entities.id'], ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('entity_id_1', 'entity_id_2')
|
|
)
|
|
op.create_index('idx_entity_cooccurrences_count', 'entity_cooccurrences', ['cooccurrence_count'], unique=False,
|
|
postgresql_ops={'cooccurrence_count': 'DESC'})
|
|
op.create_index('idx_entity_cooccurrences_entity1', 'entity_cooccurrences', ['entity_id_1'], unique=False)
|
|
op.create_index('idx_entity_cooccurrences_entity2', 'entity_cooccurrences', ['entity_id_2'], unique=False)
|
|
op.create_table('memory_units',
|
|
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
|
|
sa.Column('agent_id', sa.Text(), nullable=False),
|
|
sa.Column('document_id', sa.Text(), nullable=True),
|
|
sa.Column('text', sa.Text(), nullable=False),
|
|
sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(dim=384), nullable=True),
|
|
sa.Column('context', sa.Text(), nullable=True),
|
|
sa.Column('event_date', postgresql.TIMESTAMP(timezone=True), nullable=False),
|
|
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('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.CheckConstraint(
|
|
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR (fact_type != 'opinion' AND confidence_score IS NULL)",
|
|
name='confidence_score_fact_type_check'),
|
|
sa.CheckConstraint("fact_type IN ('world', 'agent', 'opinion')"),
|
|
sa.CheckConstraint(
|
|
'confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)'),
|
|
sa.ForeignKeyConstraint(['document_id', 'agent_id'], ['documents.id', 'documents.agent_id'],
|
|
name='memory_units_document_fkey', ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('id')
|
|
)
|
|
op.create_index('idx_memory_units_access_count', 'memory_units', ['access_count'], unique=False,
|
|
postgresql_ops={'access_count': 'DESC'})
|
|
op.create_index('idx_memory_units_agent_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
|
|
postgresql_ops={'event_date': 'DESC'})
|
|
op.create_index('idx_memory_units_agent_fact_type', 'memory_units', ['agent_id', 'fact_type'], unique=False)
|
|
op.create_index('idx_memory_units_agent_id', 'memory_units', ['agent_id'], unique=False)
|
|
op.create_index('idx_memory_units_agent_type_date', 'memory_units', ['agent_id', 'fact_type', 'event_date'],
|
|
unique=False, postgresql_ops={'event_date': 'DESC'})
|
|
op.create_index('idx_memory_units_document_id', 'memory_units', ['document_id'], unique=False)
|
|
op.create_index('idx_memory_units_embedding', 'memory_units', ['embedding'], unique=False, postgresql_using='hnsw',
|
|
postgresql_ops={'embedding': 'vector_cosine_ops'})
|
|
op.create_index('idx_memory_units_event_date', 'memory_units', ['event_date'], unique=False,
|
|
postgresql_ops={'event_date': 'DESC'})
|
|
op.create_index('idx_memory_units_fact_type', 'memory_units', ['fact_type'], unique=False)
|
|
op.create_index('idx_memory_units_opinion_confidence', 'memory_units', ['agent_id', 'confidence_score'],
|
|
unique=False, postgresql_where=sa.text("fact_type = 'opinion'"),
|
|
postgresql_ops={'confidence_score': 'DESC'})
|
|
op.create_index('idx_memory_units_opinion_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
|
|
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
|
|
op.create_table('memory_links',
|
|
sa.Column('from_unit_id', sa.UUID(), nullable=False),
|
|
sa.Column('to_unit_id', sa.UUID(), nullable=False),
|
|
sa.Column('link_type', sa.Text(), nullable=False),
|
|
sa.Column('entity_id', sa.UUID(), nullable=False),
|
|
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'], ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['from_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['to_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('from_unit_id', 'to_unit_id', 'link_type', 'entity_id')
|
|
)
|
|
op.create_index('idx_memory_links_entity', 'memory_links', ['entity_id'], unique=False,
|
|
postgresql_where=sa.text('entity_id IS NOT NULL'))
|
|
op.create_index('idx_memory_links_from', 'memory_links', ['from_unit_id'], unique=False)
|
|
op.create_index('idx_memory_links_from_weight', 'memory_links', ['from_unit_id', 'weight'], unique=False,
|
|
postgresql_where=sa.text('weight >= 0.1'), postgresql_ops={'weight': 'DESC'})
|
|
op.create_index('idx_memory_links_to', 'memory_links', ['to_unit_id'], unique=False)
|
|
op.create_index('idx_memory_links_type', 'memory_links', ['link_type'], unique=False)
|
|
op.create_table('unit_entities',
|
|
sa.Column('unit_id', sa.UUID(), nullable=False),
|
|
sa.Column('entity_id', sa.UUID(), nullable=False),
|
|
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ondelete='CASCADE'),
|
|
sa.ForeignKeyConstraint(['unit_id'], ['memory_units.id'], ondelete='CASCADE'),
|
|
sa.PrimaryKeyConstraint('unit_id', 'entity_id')
|
|
)
|
|
op.create_index('idx_unit_entities_entity', 'unit_entities', ['entity_id'], unique=False)
|
|
op.create_index('idx_unit_entities_unit', 'unit_entities', ['unit_id'], unique=False)
|
|
# ### end Alembic commands ###
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Downgrade schema."""
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
op.drop_index('idx_unit_entities_unit', table_name='unit_entities')
|
|
op.drop_index('idx_unit_entities_entity', table_name='unit_entities')
|
|
op.drop_table('unit_entities')
|
|
op.drop_index('idx_memory_links_type', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_to', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_from_weight', table_name='memory_links', postgresql_where=sa.text('weight >= 0.1'),
|
|
postgresql_ops={'weight': 'DESC'})
|
|
op.drop_index('idx_memory_links_from', table_name='memory_links')
|
|
op.drop_index('idx_memory_links_entity', table_name='memory_links',
|
|
postgresql_where=sa.text('entity_id IS NOT NULL'))
|
|
op.drop_table('memory_links')
|
|
op.drop_index('idx_memory_units_opinion_date', table_name='memory_units',
|
|
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
|
|
op.drop_index('idx_memory_units_opinion_confidence', table_name='memory_units',
|
|
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'confidence_score': 'DESC'})
|
|
op.drop_index('idx_memory_units_fact_type', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_event_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
|
|
op.drop_index('idx_memory_units_embedding', table_name='memory_units', postgresql_using='hnsw',
|
|
postgresql_ops={'embedding': 'vector_cosine_ops'})
|
|
op.drop_index('idx_memory_units_document_id', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_agent_type_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
|
|
op.drop_index('idx_memory_units_agent_id', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_agent_fact_type', table_name='memory_units')
|
|
op.drop_index('idx_memory_units_agent_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
|
|
op.drop_index('idx_memory_units_access_count', table_name='memory_units', postgresql_ops={'access_count': 'DESC'})
|
|
op.drop_table('memory_units')
|
|
op.drop_index('idx_entity_cooccurrences_entity2', table_name='entity_cooccurrences')
|
|
op.drop_index('idx_entity_cooccurrences_entity1', table_name='entity_cooccurrences')
|
|
op.drop_index('idx_entity_cooccurrences_count', table_name='entity_cooccurrences',
|
|
postgresql_ops={'cooccurrence_count': 'DESC'})
|
|
op.drop_table('entity_cooccurrences')
|
|
op.drop_index('idx_entities_type', table_name='entities')
|
|
op.drop_index('idx_entities_canonical_name', table_name='entities')
|
|
op.drop_index('idx_entities_agent_name_type', table_name='entities')
|
|
op.drop_index('idx_entities_agent_id', table_name='entities')
|
|
op.drop_table('entities')
|
|
op.drop_index('idx_documents_content_hash', table_name='documents')
|
|
op.drop_index('idx_documents_agent_id', table_name='documents')
|
|
op.drop_table('documents')
|
|
# ### end Alembic commands ###
|