fix(embedding): address PR #678 review comments
CodeRabbit + Rajat review feedback. All actionable items addressed except two false-positives (responded on PR). Bug fixes: - deriver telemetry leak: validator was called outside try/finally so shutdown_telemetry() did not run on validation failure. Moved inside. - _emit_report printed "no effect with pgvector" unconditionally, including from implicit post-apply calls. Added is_report_mode flag; only print on explicit --report. - LanceDB and Turbopuffer probes returned None when the namespace existed but its schema was malformed (no vector field / unparseable type string), silently bucketing real corruption as "missing" (lazy-create) and letting it pass the startup validator. Now raise VectorStoreError with actionable diagnostics; None remains valid only for "namespace does not exist." - Startup validator only sampled message namespaces; added a parallel Collection-row sample so document namespaces are probed too, with the same dim assertion. Mirrors the --report path. Hygiene: - StartupValidationError now subclasses HonchoException so existing exception handlers recognize it. ValidationException is @final and has 422 request-validation semantics that would be misleading here. - scripts/configure_embeddings.py main() no longer spins up two event loops. engine.dispose() moved into a try/finally inside _async_main so cleanup runs in the same loop as the pipeline. - Replaced hand-rolled retry loop with tenacity.AsyncRetrying; same fail-closed semantics, less code, before_sleep_log for visibility. - Added _validate_identifier() defense-in-depth: DB.SCHEMA and HNSW index names are regex-checked against [A-Za-z_][A-Za-z0-9_]* before SQL interpolation. Operator config + DB catalog are not user input under the current threat model, but the constraint is cheap to gate. Test + docs: - test_app_settings_accepts_non_1536_with_any_vector_store_configuration now actually exercises turbopuffer (was missing); supplies a dummy TURBOPUFFER_API_KEY to satisfy the model_validator. - changing-embeddings.mdx: hyphenated "out-of-band" per reviewer style.
This commit is contained in:
Vineeth Voruganti
parent
2de0401dc0
commit
e4da919e0b
|
|
@ -9,7 +9,7 @@ icon: "rotate"
|
|||
The embedding dimension is **machine-enforced** as immutable for the life of a deployment. The embedding model is **operator-owned** as immutable by contract. The supported way to change either is:
|
||||
|
||||
1. Stand up a new deployment at the desired configuration.
|
||||
2. Replay or re-embed your data into it out of band.
|
||||
2. Replay or re-embed your data into it out-of-band.
|
||||
3. Cut traffic over to the new deployment.
|
||||
|
||||
The rest of this page explains why, and what the safety boundaries actually are.
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import argparse
|
|||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
|
@ -49,6 +50,22 @@ 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
|
||||
|
|
@ -200,6 +217,7 @@ async def _apply_pgvector_alter(engine: AsyncEngine, plan: _PgvectorPlan) -> Non
|
|||
# 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}"'))
|
||||
|
||||
|
|
@ -288,12 +306,20 @@ async def _probe_namespace_dim(store: VectorStore, namespace: str) -> int | None
|
|||
return await store.probe_namespace_dim(namespace)
|
||||
|
||||
|
||||
async def _emit_report(engine: AsyncEngine, target_dim: int) -> int:
|
||||
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":
|
||||
print("--report has no effect with 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)
|
||||
|
|
@ -384,11 +410,21 @@ def _confirm(prompt: str) -> bool:
|
|||
|
||||
|
||||
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)
|
||||
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:
|
||||
|
|
@ -397,7 +433,7 @@ async def _async_main(args: argparse.Namespace) -> int:
|
|||
+ f" {schema}.message_embeddings.embedding already at dim {target_dim},"
|
||||
+ " skipping ALTER"
|
||||
)
|
||||
return await _emit_report(engine, target_dim)
|
||||
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
|
||||
|
|
@ -425,20 +461,14 @@ async def _async_main(args: argparse.Namespace) -> int:
|
|||
|
||||
await _apply_pgvector_alter(engine, plan)
|
||||
print(f"\npgvector schema is now at dim {target_dim}")
|
||||
return await _emit_report(engine, 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)
|
||||
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())
|
||||
return asyncio.run(_async_main(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
|
|
@ -59,11 +59,11 @@ async def run_deriver():
|
|||
# Initialize async telemetry (CloudEvents emitter)
|
||||
await initialize_telemetry_async()
|
||||
|
||||
# Fail fast if the embedding schema does not match settings — same gate
|
||||
# the API runs in its lifespan.
|
||||
await validate_embedding_schema(engine)
|
||||
|
||||
try:
|
||||
# Fail fast if the embedding schema does not match settings — same
|
||||
# gate the API runs in its lifespan. Inside the try block so the
|
||||
# telemetry buffer is still flushed if validation raises.
|
||||
await validate_embedding_schema(engine)
|
||||
await main()
|
||||
finally:
|
||||
# Shutdown telemetry (flush CloudEvents buffer)
|
||||
|
|
|
|||
|
|
@ -13,15 +13,23 @@ ones. Full enumeration is available via `uv run python scripts/configure_embeddi
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from sqlalchemy import select, text
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
from tenacity import (
|
||||
AsyncRetrying,
|
||||
RetryError,
|
||||
before_sleep_log,
|
||||
retry_if_exception_type,
|
||||
stop_after_attempt,
|
||||
wait_fixed,
|
||||
)
|
||||
|
||||
from src.config import AppSettings, settings
|
||||
from src.models import Workspace
|
||||
from src.exceptions import HonchoException
|
||||
from src.models import Collection, Workspace
|
||||
from src.vector_store import VectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -38,10 +46,15 @@ _RETRY_BACKOFF_SECONDS = 1.0
|
|||
_EXTERNAL_SAMPLE_LIMIT = 10
|
||||
|
||||
|
||||
class StartupValidationError(RuntimeError):
|
||||
class StartupValidationError(HonchoException):
|
||||
"""Raised when the embedding configuration cannot be reconciled with the
|
||||
physical schema. Always surfaced before any HTTP route is served or any
|
||||
queue task is processed.
|
||||
|
||||
Inherits from ``HonchoException`` (status_code=500) so the project's
|
||||
exception handlers recognize it consistently. Startup-time failure, not
|
||||
a per-request validation error — ``ValidationException``'s 422 semantics
|
||||
would be misleading.
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -76,24 +89,23 @@ async def _introspect_pgvector_dims_with_retry(
|
|||
columns. Fails closed on the last attempt — uncertainty is not a green
|
||||
light to serve traffic.
|
||||
"""
|
||||
last_exc: Exception | None = None
|
||||
for attempt in range(_RETRY_ATTEMPTS):
|
||||
try:
|
||||
return await _introspect_pgvector_dims_once(engine, schema)
|
||||
except SQLAlchemyError as e:
|
||||
last_exc = e
|
||||
remaining = _RETRY_ATTEMPTS - attempt - 1
|
||||
if remaining > 0:
|
||||
logger.warning(
|
||||
"Embedding schema introspection attempt %d/%d failed: %s",
|
||||
attempt + 1,
|
||||
_RETRY_ATTEMPTS,
|
||||
e,
|
||||
)
|
||||
await asyncio.sleep(_RETRY_BACKOFF_SECONDS)
|
||||
raise StartupValidationError(
|
||||
f"could not validate embedding schema: {last_exc}"
|
||||
) from last_exc
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(_RETRY_ATTEMPTS),
|
||||
wait=wait_fixed(_RETRY_BACKOFF_SECONDS),
|
||||
retry=retry_if_exception_type(SQLAlchemyError),
|
||||
before_sleep=before_sleep_log(logger, logging.WARNING),
|
||||
reraise=False,
|
||||
):
|
||||
with attempt:
|
||||
return await _introspect_pgvector_dims_once(engine, schema)
|
||||
except RetryError as e:
|
||||
underlying = e.last_attempt.exception()
|
||||
raise StartupValidationError(
|
||||
f"could not validate embedding schema: {underlying}"
|
||||
) from underlying
|
||||
# Unreachable: AsyncRetrying either returns from inside the loop or raises.
|
||||
raise StartupValidationError("embedding schema introspection did not run")
|
||||
|
||||
|
||||
async def _introspect_pgvector_dims_once(
|
||||
|
|
@ -159,15 +171,23 @@ async def _sample_external_namespaces(engine: AsyncEngine, *, target_dim: int) -
|
|||
External stores in this codebase are per-workspace and lazy-created on
|
||||
first write (see ``src.vector_store.get_vector_namespace``), so there is
|
||||
no canonical deployment-wide namespace to introspect. We enumerate up to
|
||||
``_EXTERNAL_SAMPLE_LIMIT`` workspaces from the application DB and probe
|
||||
their derived namespaces. Missing namespaces are OK; mismatched dims
|
||||
crash startup. Run ``configure_embeddings --report`` for full
|
||||
enumeration when a hard guarantee is needed.
|
||||
``_EXTERNAL_SAMPLE_LIMIT`` of each namespace category from the application
|
||||
DB and probe each:
|
||||
|
||||
- Message namespaces — one per workspace.
|
||||
- Document namespaces — one per existing ``(workspace, observer, observed)``
|
||||
collection triple.
|
||||
|
||||
Missing namespaces are OK; mismatched dims crash startup. Run
|
||||
``configure_embeddings --report`` for full enumeration when a hard
|
||||
guarantee is needed.
|
||||
"""
|
||||
workspace_names = await _sample_workspace_names(engine, _EXTERNAL_SAMPLE_LIMIT)
|
||||
if not workspace_names:
|
||||
collection_keys = await _sample_collection_keys(engine, _EXTERNAL_SAMPLE_LIMIT)
|
||||
if not workspace_names and not collection_keys:
|
||||
logger.info(
|
||||
"External-store validator: no workspaces exist yet, skipping sample"
|
||||
"External-store validator: no workspaces or collections exist yet,"
|
||||
+ " skipping sample"
|
||||
)
|
||||
return
|
||||
|
||||
|
|
@ -180,9 +200,21 @@ async def _sample_external_namespaces(engine: AsyncEngine, *, target_dim: int) -
|
|||
# That is its own problem and not for this validator to swallow.
|
||||
return
|
||||
|
||||
mismatches: list[tuple[str, int]] = []
|
||||
candidates: list[str] = []
|
||||
for workspace_name in workspace_names:
|
||||
namespace = store.get_vector_namespace("message", workspace_name)
|
||||
candidates.append(store.get_vector_namespace("message", workspace_name))
|
||||
for workspace_name, observer, observed in collection_keys:
|
||||
candidates.append(
|
||||
store.get_vector_namespace(
|
||||
"document",
|
||||
workspace_name,
|
||||
observer=observer,
|
||||
observed=observed,
|
||||
)
|
||||
)
|
||||
|
||||
mismatches: list[tuple[str, int]] = []
|
||||
for namespace in candidates:
|
||||
actual_dim = await _probe_namespace_dim(store, namespace)
|
||||
if actual_dim is not None and actual_dim != target_dim:
|
||||
mismatches.append((namespace, actual_dim))
|
||||
|
|
@ -209,6 +241,22 @@ async def _sample_workspace_names(engine: AsyncEngine, limit: int) -> list[str]:
|
|||
return [row[0] for row in result]
|
||||
|
||||
|
||||
async def _sample_collection_keys(
|
||||
engine: AsyncEngine, limit: int
|
||||
) -> list[tuple[str, str, str]]:
|
||||
"""Pull up to ``limit`` ``(workspace_name, observer, observed)`` triples,
|
||||
one per existing collection row. Each triple corresponds to a document
|
||||
namespace that may exist in the external store."""
|
||||
stmt = (
|
||||
select(Collection.workspace_name, Collection.observer, Collection.observed)
|
||||
.order_by(Collection.created_at.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
async with engine.connect() as conn:
|
||||
result = await conn.execute(stmt)
|
||||
return [(row[0], row[1], row[2]) for row in result]
|
||||
|
||||
|
||||
async def _probe_namespace_dim(store: VectorStore, namespace: str) -> int | None:
|
||||
"""Return the namespace's declared dim, or ``None`` if not present.
|
||||
|
||||
|
|
|
|||
|
|
@ -371,7 +371,15 @@ class LanceDBVectorStore(VectorStore):
|
|||
logger.debug("LanceDB connection closed")
|
||||
|
||||
async def probe_namespace_dim(self, namespace: str) -> int | None:
|
||||
"""Inspect a LanceDB table's vector column to recover its declared dim."""
|
||||
"""Inspect a LanceDB table's vector column to recover its declared dim.
|
||||
|
||||
Returns ``None`` only when the table does not exist (lazy-create
|
||||
model, expected case). When the table exists but its schema does
|
||||
not include a ``vector`` field with a fixed ``list_size``, raises
|
||||
``VectorStoreError`` — that is a malformed table, not a missing one,
|
||||
and silently bucketing it as "missing" would let real corruption
|
||||
through the startup validator.
|
||||
"""
|
||||
db = await self._get_db()
|
||||
table_names = await db.table_names()
|
||||
if namespace not in table_names:
|
||||
|
|
@ -381,4 +389,8 @@ class LanceDBVectorStore(VectorStore):
|
|||
for field in schema:
|
||||
if field.name == "vector" and hasattr(field.type, "list_size"):
|
||||
return int(field.type.list_size)
|
||||
return None
|
||||
raise VectorStoreError(
|
||||
f"LanceDB table {namespace!r} exists but has no 'vector' field"
|
||||
+ " with a fixed dimension; cannot probe dim. Schema may be"
|
||||
+ " corrupted — inspect with `lancedb` CLI before retrying."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -317,6 +317,12 @@ class TurbopufferVectorStore(VectorStore):
|
|||
to ``AttributeSchemaConfig``; the vector field's ``type`` string is
|
||||
a bracket-prefixed dim with a width suffix, e.g. ``"[768]f32"``,
|
||||
``"[1536]f16"``, ``"[256]i8"``.
|
||||
|
||||
Returns ``None`` only when the namespace does not exist yet
|
||||
(NotFoundError or ``exists() == False``). When the namespace
|
||||
exists but its schema lacks a parseable ``vector`` attribute,
|
||||
raises ``VectorStoreError`` — silently bucketing that as "missing"
|
||||
would let a corrupt namespace pass the startup validator.
|
||||
"""
|
||||
ns = self._get_namespace(namespace)
|
||||
try:
|
||||
|
|
@ -332,6 +338,16 @@ class TurbopufferVectorStore(VectorStore):
|
|||
|
||||
vector_attr = schema.get("vector")
|
||||
if vector_attr is None:
|
||||
return None
|
||||
match = re.search(r"\[(\d+)\]", str(vector_attr.type))
|
||||
return int(match.group(1)) if match else None
|
||||
raise VectorStoreError(
|
||||
f"Turbopuffer namespace {namespace!r} exists but its schema"
|
||||
+ " has no 'vector' attribute; cannot probe dim."
|
||||
)
|
||||
type_str = str(vector_attr.type)
|
||||
match = re.search(r"\[(\d+)\]", type_str)
|
||||
if match is None:
|
||||
raise VectorStoreError(
|
||||
f"Turbopuffer namespace {namespace!r} has an unparseable"
|
||||
+ f" vector type {type_str!r}; expected `[<dim>]<width>`"
|
||||
+ " (e.g. `[768]f32`). SDK format may have changed."
|
||||
)
|
||||
return int(match.group(1))
|
||||
|
|
|
|||
|
|
@ -289,11 +289,19 @@ def test_app_settings_accepts_non_1536_with_any_vector_store_configuration() ->
|
|||
("pgvector", False),
|
||||
("lancedb", True),
|
||||
("lancedb", False),
|
||||
("turbopuffer", True),
|
||||
("turbopuffer", False),
|
||||
]
|
||||
for store_type, migrated in combos:
|
||||
# Turbopuffer's model_validator requires TURBOPUFFER_API_KEY whenever
|
||||
# TYPE="turbopuffer"; supply a dummy value so the test exercises the
|
||||
# dim-acceptance path rather than the api-key guard.
|
||||
vs_kwargs: dict[str, Any] = {"TYPE": store_type, "MIGRATED": migrated}
|
||||
if store_type == "turbopuffer":
|
||||
vs_kwargs["TURBOPUFFER_API_KEY"] = "test-key"
|
||||
settings = AppSettings(
|
||||
EMBEDDING=EmbeddingSettings(VECTOR_DIMENSIONS=768),
|
||||
VECTOR_STORE=VectorStoreSettings(TYPE=store_type, MIGRATED=migrated),
|
||||
VECTOR_STORE=VectorStoreSettings(**vs_kwargs),
|
||||
)
|
||||
assert settings.EMBEDDING.VECTOR_DIMENSIONS == 768
|
||||
assert store_type == settings.VECTOR_STORE.TYPE
|
||||
|
|
|
|||
Loading…
Reference in New Issue