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

feat(scripts): add configure_embeddings bootstrap CLI 

Adds scripts/configure_embeddings.py alongside the other one-off scripts
  (provision_db, migrate_db, generate_jwt_secret, etc.). Invoked as
  `uv run python scripts/configure_embeddings.py` — same convention as the
  existing scripts in that directory, including the sys.path shim that
  lets src.* imports resolve when run directly.

  Bootstrap step for self-hosted installs at a non-default
  EMBEDDING_VECTOR_DIMENSIONS — runs between `alembic upgrade head` and
  starting the API/deriver.

  pgvector ALTER safety (single transaction):
  - LOCK TABLE {schema}.documents, {schema}.message_embeddings IN ACCESS
    EXCLUSIVE MODE — closes the TOCTOU window between population check
    and ALTER
  - COUNT(*) WHERE embedding IS NOT NULL on both tables; refuse with a
    non-zero exit if either is populated (ALTER ... USING NULL would
    silently wipe those vectors)
  - Snapshot HNSW index DDL from pg_indexes; drop, ALTER, recreate from
    the captured DDL so operator-set HNSW params (m, ef_construction)
    survive the round trip

  External vector stores (turbopuffer, lancedb) are never created or
  modified — namespaces are per-workspace and lazy-created on first write.
  The --report mode enumerates workspaces and collections from the
  application DB, derives the expected namespaces via
  get_vector_namespace(), and prints a per-namespace status table.

  CLI modes (mutually exclusive):
  - (default) interactive: print plan, prompt to confirm
  - --dry-run: print plan and exit 0 without touching the DB
  - --yes: apply without prompt
  - --report: print external-store namespace inventory and exit

  Also updates src/startup/embedding_validator.py error-message paths and
  docs/v3/contributing/configuration.mdx invocations to point at the new
  script location.

  Tests cover plan no-op, plan needs-alter, plan raises on missing column,
  ALTER + HNSW round-trip, refuse-when-populated (monkeypatched count to
  avoid wiring the full workspace/peer/collection/document FK chain just
  to land one vector row), and idempotency.
This commit is contained in:
Vineeth Voruganti 2026-05-12 18:08:28 -04:00
parent 405f625595
commit 130b49884b
5 changed files with 662 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
docs/v3/contributing/configuration.mdx
View File

@ -218,8 +218,48 @@ EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=http://localhost:8000/v1
EMBEDDING_MODEL_CONFIG__OVERRIDES__API_KEY_ENV=EMBEDDING_CUSTOM_API_KEY
```
Current constraint:
- `EMBEDDING_VECTOR_DIMENSIONS` can be changed for fully migrated external vector stores, but pgvector and dual-write mode still require `1536` until the schema migration lands.
Forwarding `dimensions=` to OpenAI-compatible providers is controlled by `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE`:
- `auto` (default): forwards `dimensions=` when you have explicitly set `EMBEDDING_VECTOR_DIMENSIONS` to a non-default value, and the configured model is not on the known-rejecting list (currently `text-embedding-ada-002`). Deployments that stay on the default keep their existing behavior.
- `always`: always forward. Use for OpenAI-compatible self-hosted providers that require it, and for same-as-default truncation (e.g. `text-embedding-3-large` truncated to 1536 — `auto` would skip the parameter because 1536 is the default).
- `never`: never forward. Explicit opt-out for providers that reject the parameter.
#### Bootstrapping non-default dimensions
`EMBEDDING_VECTOR_DIMENSIONS` is treated as immutable for the life of a deployment. The pgvector schema is dim-pinned by Alembic at `1536` by default; if you want a different dim, you must ALTER the empty columns once at bootstrap time.
Install order for a non-default dim:
```bash
# 1. Apply migrations (creates default vector(1536) schema)
alembic upgrade head
# 2. Set the dim you want
export EMBEDDING_VECTOR_DIMENSIONS=768
# 3. ALTER the empty columns to the target dim
uv run python scripts/configure_embeddings.py --dry-run # preview
uv run python scripts/configure_embeddings.py --yes # apply
# 4. Start API and deriver — both run the startup validator and refuse
# to serve traffic if the schema and EMBEDDING_VECTOR_DIMENSIONS disagree.
```
Existing deployments at 1536 with `text-embedding-3-small` need no action — step 3 detects matching dims and skips.
The script refuses to ALTER tables that already contain non-null embeddings. To switch dim or model on a populated deployment, stand up a new deployment at the new configuration and migrate data out of band; there is no in-place re-embedding affordance.
External vector stores (Turbopuffer, LanceDB) do not need bootstrap setup. Namespaces are per-workspace and lazy-created on first write at whatever dim the embedding client returns. Use `--report` to inventory the existing namespaces against the configured dim:
```bash
uv run python scripts/configure_embeddings.py --report
```
The startup validator at `src/startup/embedding_validator.py` enforces the dim invariant at boot for both the API (`src/main.py` lifespan) and the deriver (`src/deriver/__main__.py`). A mismatch crashes the process with an actionable error before any HTTP route is served or any queue task is processed.
`VECTOR_STORE_DIMENSIONS` is **deprecated**. `EMBEDDING_VECTOR_DIMENSIONS` is the single source of truth; setting `VECTOR_STORE_DIMENSIONS` explicitly emits a startup warning and is otherwise ignored. The field will be removed in a future release; drop it from your `.env` to silence the warning.
The `VECTOR_STORE_MIGRATED` flag still exists and still controls dual-write / cutover semantics for legacy tenants moving between storage backends (pgvector ↔ turbopuffer ↔ lancedb). It is unrelated to dimension configuration after this release.
### Feature-Specific Model Configuration

438
scripts/configure_embeddings.py Normal file
View File

@ -0,0 +1,438 @@
"""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 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 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
logger = logging.getLogger(__name__)
_EMBEDDING_TABLES: tuple[str, ...] = ("documents", "message_embeddings")
# ---------------------------------------------------------------------------
# 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", "unknown"
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:
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]:
result = await conn.execute(text("SELECT name FROM workspaces ORDER BY created_at"))
return [row.name for row in result]
async def _enumerate_collections(
conn: AsyncConnection,
) -> list[tuple[str, str, str]]:
result = await conn.execute(
text("SELECT workspace_name, observer, observed FROM collections")
)
return [(row.workspace_name, row.observer, row.observed) 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: object, namespace: str) -> int | None:
"""Best-effort: return the namespace's declared dim if introspectable.
Returns ``None`` if the SDK does not expose a uniform dim accessor for
the configured store. Future work can specialize per store; today the
pgvector validator is the load-bearing dim safety.
"""
_ = (store, namespace)
return None
async def _emit_report(engine: AsyncEngine, target_dim: int) -> 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.
"""
if settings.VECTOR_STORE.TYPE == "pgvector":
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()
records: list[_NamespaceRecord] = []
for _ns_type, namespace in inventory:
actual = await _probe_namespace_dim(store, namespace) if store else None
if actual is None:
status = "unknown"
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:
target_dim = settings.EMBEDDING.VECTOR_DIMENSIONS
schema = settings.DB.SCHEMA
if args.report:
return await _emit_report(engine, target_dim)
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)
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)
def main(argv: list[str] | None = None) -> int:
logging.basicConfig(level=logging.INFO, format="%(message)s")
parser = _build_parser()
args = parser.parse_args(argv)
try:
return asyncio.run(_async_main(args))
finally:
# Best-effort cleanup; the script is short-lived so a leak here is
# harmless but keeping the dispose explicit avoids "engine not
# disposed" warnings during tests.
asyncio.run(engine.dispose())
if __name__ == "__main__":
raise SystemExit(main())

9
src/startup/embedding_validator.py
View File

@ -8,8 +8,7 @@ actual state instead of operator-asserted state.
For external stores (turbopuffer, lancedb) the check is best-effort: namespaces
are per-workspace and lazy-created, so this validator can only sample existing
ones. Full enumeration is available via `uv run python -m
src.scripts.configure_embeddings --report` (Phase 3).
ones. Full enumeration is available via `uv run python scripts/configure_embeddings.py --report`.
"""
from __future__ import annotations
@ -139,7 +138,7 @@ def _assert_pgvector_dims_match(
raise StartupValidationError(
f"{schema}.{table}.embedding has no declared vector dimension"
+ " (unbounded typmod). Run"
+ " `uv run python -m src.scripts.configure_embeddings`."
+ " `uv run python scripts/configure_embeddings.py`."
)
# pgvector stores the declared dim directly in atttypmod (no VARHDRSZ).
actual = atttypmod
@ -147,7 +146,7 @@ def _assert_pgvector_dims_match(
raise StartupValidationError(
f"{schema}.{table}.embedding dim ({actual}) does not match"
+ f" EMBEDDING_VECTOR_DIMENSIONS ({target_dim}). Run"
+ " `uv run python -m src.scripts.configure_embeddings`"
+ " `uv run python scripts/configure_embeddings.py`"
+ " or fix EMBEDDING_VECTOR_DIMENSIONS."
)
@ -191,7 +190,7 @@ async def _sample_external_namespaces(engine: AsyncEngine, *, target_dim: int) -
raise StartupValidationError(
f"Existing external-store namespaces have dim != {target_dim}:"
+ f" {formatted}. Run"
+ " `uv run python -m src.scripts.configure_embeddings --report`."
+ " `uv run python scripts/configure_embeddings.py --report`."
)

0
tests/scripts/__init__.py Normal file
View File

178
tests/scripts/test_configure_embeddings.py Normal file
View File

@ -0,0 +1,178 @@
"""Phase 3: configure_embeddings script tests."""
from __future__ import annotations
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
import pytest
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine
from scripts.configure_embeddings import (
_apply_pgvector_alter, # pyright: ignore[reportPrivateUsage]
_build_pgvector_plan, # pyright: ignore[reportPrivateUsage]
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@asynccontextmanager
async def _restore_schema_to(db_engine: AsyncEngine, dim: int) -> AsyncGenerator[None]:
"""ALTER both embedding columns back to ``dim`` on exit so this test
leaves the shared test DB in a consistent state for subsequent tests."""
try:
yield
finally:
async with db_engine.begin() as conn:
for table in ("documents", "message_embeddings"):
await conn.execute(
text(
f"ALTER TABLE {table} ALTER COLUMN embedding"
+ f" TYPE vector({dim}) USING NULL"
)
)
async def _current_dims(db_engine: AsyncEngine) -> dict[str, int]:
async with db_engine.connect() as conn:
result = await conn.execute(
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 = 'public'
AND c.relname = ANY(:tables)
AND a.attname = 'embedding'
"""
),
{"tables": ["documents", "message_embeddings"]},
)
return {row.table_name: row.typmod for row in result}
async def _hnsw_indexes(db_engine: AsyncEngine) -> set[str]:
async with db_engine.connect() as conn:
result = await conn.execute(
text(
"""
SELECT indexname
FROM pg_indexes
WHERE schemaname = 'public'
AND tablename IN ('documents', 'message_embeddings')
AND indexdef ILIKE '%USING hnsw%'
"""
)
)
return {row.indexname for row in result}
# ---------------------------------------------------------------------------
# Plan
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_plan_no_alter_needed_when_dims_already_match(
db_engine: AsyncEngine,
) -> None:
plan = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public")
assert plan.needs_alter is False
assert plan.current_dims == {"documents": 1536, "message_embeddings": 1536}
@pytest.mark.asyncio
async def test_plan_needs_alter_when_target_differs(db_engine: AsyncEngine) -> None:
plan = await _build_pgvector_plan(db_engine, target_dim=768, schema="public")
assert plan.needs_alter is True
@pytest.mark.asyncio
async def test_plan_raises_on_missing_column(db_engine: AsyncEngine) -> None:
with pytest.raises(SystemExit, match="required vector columns missing"):
await _build_pgvector_plan(db_engine, target_dim=1536, schema="no_such_schema")
# ---------------------------------------------------------------------------
# Apply
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_apply_alters_dims_and_recreates_hnsw_indexes(
db_engine: AsyncEngine,
) -> None:
# 768 is the canonical "small" dim used in non-1536 deployments and is
# well below pgvector's 2000-dim HNSW limit.
target = 768
async with _restore_schema_to(db_engine, dim=1536):
before_indexes = await _hnsw_indexes(db_engine)
assert before_indexes, "test fixture should have HNSW indexes pre-alter"
plan = await _build_pgvector_plan(db_engine, target_dim=target, schema="public")
assert plan.needs_alter is True
await _apply_pgvector_alter(db_engine, plan)
after_dims = await _current_dims(db_engine)
assert after_dims == {"documents": target, "message_embeddings": target}
after_indexes = await _hnsw_indexes(db_engine)
assert (
after_indexes == before_indexes
), "HNSW indexes should be recreated with the same names"
@pytest.mark.asyncio
async def test_apply_refuses_when_embeddings_populated(
db_engine: AsyncEngine,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""ALTER ... USING NULL would silently wipe non-null embeddings, so the
pre-check must abort the transaction before any destructive action.
We patch the count helper to simulate populated tables rather than wire
up the full FK chain of workspace/peer/collection/document just to land
one vector row.
"""
async def fake_count(_conn: object, _schema: str, table: str) -> int:
return 7 if table == "documents" else 0
monkeypatch.setattr(
"scripts.configure_embeddings._count_non_null_embeddings",
fake_count,
)
async with _restore_schema_to(db_engine, dim=1536):
plan = await _build_pgvector_plan(db_engine, target_dim=768, schema="public")
with pytest.raises(
SystemExit, match="refusing to ALTER populated embedding tables"
):
await _apply_pgvector_alter(db_engine, plan)
# The SystemExit aborts the transaction; nothing should have changed.
dims_after_refuse = await _current_dims(db_engine)
assert dims_after_refuse == {
"documents": 1536,
"message_embeddings": 1536,
}
# ---------------------------------------------------------------------------
# Idempotency
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_idempotent_apply_is_a_noop(db_engine: AsyncEngine) -> None:
"""Build plan twice with the matching dim — second call should still
return needs_alter=False without raising or making any changes."""
plan_a = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public")
plan_b = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public")
assert plan_a.needs_alter is False
assert plan_b.needs_alter is False
assert plan_a.current_dims == plan_b.current_dims