RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: remove destructive embedding migration

This commit is contained in:
Rajat Ahuja 2025-12-04 14:05:00 -05:00
parent ca218da615
commit 0dd3ee4e45
4 changed files with 80 additions and 136 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py Normal file
View File

@ -0,0 +1,54 @@
"""add chunk_index to message_embeddings
This migration adds the chunk_index column to message_embeddings table for tracking
chunked message embeddings in external vector stores (turbopuffer/lancedb).
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 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 for tracking chunked message embeddings."""
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,
)
def downgrade() -> None:
"""Remove chunk_index column from message_embeddings."""
inspector = sa.inspect(op.get_bind())
# Remove chunk_index column if it exists
if column_exists("message_embeddings", "chunk_index", inspector):
op.drop_column("message_embeddings", "chunk_index", schema=schema)

136
migrations/versions/f1a2b3c4d5e6_remove_embedding_columns_for_vector_store.py
View File

@ -1,136 +0,0 @@
"""remove embedding columns for vector store migration
This migration removes the embedding columns from the message_embeddings and documents
tables as part of the migration from pgvector to external vector stores (turbopuffer/lancedb).
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, index_exists
# 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:
"""Remove embedding columns and HNSW indexes from message_embeddings and documents tables.
Also add chunk_index column to message_embeddings for tracking chunked message embeddings.
"""
inspector = sa.inspect(op.get_bind())
# === message_embeddings table ===
# Drop HNSW index on message_embeddings.embedding if it exists
# Check for both possible index names (old naming vs new naming convention)
for index_name in [
"ix_message_embeddings_embedding_hnsw",
"idx_message_embeddings_embedding_hnsw",
]:
if index_exists("message_embeddings", index_name, inspector):
op.drop_index(index_name, table_name="message_embeddings", schema=schema)
# Drop embedding column from message_embeddings if it exists
if column_exists("message_embeddings", "embedding", inspector):
op.drop_column("message_embeddings", "embedding", schema=schema)
# 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,
)
# === documents table ===
# Drop HNSW index on documents.embedding if it exists
# Check for both possible index names (old naming vs new naming convention)
for index_name in [
"ix_documents_embedding_hnsw",
"idx_documents_embedding_hnsw",
]:
if index_exists("documents", index_name, inspector):
op.drop_index(index_name, table_name="documents", schema=schema)
# Drop embedding column from documents if it exists
if column_exists("documents", "embedding", inspector):
op.drop_column("documents", "embedding", schema=schema)
def downgrade() -> None:
"""Restore embedding columns and HNSW indexes, remove chunk_index.
Note: This downgrade will create empty embedding columns. The actual embeddings
would need to be restored from a backup or regenerated if rolling back this migration.
"""
inspector = sa.inspect(op.get_bind())
# === documents table ===
# Add embedding column back to documents if it doesn't exist
if not column_exists("documents", "embedding", inspector):
op.add_column(
"documents",
sa.Column("embedding", Vector(1536), nullable=True),
schema=schema,
)
# Recreate HNSW index on documents.embedding
if not index_exists("documents", "ix_documents_embedding_hnsw", inspector):
op.execute(
f"""
CREATE INDEX ix_documents_embedding_hnsw
ON {schema}.documents
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64)
"""
)
# === message_embeddings table ===
# Remove chunk_index column if it exists
if column_exists("message_embeddings", "chunk_index", inspector):
op.drop_column("message_embeddings", "chunk_index", schema=schema)
# Add embedding column back to message_embeddings if it doesn't exist
if not column_exists("message_embeddings", "embedding", inspector):
op.add_column(
"message_embeddings",
sa.Column("embedding", Vector(1536), nullable=True),
schema=schema,
)
# Recreate HNSW index on message_embeddings.embedding
if not index_exists(
"message_embeddings", "ix_message_embeddings_embedding_hnsw", inspector
):
op.execute(
f"""
CREATE INDEX ix_message_embeddings_embedding_hnsw
ON {schema}.message_embeddings
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64)
"""
)

2
tests/alembic/revisions/__init__.py
View File

@ -21,6 +21,7 @@ from . import (
test_d429de0e5338_adopt_peer_paradigm,
test_e9b705f9adf9_add_server_defaults_to_timestamp_,
test_ec8f94139b02_codify_workspace_name_and_message_id_in_,
test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings,
)
__all__ = [
@ -44,4 +45,5 @@ __all__ = [
"test_d429de0e5338_adopt_peer_paradigm",
"test_e9b705f9adf9_add_server_defaults_to_timestamp_",
"test_ec8f94139b02_codify_workspace_name_and_message_id_in_",
"test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings",
]

24
tests/alembic/revisions/test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py Normal file
View File

@ -0,0 +1,24 @@
"""Hooks for revision f1a2b3c4d5e6 (add_chunk_index_to_message_embeddings)."""
from __future__ import annotations
from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier
@register_before_upgrade("f1a2b3c4d5e6")
def prepare_add_chunk_index_to_message_embeddings(
verifier: MigrationVerifier,
) -> None:
"""Seed state and assertions before upgrading to f1a2b3c4d5e6."""
# Verify chunk_index column doesn't exist before migration
verifier.assert_column_exists("message_embeddings", "chunk_index", exists=False)
@register_after_upgrade("f1a2b3c4d5e6")
def verify_add_chunk_index_to_message_embeddings(
verifier: MigrationVerifier,
) -> None:
"""Add assertions validating the effects of f1a2b3c4d5e6."""
# Verify chunk_index column was added with correct properties
verifier.assert_column_exists("message_embeddings", "chunk_index", nullable=False)