98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
"""add chunk_index to message_embeddings and make embeddings nullable
|
|
|
|
This migration:
|
|
1. Adds the chunk_index column to message_embeddings table for tracking
|
|
chunked message embeddings in external vector stores (turbopuffer/lancedb).
|
|
2. Makes embedding columns nullable in both message_embeddings and documents tables
|
|
since embeddings are now stored in external vector stores instead of PostgreSQL.
|
|
|
|
Revision ID: f1a2b3c4d5e6
|
|
Revises: baa22cad81e2
|
|
Create Date: 2025-11-24 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 column_exists, get_schema
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "f1a2b3c4d5e6"
|
|
down_revision: str | None = "baa22cad81e2"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
schema = get_schema()
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Add chunk_index column to message_embeddings and make embeddings nullable."""
|
|
inspector = sa.inspect(op.get_bind())
|
|
|
|
# Add chunk_index column to message_embeddings if it doesn't exist
|
|
# This is needed to track which chunk of a message this embedding represents
|
|
# Vector ID format: {message_public_id}_{chunk_index}
|
|
if not column_exists("message_embeddings", "chunk_index", inspector):
|
|
op.add_column(
|
|
"message_embeddings",
|
|
sa.Column(
|
|
"chunk_index",
|
|
sa.Integer(),
|
|
nullable=False,
|
|
server_default="0",
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
# Make message_embeddings.embedding nullable since embeddings are now stored
|
|
# in external vector stores (turbopuffer/lancedb) instead of PostgreSQL
|
|
op.alter_column(
|
|
"message_embeddings",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=True,
|
|
schema=schema,
|
|
)
|
|
|
|
# Make documents.embedding nullable for the same reason
|
|
# (this should already be nullable, but ensure it for consistency)
|
|
op.alter_column(
|
|
"documents",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=True,
|
|
schema=schema,
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Remove chunk_index column and revert embedding columns to non-nullable."""
|
|
inspector = sa.inspect(op.get_bind())
|
|
|
|
# Revert documents.embedding back to nullable=True (it was originally nullable=True)
|
|
op.alter_column(
|
|
"documents",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=True, # Keep as nullable since it was nullable in the original schema
|
|
schema=schema,
|
|
)
|
|
|
|
# Revert message_embeddings.embedding back to nullable=False
|
|
# Note: This may fail if there are NULL values in the database
|
|
op.alter_column(
|
|
"message_embeddings",
|
|
"embedding",
|
|
existing_type=Vector(1536),
|
|
nullable=False,
|
|
schema=schema,
|
|
)
|
|
|
|
# Remove chunk_index column if it exists
|
|
if column_exists("message_embeddings", "chunk_index", inspector):
|
|
op.drop_column("message_embeddings", "chunk_index", schema=schema)
|