118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
"""add messageembedding table
|
|
|
|
Revision ID: 917195d9b5e9
|
|
Revises: d429de0e5338
|
|
Create Date: 2024-01-01 12:00:00.000000
|
|
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from pgvector.sqlalchemy import Vector
|
|
|
|
from migrations.utils import index_exists, table_exists
|
|
from src.config import settings
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "917195d9b5e9"
|
|
down_revision: str | None = "d429de0e5338"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
schema = settings.DB.SCHEMA
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.create_table(
|
|
"message_embeddings",
|
|
sa.Column("id", sa.BigInteger(), sa.Identity(), nullable=False),
|
|
sa.Column("content", sa.Text(), nullable=False),
|
|
sa.Column("embedding", Vector(1536), nullable=False),
|
|
sa.Column("message_id", sa.Text(), nullable=False),
|
|
sa.Column("workspace_name", sa.Text(), nullable=False),
|
|
sa.Column("session_name", sa.Text(), nullable=True),
|
|
sa.Column("peer_name", sa.Text(), nullable=False),
|
|
sa.Column(
|
|
"created_at",
|
|
sa.DateTime(timezone=True),
|
|
nullable=False,
|
|
server_default=sa.func.now(),
|
|
),
|
|
# Foreign key constraints
|
|
sa.ForeignKeyConstraint(["message_id"], [f"{schema}.messages.public_id"]),
|
|
sa.ForeignKeyConstraint(["workspace_name"], [f"{schema}.workspaces.name"]),
|
|
sa.ForeignKeyConstraint(
|
|
["session_name", "workspace_name"],
|
|
[f"{schema}.sessions.name", f"{schema}.sessions.workspace_name"],
|
|
),
|
|
sa.ForeignKeyConstraint(
|
|
["peer_name", "workspace_name"],
|
|
[f"{schema}.peers.name", f"{schema}.peers.workspace_name"],
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
# Create indexes
|
|
op.create_index(
|
|
"idx_message_embeddings_message_id",
|
|
"message_embeddings",
|
|
["message_id"],
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"idx_message_embeddings_workspace_name",
|
|
"message_embeddings",
|
|
["workspace_name"],
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"idx_message_embeddings_session_name",
|
|
"message_embeddings",
|
|
["session_name"],
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"idx_message_embeddings_peer_name",
|
|
"message_embeddings",
|
|
["peer_name"],
|
|
schema=schema,
|
|
)
|
|
op.create_index(
|
|
"idx_message_embeddings_created_at",
|
|
"message_embeddings",
|
|
["created_at"],
|
|
schema=schema,
|
|
)
|
|
|
|
# Create HNSW index for vector similarity search
|
|
op.execute(f"""
|
|
CREATE INDEX idx_message_embeddings_embedding_hnsw
|
|
ON {schema}.message_embeddings
|
|
USING hnsw (embedding vector_cosine_ops)
|
|
WITH (m = 16, ef_construction = 64)
|
|
""")
|
|
|
|
|
|
def downgrade() -> None:
|
|
inspector = sa.inspect(op.get_bind())
|
|
if not table_exists("message_embeddings", inspector):
|
|
return
|
|
|
|
# Drop indexes defensively
|
|
indexes_to_drop = [
|
|
"idx_message_embeddings_embedding_hnsw",
|
|
"idx_message_embeddings_message_id",
|
|
"idx_message_embeddings_workspace_name",
|
|
"idx_message_embeddings_session_name",
|
|
"idx_message_embeddings_peer_name",
|
|
"idx_message_embeddings_created_at",
|
|
]
|
|
|
|
for index_name in indexes_to_drop:
|
|
if index_exists("message_embeddings", index_name, inspector):
|
|
op.drop_index(index_name, table_name="message_embeddings", schema=schema)
|
|
|
|
# Drop table (this will also drop foreign keys and check constraints)
|
|
op.drop_table("message_embeddings", schema=schema)
|