diff --git a/docs/v3/contributing/changing-embeddings.mdx b/docs/v3/contributing/changing-embeddings.mdx index 24491038..b93e30ec 100644 --- a/docs/v3/contributing/changing-embeddings.mdx +++ b/docs/v3/contributing/changing-embeddings.mdx @@ -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. diff --git a/scripts/configure_embeddings.py b/scripts/configure_embeddings.py index 389876a4..490c0580 100644 --- a/scripts/configure_embeddings.py +++ b/scripts/configure_embeddings.py @@ -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__": diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index d017a813..20ffbb66 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -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) diff --git a/src/startup/embedding_validator.py b/src/startup/embedding_validator.py index ca996c16..a397c9e2 100644 --- a/src/startup/embedding_validator.py +++ b/src/startup/embedding_validator.py @@ -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. diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 8004d8e8..1a71ff83 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -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." + ) diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index 0e72c785..b7a0a5c5 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -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 `[]`" + + " (e.g. `[768]f32`). SDK format may have changed." + ) + return int(match.group(1)) diff --git a/tests/llm/test_model_config.py b/tests/llm/test_model_config.py index 37c1ee68..bfc1f392 100644 --- a/tests/llm/test_model_config.py +++ b/tests/llm/test_model_config.py @@ -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