PMFlow/backend/alembic/versions/20260311_add_kb_chunk_ivffl...

43 lines
944 B
Python

"""add ivfflat index for knowledge chunk embedding
Revision ID: add_kb_chunk_ivfflat_index
Revises: add_pg_trgm_text_indexes
Create Date: 2026-03-11
"""
from __future__ import annotations
from alembic import op
revision = "add_kb_chunk_ivfflat_index"
down_revision = "add_pg_trgm_text_indexes"
branch_labels = None
depends_on = None
def _is_postgresql() -> bool:
bind = op.get_bind()
return str(getattr(getattr(bind, "dialect", None), "name", "")).lower().startswith("postgres")
def upgrade() -> None:
if not _is_postgresql():
return
op.execute(
"""
CREATE INDEX IF NOT EXISTS ix_knowledge_chunks_embedding_ivfflat
ON knowledge_base_chunks
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100)
"""
)
def downgrade() -> None:
if not _is_postgresql():
return
op.execute("DROP INDEX IF EXISTS ix_knowledge_chunks_embedding_ivfflat")