476 lines
17 KiB
Python
476 lines
17 KiB
Python
"""Configure pgvector schema dim to match EMBEDDING_VECTOR_DIMENSIONS.
|
|
|
|
Usage::
|
|
|
|
uv run python scripts/configure_embeddings.py # interactive
|
|
uv run python scripts/configure_embeddings.py --dry-run # print intent, no DB write
|
|
uv run python scripts/configure_embeddings.py --yes # apply without prompt
|
|
uv run python scripts/configure_embeddings.py --report # full external-store inventory
|
|
|
|
The bootstrap sequence for a self-hosted install is:
|
|
|
|
1. alembic upgrade head # creates default vector(1536) schema
|
|
2. uv run python scripts/configure_embeddings.py # ALTER columns to target dim
|
|
3. start the API and deriver # validators refuse to start on mismatch
|
|
|
|
Existing 1536 deployments need no action — step 2 is a no-op when settings
|
|
already match the schema.
|
|
|
|
This script never creates or modifies external-store namespaces. Turbopuffer
|
|
and LanceDB namespaces are per-workspace and lazy-created on first write by
|
|
application code; their dim is implicitly pinned at that point. Use
|
|
``--report`` to enumerate existing namespaces against the configured dim.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
from dataclasses import dataclass
|
|
|
|
# Match the path-shim convention used by the other scripts in this directory
|
|
# so `src.*` imports resolve when the script is run directly.
|
|
_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
if _PROJECT_ROOT not in sys.path:
|
|
sys.path.insert(0, _PROJECT_ROOT)
|
|
|
|
from sqlalchemy import select, text # noqa: E402
|
|
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine # noqa: E402
|
|
|
|
from src.config import settings # noqa: E402
|
|
from src.db import engine # noqa: E402
|
|
from src.models import Collection, Workspace # noqa: E402
|
|
from src.vector_store import VectorStore # noqa: E402
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_EMBEDDING_TABLES: tuple[str, ...] = ("documents", "message_embeddings")
|
|
|
|
# Defense-in-depth for the dynamic SQL paths in this script. The values we
|
|
# interpolate (schema from settings, index names from pg_indexes, table
|
|
# names from a hardcoded constant) are not user input under any threat
|
|
# model we currently care about, but validating once at the top of the
|
|
# pipeline keeps the constraint explicit and the surface tight.
|
|
_SAFE_IDENT_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
|
|
|
|
|
def _validate_identifier(name: str, *, kind: str) -> None:
|
|
if not _SAFE_IDENT_PATTERN.fullmatch(name):
|
|
raise SystemExit(
|
|
f"error: refusing to interpolate {kind} {name!r} into SQL —"
|
|
+ " expected a SQL identifier of the form [A-Za-z_][A-Za-z0-9_]*."
|
|
+ " Reconfigure your environment and re-run."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Result types
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _PgvectorPlan:
|
|
target_dim: int
|
|
schema: str
|
|
current_dims: dict[str, int]
|
|
needs_alter: bool
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _NamespaceRecord:
|
|
"""A single row in the --report output."""
|
|
|
|
namespace: str
|
|
status: str # one of: "ok", "missing", "mismatch"
|
|
actual_dim: int | None
|
|
target_dim: int
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# pgvector phase
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _introspect_pgvector(conn: AsyncConnection, schema: str) -> dict[str, int]:
|
|
"""Return ``{table_name: atttypmod}`` for embedding columns in ``schema``.
|
|
|
|
Tables not present in the result dict are absent from the schema.
|
|
pgvector stores the declared dim directly in ``atttypmod`` (no VARHDRSZ).
|
|
"""
|
|
query = text(
|
|
"""
|
|
SELECT c.relname AS table_name, a.atttypmod AS typmod
|
|
FROM pg_attribute a
|
|
JOIN pg_class c ON a.attrelid = c.oid
|
|
JOIN pg_namespace n ON c.relnamespace = n.oid
|
|
WHERE n.nspname = :schema
|
|
AND c.relname = ANY(:tables)
|
|
AND a.attname = 'embedding'
|
|
"""
|
|
)
|
|
result = await conn.execute(
|
|
query,
|
|
{"schema": schema, "tables": list(_EMBEDDING_TABLES)},
|
|
)
|
|
return {row.table_name: row.typmod for row in result}
|
|
|
|
|
|
async def _build_pgvector_plan(
|
|
engine: AsyncEngine, target_dim: int, schema: str
|
|
) -> _PgvectorPlan:
|
|
"""Build a plan describing what (if anything) the script will change."""
|
|
async with engine.connect() as conn:
|
|
current = await _introspect_pgvector(conn, schema)
|
|
|
|
missing = set(_EMBEDDING_TABLES) - current.keys()
|
|
if missing:
|
|
listing = ", ".join(sorted(f"{schema}.{t}.embedding" for t in missing))
|
|
raise SystemExit(
|
|
f"error: required vector columns missing: {listing}."
|
|
+ " Run `alembic upgrade head` first."
|
|
)
|
|
for table, typmod in current.items():
|
|
if typmod == -1:
|
|
raise SystemExit(
|
|
f"error: {schema}.{table}.embedding has no declared vector"
|
|
+ " dimension (unbounded typmod). Drop and recreate the column"
|
|
+ " or restore from a versioned migration before re-running."
|
|
)
|
|
|
|
needs_alter = any(typmod != target_dim for typmod in current.values())
|
|
return _PgvectorPlan(
|
|
target_dim=target_dim,
|
|
schema=schema,
|
|
current_dims=current,
|
|
needs_alter=needs_alter,
|
|
)
|
|
|
|
|
|
async def _count_non_null_embeddings(
|
|
conn: AsyncConnection, schema: str, table: str
|
|
) -> int:
|
|
query = text(
|
|
f'SELECT COUNT(*) AS n FROM "{schema}"."{table}" WHERE embedding IS NOT NULL'
|
|
)
|
|
result = await conn.execute(query)
|
|
row = result.first()
|
|
return int(row.n) if row is not None else 0
|
|
|
|
|
|
async def _fetch_hnsw_index_defs(
|
|
conn: AsyncConnection, schema: str
|
|
) -> list[tuple[str, str]]:
|
|
"""Return ``(index_name, CREATE INDEX ...)`` for HNSW indices on the
|
|
embedding columns. We re-CREATE them after the ALTER using these exact
|
|
definitions, preserving operator-set params (m, ef_construction, etc.)."""
|
|
query = text(
|
|
"""
|
|
SELECT indexname AS name, indexdef AS ddl
|
|
FROM pg_indexes
|
|
WHERE schemaname = :schema
|
|
AND tablename = ANY(:tables)
|
|
AND indexdef ILIKE '%USING hnsw%'
|
|
"""
|
|
)
|
|
result = await conn.execute(
|
|
query, {"schema": schema, "tables": list(_EMBEDDING_TABLES)}
|
|
)
|
|
return [(row.name, row.ddl) for row in result]
|
|
|
|
|
|
async def _apply_pgvector_alter(engine: AsyncEngine, plan: _PgvectorPlan) -> None:
|
|
"""ALTER the embedding columns to ``plan.target_dim`` in a single
|
|
transaction. Refuses to proceed if any non-null embeddings exist.
|
|
|
|
Sequence (inside the transaction):
|
|
1. LOCK TABLE ... IN ACCESS EXCLUSIVE MODE — closes the TOCTOU window
|
|
between the population check and the ALTER.
|
|
2. SELECT COUNT(embedding IS NOT NULL) per table — refuse if any > 0.
|
|
3. Save HNSW index definitions, then DROP them (cannot ALTER under HNSW).
|
|
4. ALTER ... ALTER COLUMN embedding TYPE vector(N) USING NULL.
|
|
5. Recreate HNSW indices from saved definitions.
|
|
"""
|
|
async with engine.begin() as conn:
|
|
# Step 1: lock both tables for the duration of the transaction.
|
|
for table in _EMBEDDING_TABLES:
|
|
await conn.execute(
|
|
text(f'LOCK TABLE "{plan.schema}"."{table}" IN ACCESS EXCLUSIVE MODE')
|
|
)
|
|
|
|
# Step 2: population check.
|
|
counts: dict[str, int] = {}
|
|
for table in _EMBEDDING_TABLES:
|
|
counts[table] = await _count_non_null_embeddings(conn, plan.schema, table)
|
|
populated = {t: n for t, n in counts.items() if n > 0}
|
|
if populated:
|
|
detail = ", ".join(f"{t}: {n} rows" for t, n in sorted(populated.items()))
|
|
raise SystemExit(
|
|
f"error: refusing to ALTER populated embedding tables ({detail})."
|
|
+ " This script only configures empty tables. Re-embed out-of-band"
|
|
+ " into a fresh deployment, then cut over."
|
|
)
|
|
|
|
# Step 3: snapshot + drop HNSW indices.
|
|
index_defs = await _fetch_hnsw_index_defs(conn, plan.schema)
|
|
for index_name, _ddl in index_defs:
|
|
_validate_identifier(index_name, kind="HNSW index name")
|
|
logger.info("dropping HNSW index %s", index_name)
|
|
await conn.execute(text(f'DROP INDEX "{plan.schema}"."{index_name}"'))
|
|
|
|
# Step 4: ALTER columns.
|
|
for table in _EMBEDDING_TABLES:
|
|
logger.info(
|
|
"altering %s.%s.embedding to vector(%d)",
|
|
plan.schema,
|
|
table,
|
|
plan.target_dim,
|
|
)
|
|
await conn.execute(
|
|
text(
|
|
f'ALTER TABLE "{plan.schema}"."{table}"'
|
|
+ f" ALTER COLUMN embedding TYPE vector({plan.target_dim})"
|
|
+ " USING NULL"
|
|
)
|
|
)
|
|
|
|
# Step 5: recreate HNSW indices from the saved definitions.
|
|
for index_name, ddl in index_defs:
|
|
logger.info("recreating HNSW index %s", index_name)
|
|
await conn.execute(text(ddl))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# External-store report
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _enumerate_workspaces(conn: AsyncConnection) -> list[str]:
|
|
"""All workspace names, ordered by creation. Uses the ORM so
|
|
``Base.metadata.schema`` (configured from ``DB.SCHEMA``) is honored —
|
|
non-public schema deployments must not sample the wrong table."""
|
|
stmt = select(Workspace.name).order_by(Workspace.created_at)
|
|
result = await conn.execute(stmt)
|
|
return [row[0] for row in result]
|
|
|
|
|
|
async def _enumerate_collections(
|
|
conn: AsyncConnection,
|
|
) -> list[tuple[str, str, str]]:
|
|
"""Every (workspace_name, observer, observed) triple that has a row in
|
|
the collections table — these are the document namespaces that could
|
|
exist in an external store."""
|
|
stmt = select(Collection.workspace_name, Collection.observer, Collection.observed)
|
|
result = await conn.execute(stmt)
|
|
return [(row[0], row[1], row[2]) for row in result]
|
|
|
|
|
|
async def _build_external_namespace_inventory(
|
|
engine: AsyncEngine,
|
|
) -> list[tuple[str, str]]:
|
|
"""Return ``(namespace_type, namespace_name)`` pairs for every namespace
|
|
that should exist based on the application DB. Message namespaces are
|
|
derived per workspace, document namespaces per collection row.
|
|
"""
|
|
from src.vector_store import get_external_vector_store
|
|
|
|
store = get_external_vector_store()
|
|
if store is None:
|
|
return []
|
|
|
|
async with engine.connect() as conn:
|
|
workspace_names = await _enumerate_workspaces(conn)
|
|
collection_keys = await _enumerate_collections(conn)
|
|
|
|
pairs: list[tuple[str, str]] = []
|
|
for workspace_name in workspace_names:
|
|
pairs.append(("message", store.get_vector_namespace("message", workspace_name)))
|
|
for workspace_name, observer, observed in collection_keys:
|
|
pairs.append(
|
|
(
|
|
"document",
|
|
store.get_vector_namespace(
|
|
"document", workspace_name, observer=observer, observed=observed
|
|
),
|
|
)
|
|
)
|
|
return pairs
|
|
|
|
|
|
async def _probe_namespace_dim(store: VectorStore, namespace: str) -> int | None:
|
|
"""Return the namespace's declared dim, or ``None`` if the namespace
|
|
does not exist yet. Delegates to the store-specific probe."""
|
|
return await store.probe_namespace_dim(namespace)
|
|
|
|
|
|
async def _emit_report(
|
|
engine: AsyncEngine, target_dim: int, *, is_report_mode: bool
|
|
) -> int:
|
|
"""Print the per-namespace inventory and return an exit code. 0 on a
|
|
clean report (all matching or missing); non-zero on any mismatch.
|
|
|
|
``is_report_mode=True`` means the operator explicitly invoked ``--report``
|
|
— only then do we print the "no effect with pgvector" notice. Implicit
|
|
post-apply calls from interactive/dry-run/yes mode stay silent when the
|
|
deployment is on pgvector.
|
|
"""
|
|
if settings.VECTOR_STORE.TYPE == "pgvector":
|
|
if is_report_mode:
|
|
print("--report has no effect with VECTOR_STORE_TYPE=pgvector")
|
|
return 0
|
|
|
|
inventory = await _build_external_namespace_inventory(engine)
|
|
if not inventory:
|
|
print(
|
|
"no external namespaces to inventory"
|
|
+ " (no workspaces/collections exist yet, or no external store configured)"
|
|
)
|
|
return 0
|
|
|
|
from src.vector_store import get_external_vector_store
|
|
|
|
store = get_external_vector_store()
|
|
if store is None:
|
|
print("no external store configured; nothing to report")
|
|
return 0
|
|
|
|
records: list[_NamespaceRecord] = []
|
|
for _ns_type, namespace in inventory:
|
|
actual = await _probe_namespace_dim(store, namespace)
|
|
if actual is None:
|
|
# Namespace has not been written to yet (lazy-create model).
|
|
status = "missing"
|
|
elif actual == target_dim:
|
|
status = "ok"
|
|
else:
|
|
status = "mismatch"
|
|
records.append(
|
|
_NamespaceRecord(
|
|
namespace=namespace,
|
|
status=status,
|
|
actual_dim=actual,
|
|
target_dim=target_dim,
|
|
)
|
|
)
|
|
|
|
width = max(len(r.namespace) for r in records)
|
|
print(f"{'namespace'.ljust(width)} status dim")
|
|
print(f"{'-' * width} --------- ------")
|
|
for r in records:
|
|
dim_str = "?" if r.actual_dim is None else str(r.actual_dim)
|
|
print(f"{r.namespace.ljust(width)} {r.status:<9} {dim_str}")
|
|
|
|
mismatches = [r for r in records if r.status == "mismatch"]
|
|
if mismatches:
|
|
print(
|
|
f"\nerror: {len(mismatches)} namespace(s) have dim != {target_dim}",
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
prog="configure_embeddings",
|
|
description=(
|
|
"Configure pgvector schema dim to match EMBEDDING_VECTOR_DIMENSIONS."
|
|
),
|
|
)
|
|
mode = parser.add_mutually_exclusive_group()
|
|
mode.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="print intended changes and exit without touching the DB",
|
|
)
|
|
mode.add_argument(
|
|
"--yes",
|
|
action="store_true",
|
|
help="apply changes without an interactive prompt",
|
|
)
|
|
mode.add_argument(
|
|
"--report",
|
|
action="store_true",
|
|
help="print external-store namespace inventory and exit",
|
|
)
|
|
return parser
|
|
|
|
|
|
def _confirm(prompt: str) -> bool:
|
|
response = input(f"{prompt} [y/N]: ").strip().lower()
|
|
return response in {"y", "yes"}
|
|
|
|
|
|
async def _async_main(args: argparse.Namespace) -> int:
|
|
try:
|
|
return await _run_pipeline(args)
|
|
finally:
|
|
# Dispose inside the same event loop so cleanup doesn't spin up a
|
|
# second loop just to await engine.dispose().
|
|
await engine.dispose()
|
|
|
|
|
|
async def _run_pipeline(args: argparse.Namespace) -> int:
|
|
target_dim = settings.EMBEDDING.VECTOR_DIMENSIONS
|
|
schema = settings.DB.SCHEMA
|
|
_validate_identifier(schema, kind="DB.SCHEMA")
|
|
|
|
if args.report:
|
|
return await _emit_report(engine, target_dim, is_report_mode=True)
|
|
|
|
plan = await _build_pgvector_plan(engine, target_dim, schema)
|
|
if not plan.needs_alter:
|
|
print(
|
|
f"pgvector: {schema}.documents.embedding and"
|
|
+ f" {schema}.message_embeddings.embedding already at dim {target_dim},"
|
|
+ " skipping ALTER"
|
|
)
|
|
return await _emit_report(engine, target_dim, is_report_mode=False)
|
|
|
|
current_summary = ", ".join(
|
|
f"{schema}.{t}.embedding={plan.current_dims[t]}" for t in _EMBEDDING_TABLES
|
|
)
|
|
print(f"target dim: {target_dim}")
|
|
print(f"current: {current_summary}")
|
|
print("planned operations (single transaction):")
|
|
print(f" - LOCK TABLE {schema}.documents IN ACCESS EXCLUSIVE MODE")
|
|
print(f" - LOCK TABLE {schema}.message_embeddings IN ACCESS EXCLUSIVE MODE")
|
|
print(" - refuse if any non-null embeddings exist")
|
|
print(" - DROP existing HNSW indices on the embedding columns")
|
|
print(
|
|
f" - ALTER COLUMN embedding TYPE vector({target_dim}) USING NULL"
|
|
+ " on both tables"
|
|
)
|
|
print(" - CREATE HNSW indices from snapshotted definitions")
|
|
|
|
if args.dry_run:
|
|
print("\n--dry-run: no changes applied")
|
|
return 0
|
|
|
|
if not args.yes and not _confirm("apply?"):
|
|
print("aborted")
|
|
return 1
|
|
|
|
await _apply_pgvector_alter(engine, plan)
|
|
print(f"\npgvector schema is now at dim {target_dim}")
|
|
return await _emit_report(engine, target_dim, is_report_mode=False)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
parser = _build_parser()
|
|
args = parser.parse_args(argv)
|
|
return asyncio.run(_async_main(args))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|