feat: persist isolated vector snapshots and exact retrieval
This commit is contained in:
@@ -0,0 +1,122 @@
|
||||
"""Persist dimension-checked vector snapshots with composite scope constraints."""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "0006_vectors"
|
||||
down_revision = "0005_chunk_indexes"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
class ExtensionVector(sa.types.UserDefinedType):
|
||||
"""Qualify extension type without adding public business tables to search_path."""
|
||||
|
||||
def get_col_spec(self, **kw):
|
||||
return "public.vector"
|
||||
|
||||
|
||||
def upgrade():
|
||||
schema = op.get_bind().scalar(
|
||||
sa.text("SELECT extnamespace::regnamespace::text FROM pg_extension WHERE extname='vector'")
|
||||
)
|
||||
if schema != "public":
|
||||
raise RuntimeError("vector 扩展须安装在 public;请由数据库管理员核对,迁移不会移动共享扩展")
|
||||
op.create_unique_constraint("uq_chunk_text_scope", "chunks", ["id", "index_version_id"])
|
||||
op.create_table(
|
||||
"embedding_spaces",
|
||||
sa.Column("id", sa.String(64), primary_key=True),
|
||||
sa.Column("dimension", sa.Integer(), nullable=False),
|
||||
sa.Column("spec", sa.JSON(), nullable=False),
|
||||
sa.UniqueConstraint("id", "dimension", name="uq_embedding_space_dimension"),
|
||||
sa.CheckConstraint("dimension BETWEEN 1 AND 2000", name="ck_embedding_space_dimension"),
|
||||
)
|
||||
op.create_table(
|
||||
"vector_versions",
|
||||
sa.Column("id", sa.String(36), primary_key=True),
|
||||
sa.Column(
|
||||
"text_index_version_id",
|
||||
sa.String(36),
|
||||
sa.ForeignKey("index_versions.id"),
|
||||
nullable=False,
|
||||
),
|
||||
sa.Column("space_id", sa.String(64), nullable=False),
|
||||
sa.Column("dimension", sa.Integer(), nullable=False),
|
||||
sa.Column("chunk_count", sa.Integer(), nullable=False),
|
||||
sa.Column("status", sa.String(20), nullable=False),
|
||||
sa.Column("created_by", sa.Integer(), sa.ForeignKey("users.id"), nullable=False),
|
||||
sa.Column(
|
||||
"created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
|
||||
),
|
||||
sa.Column("published_at", sa.DateTime(timezone=True)),
|
||||
sa.ForeignKeyConstraint(
|
||||
["space_id", "dimension"],
|
||||
["embedding_spaces.id", "embedding_spaces.dimension"],
|
||||
name="fk_vector_version_space",
|
||||
),
|
||||
sa.UniqueConstraint(
|
||||
"id", "text_index_version_id", "space_id", "dimension", name="uq_vector_version_scope"
|
||||
),
|
||||
sa.CheckConstraint(
|
||||
"status IN ('ready', 'published', 'superseded')", name="ck_vector_version_status"
|
||||
),
|
||||
sa.CheckConstraint("chunk_count BETWEEN 0 AND 2000", name="ck_vector_version_count"),
|
||||
)
|
||||
op.create_table(
|
||||
"chunk_vectors",
|
||||
sa.Column("version_id", sa.String(36), primary_key=True),
|
||||
sa.Column("chunk_id", sa.String(36), primary_key=True),
|
||||
sa.Column("text_index_version_id", sa.String(36), nullable=False),
|
||||
sa.Column("space_id", sa.String(64), nullable=False),
|
||||
sa.Column("dimension", sa.Integer(), nullable=False),
|
||||
sa.Column("embedding", ExtensionVector(), nullable=False),
|
||||
sa.ForeignKeyConstraint(
|
||||
["version_id", "text_index_version_id", "space_id", "dimension"],
|
||||
[
|
||||
"vector_versions.id",
|
||||
"vector_versions.text_index_version_id",
|
||||
"vector_versions.space_id",
|
||||
"vector_versions.dimension",
|
||||
],
|
||||
name="fk_chunk_vector_scope",
|
||||
),
|
||||
sa.ForeignKeyConstraint(
|
||||
["chunk_id", "text_index_version_id"],
|
||||
["chunks.id", "chunks.index_version_id"],
|
||||
name="fk_chunk_vector_text",
|
||||
),
|
||||
sa.CheckConstraint("public.vector_dims(embedding) = dimension", name="ck_vector_dimension"),
|
||||
sa.CheckConstraint(
|
||||
"public.vector_norm(embedding) BETWEEN 0.9999 AND 1.0001", name="ck_vector_unit_norm"
|
||||
),
|
||||
)
|
||||
op.create_table(
|
||||
"vector_state",
|
||||
sa.Column("id", sa.Integer(), primary_key=True),
|
||||
sa.Column("active_version_id", sa.String(36), sa.ForeignKey("vector_versions.id")),
|
||||
sa.Column("generation", sa.Integer(), nullable=False),
|
||||
sa.CheckConstraint("id = 1", name="ck_vector_state_singleton"),
|
||||
)
|
||||
op.execute(sa.text("INSERT INTO vector_state (id, generation) VALUES (1, 0)"))
|
||||
op.create_table(
|
||||
"vector_publications",
|
||||
sa.Column("generation", sa.Integer(), primary_key=True),
|
||||
sa.Column("previous_version_id", sa.String(36), sa.ForeignKey("vector_versions.id")),
|
||||
sa.Column("version_id", sa.String(36), sa.ForeignKey("vector_versions.id"), nullable=False),
|
||||
sa.Column("actor_id", sa.Integer(), sa.ForeignKey("users.id"), nullable=False),
|
||||
sa.Column(
|
||||
"published_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
for table in [
|
||||
"vector_publications",
|
||||
"vector_state",
|
||||
"chunk_vectors",
|
||||
"vector_versions",
|
||||
"embedding_spaces",
|
||||
]:
|
||||
op.drop_table(table)
|
||||
op.drop_constraint("uq_chunk_text_scope", "chunks", type_="unique")
|
||||
Reference in New Issue
Block a user