Make embeddings configurable (#678)
* feat(embedding): add dimensions_mode for OpenAI dimensions= forwarding
Add EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE (auto|always|never) controlling
whether the dimensions= parameter is forwarded on OpenAI embeddings.create
calls. auto (default) sends it when the operator explicitly set
EMBEDDING_VECTOR_DIMENSIONS and the configured model is not on the
known-rejecting allowlist (currently text-embedding-ada-002).
The provenance check (was VECTOR_DIMENSIONS explicitly set?) lives as
EmbeddingSettings.resolve_send_dimensions() because it needs access to
model_fields_set, which the standalone resolver does not have. The
resolved boolean is passed into _EmbeddingClient at construction time;
the client never inspects mode or provenance.
Also pins cloudevents <2.0 — 2.0.0 reorganized the package and dropped
cloudevents.conversion and cloudevents.http, which src/telemetry/emitter.py
imports. The original `>=1.12.0` constraint allowed the broken 2.0 resolve.
With the pin, the imports resolve cleanly and the basedpyright warning
cascade (37+ warnings about unknown types) disappears.
Drive-by cleanups (all unnecessary cast/ignore comments flagged by
basedpyright after the cloudevents downgrade):
- vector_store/lancedb.py, tests/conftest.py, and
tests/deriver/test_vector_reconciliation.py — drop dead pyright ignores
- sdks/python/src/honcho/http/{async_,}client.py — drop unnecessary
cast(datetime, ...) (parsedate_to_datetime already returns datetime)
- vector_store/turbopuffer.py — cast(Any, rows) for the upsert_rows
TypedDict that the SDK exposes but our row builder doesn't satisfy
- tests/test_datetime_parsing.py — ignore reportArgumentType on the
test that deliberately passes wrong types to assert raises
* feat(models): honor EMBEDDING_VECTOR_DIMENSIONS in pgvector columns
* feat(startup): atomic swap dim-vs-MIGRATED guard for runtime schema validator
Add src/startup/embedding_validator.py that introspects the actual pgvector
column dim at boot and refuses to start if it does not match
EMBEDDING_VECTOR_DIMENSIONS. Runs after the DB pool is up and before the
embedding client is constructed, in both src/main.py (FastAPI lifespan) and
src/deriver/__main__.py.
Implementation details:
- Schema-qualified pg_attribute join through pg_class/pg_namespace respects
DB.SCHEMA rather than relying on search_path
- Bounded retry (3 attempts, 1s backoff) for transient introspection failure,
then fail-closed with "could not validate embedding schema" — uncertainty
is not a green light to serve traffic
- External-store sampler (turbopuffer, lancedb) enumerates workspaces from
the application DB and probes their lazy-created namespaces; current
per-namespace probe is a no-op stub since the SDKs do not expose
uniform dim introspection — full enumeration is left to
`configure_embeddings --report` in Phase 3
Atomic guard swap: deletes the old dim-vs-MIGRATED config validator (which
forbade non-1536 pgvector unless MIGRATED=True) in the same commit as the
new runtime validator. There is no release window where non-1536 pgvector
can start unprotected. The 9 dual-write branches that use VECTOR_STORE.MIGRATED
remain untouched and load-bearing for legacy-tenant backend swaps.
VECTOR_STORE_DIMENSIONS deprecation: drop the "must match" raise; in
propagate_namespace, check model_fields_set and emit logger.warning +
DeprecationWarning (DeprecationWarning alone is filtered by Python's default
config and would not reach operators). Always overwrite with
EMBEDDING.VECTOR_DIMENSIONS regardless.
Test changes:
- tests/test_models_vector_dim.py: Phase 1's VECTOR_STORE_TYPE=lancedb +
MIGRATED=true escape hatches removed; the test now passes on plain
EMBEDDING_VECTOR_DIMENSIONS=768
- tests/llm/test_model_config.py: the two tests asserting the old guards
replaced with tests for the new deprecation + acceptance behavior
- tests/startup/test_embedding_validator.py: 10 new tests — dim assertion
logic (pass/mismatch/missing/unbounded/non-public-schema), fail-closed
retry budget, real-test-DB pass, real-DB ALTER-then-validate, deprecation
warning capture, non-1536 + pgvector + MIGRATED=false at config time
* 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.
* docs: add changing-embeddings operations page
Document the supported way to change EMBEDDING_VECTOR_DIMENSIONS or
EMBEDDING_MODEL_CONFIG__MODEL on a Honcho deployment: provision a new
deployment at the desired configuration, replay source data out of
band, cut over at the application layer.
The page explains the asymmetry:
- Dimension is machine-enforced as immutable. The startup validator
introspects pg_attribute and crashes the API/deriver on mismatch.
- Model is operator-owned. There is no persistent metadata recording
which model produced each vector, so a same-dim model swap is
silently undetectable — flagged with a Warning callout.
Also documents the truncation edge case (text-embedding-3-large truncated to 1536 with EMBEDDING_VECTOR_DIMENSIONS left at default)
and the DIMENSIONS_MODE=always mitigation, plus a pointer that
storage-backend swap (VECTOR_STORE_MIGRATED + reconciler) is a distinct operation unaffected by this work.
Registers the page in docs/docs.json under the Self-Hosting nav group
and cross-links from configuration.mdx.
* fix(embedding): correct turbopuffer regex + tighten DIMENSIONS_MODE docs
- Turbopuffer attribute type for a vector column is `[N]f32` / `[N]f16` /
`[N]i8`, not `f32_vector(N)` as the earlier probe assumed. The earlier
regex returned None for the real SDK format, so existing Turbopuffer
namespaces would have been reported as "missing" instead of validated
for mismatch. Regex switched to `\[(\d+)\]` which is the
vendor-stable shape. Test cases rewritten to lock the actual format.
- docs/v3/contributing/configuration.mdx had a contradictory pair of
bullets: 223 said explicit 1536 makes `auto` forward dimensions=, 224
said `auto` would skip the parameter because 1536 is the default.
Operators reading both would (rightly) conclude they need `always`
even when `auto` would work. Rewrote both bullets so:
- `auto` is provenance-driven (explicit-set, not non-default-value).
- `always` is positioned as defense-in-depth for config layers that
might strip explicit default-valued envs, not the only path for
same-as-default truncation.
* 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.
* fix: modify conftest to fix ci
* fix: ci tests for typescript server
This commit is contained in:
Vineeth Voruganti
GitHub
parent
b3f371ba51
commit
b84da15d03
|
|
@ -279,8 +279,8 @@ VECTOR_STORE_MIGRATED=false
|
|||
# - Messages: {NAMESPACE}.msg.{hash(workspace)}
|
||||
# VECTOR_STORE_NAMESPACE=honcho # Inherits from NAMESPACE if not set
|
||||
|
||||
# Embedding dimensions (default: 1536 for OpenAI text-embedding-3-small)
|
||||
# VECTOR_STORE_DIMENSIONS=1536
|
||||
# Embedding dimensions are configured via EMBEDDING_VECTOR_DIMENSIONS (see top
|
||||
# of this file). VECTOR_STORE_DIMENSIONS is deprecated and ignored.
|
||||
|
||||
# Turbopuffer-specific settings (required if TYPE is "turbopuffer")
|
||||
# VECTOR_STORE_TURBOPUFFER_API_KEY=your-turbopuffer-api-key
|
||||
|
|
|
|||
|
|
@ -255,9 +255,7 @@ TYPE = "pgvector"
|
|||
# Migration flag: set to true when migration from pgvector is complete
|
||||
MIGRATED = false
|
||||
NAMESPACE = "honcho"
|
||||
# This should match embedding.vector_dimensions. pgvector and dual-write mode
|
||||
# currently still require 1536 until a schema migration lands.
|
||||
DIMENSIONS = 1536
|
||||
# DIMENSIONS is deprecated; embedding.vector_dimensions is authoritative.
|
||||
# TURBOPUFFER_API_KEY = "your-turbopuffer-api-key"
|
||||
# TURBOPUFFER_REGION = "us-east-1"
|
||||
LANCEDB_PATH = "./lancedb_data"
|
||||
|
|
|
|||
|
|
@ -149,6 +149,7 @@
|
|||
"pages": [
|
||||
"v3/contributing/self-hosting",
|
||||
"v3/contributing/configuration",
|
||||
"v3/contributing/changing-embeddings",
|
||||
"v3/contributing/troubleshooting"
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -0,0 +1,62 @@
|
|||
---
|
||||
title: "Changing Embeddings"
|
||||
description: "How to switch embedding dimension or model on a Honcho deployment"
|
||||
icon: "rotate"
|
||||
---
|
||||
|
||||
## Short answer: you can't, in place.
|
||||
|
||||
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.
|
||||
3. Cut traffic over to the new deployment.
|
||||
|
||||
The rest of this page explains why, and what the safety boundaries actually are.
|
||||
|
||||
## Why dimension is enforced and model is not
|
||||
|
||||
On boot, both the API (`src/main.py` lifespan) and the deriver (`src/deriver/__main__.py`) run the validator in `src/startup/embedding_validator.py`. It does a schema-qualified `pg_attribute` lookup against `documents.embedding` and `message_embeddings.embedding`, decodes the declared `atttypmod`, and compares it to `EMBEDDING_VECTOR_DIMENSIONS`. A mismatch crashes the process with an actionable error before any HTTP route is served or any queue task is processed.
|
||||
|
||||
There is no equivalent check for the model. The pgvector column does not record what model produced the vectors inside it, and this design intentionally avoids adding new persistent metadata fields. The runtime has no way to detect that you swapped `text-embedding-3-small` for a different model that emits the same dimension.
|
||||
|
||||
That last point is a real footgun:
|
||||
|
||||
<Warning>
|
||||
Changing `EMBEDDING_MODEL_CONFIG__MODEL` to a different model at the **same dimension** (for example `text-embedding-3-small@1536` → `text-embedding-3-large` truncated to 1536) will silently succeed. New writes will use the new model; existing rows still hold vectors from the old model; recall quality will degrade with no startup or runtime warning.
|
||||
|
||||
Treat model identity as a contract you own. If you need to change it, follow the destroy + rebuild path below.
|
||||
</Warning>
|
||||
|
||||
## Recipe: changing dim or model
|
||||
|
||||
Concretely, for either a dim change or a model change:
|
||||
|
||||
1. **Provision the new deployment** with the target environment.
|
||||
|
||||
```bash
|
||||
# On the new deployment:
|
||||
export EMBEDDING_VECTOR_DIMENSIONS=768
|
||||
export EMBEDDING_MODEL_CONFIG__TRANSPORT=openai
|
||||
export EMBEDDING_MODEL_CONFIG__MODEL=nomic-embed-text
|
||||
export EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=http://your-ollama:11434/v1
|
||||
alembic upgrade head
|
||||
uv run python scripts/configure_embeddings.py --dry-run
|
||||
uv run python scripts/configure_embeddings.py --yes
|
||||
```
|
||||
|
||||
2. **Replay your source data** (messages, documents, ingested content) into the new deployment via your normal application path. Honcho's existing message-creation API will re-derive embeddings using the new configuration. There is no in-place re-embedding tool — that would be a separate spec covering atomicity, cost-per-token, and dialectic-during-migration semantics.
|
||||
|
||||
3. **Cut over** at your application layer (DNS, load balancer, feature flag — whatever you use). The old deployment can stay running until you are confident in the new one; this design does not require an atomic switch.
|
||||
|
||||
The startup validator on the new deployment will refuse to start if step 1's `configure_embeddings.py` did not run, so a misconfiguration cannot quietly write wrong-dim vectors into the new schema.
|
||||
|
||||
## Edge case: truncation at the default dimension
|
||||
|
||||
If you are using `text-embedding-3-large` but truncating to 1536 (the default), be aware that `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE=auto` will **not** forward `dimensions=` to the API — `auto` interprets the default as "operator did not opt into a non-default dim." The provider will return native 3072, the response-dim validator will reject it, and the request will fail.
|
||||
|
||||
For this case, either set `EMBEDDING_VECTOR_DIMENSIONS=1536` explicitly (so `auto` knows the operator opted in), or set `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE=always`.
|
||||
|
||||
## Backend swap (turbopuffer ↔ lancedb ↔ pgvector) is a different operation
|
||||
|
||||
Switching the *storage backend* at constant dim/model — for example moving from pgvector to Turbopuffer — is supported via `src/reconciler/sync_vectors.py` and `VECTOR_STORE_MIGRATED`. That flow is unchanged by the embedding-pipeline work and is documented separately. It is **not** the destroy + rebuild path described above.
|
||||
|
|
@ -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 **the operator has explicitly set `EMBEDDING_VECTOR_DIMENSIONS`** — provenance, not value — and the configured model is not on the known-rejecting list (currently `text-embedding-ada-002`). Explicit `EMBEDDING_VECTOR_DIMENSIONS=1536` *does* trigger the forward; this is how `text-embedding-3-large` truncation to 1536 is expressed. Deployments that leave the setting unset get their existing behavior (`dimensions=` is not forwarded).
|
||||
- `always`: always forward, regardless of whether `EMBEDDING_VECTOR_DIMENSIONS` was set. Use for OpenAI-compatible self-hosted providers that require it. Do not pick `always` *just* for same-as-default truncation — `auto` handles that case correctly as long as you set `EMBEDDING_VECTOR_DIMENSIONS=1536` explicitly in your environment. `always` is the right answer when your config layer might strip explicit "default-valued" envs, or when you want defense-in-depth.
|
||||
- `never`: never forward. Explicit opt-out for providers that reject the parameter (e.g. `text-embedding-ada-002` if it slips past the known-rejecting allowlist).
|
||||
|
||||
#### 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. See [Changing Embeddings](./changing-embeddings) for the destroy + rebuild recipe and the same-dim model-swap caveat.
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -459,7 +499,8 @@ WEBHOOK_MAX_WORKSPACE_LIMIT=10
|
|||
VECTOR_STORE_TYPE=pgvector # Options: pgvector, turbopuffer, lancedb
|
||||
VECTOR_STORE_MIGRATED=false
|
||||
VECTOR_STORE_NAMESPACE=honcho
|
||||
VECTOR_STORE_DIMENSIONS=1536
|
||||
# Embedding dim is configured via EMBEDDING_VECTOR_DIMENSIONS — see the
|
||||
# Embedding Configuration section. VECTOR_STORE_DIMENSIONS is deprecated.
|
||||
|
||||
# Turbopuffer-specific
|
||||
VECTOR_STORE_TURBOPUFFER_API_KEY=your-turbopuffer-api-key
|
||||
|
|
@ -637,4 +678,4 @@ uv run alembic revision --autogenerate -m "Description" # Create new migration
|
|||
|
||||
5. **Dialectic level issues** — Unset level fields inherit from the built-in defaults. For Anthropic, `THINKING_BUDGET_TOKENS` must be >= 1024 when enabled. For providers without budgeted thinking, omit it or set it to `0`. `MAX_OUTPUT_TOKENS` must exceed `THINKING_BUDGET_TOKENS`.
|
||||
|
||||
6. **Vector store issues** — For Turbopuffer, set the API key. Check `VECTOR_STORE_DIMENSIONS` matches your embedding model.
|
||||
6. **Vector store issues** — For Turbopuffer, set the API key. Check that `EMBEDDING_VECTOR_DIMENSIONS` matches your embedding model — the startup validator will refuse to boot on a mismatch.
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ dependencies = [
|
|||
"cashews[redis]==7.4.4",
|
||||
"scikit-learn>=1.6.0",
|
||||
"prometheus_client>=0.21.0",
|
||||
"cloudevents>=1.12.0",
|
||||
"cloudevents>=1.12.0,<2.0",
|
||||
]
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
|
|
|
|||
|
|
@ -0,0 +1,475 @@
|
|||
"""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())
|
||||
|
|
@ -370,7 +370,7 @@ class AsyncHonchoHTTPClient:
|
|||
from datetime import datetime
|
||||
from email.utils import parsedate_to_datetime
|
||||
|
||||
dt: datetime = cast(datetime, parsedate_to_datetime(header))
|
||||
dt: datetime = parsedate_to_datetime(header)
|
||||
timestamp: float = dt.timestamp()
|
||||
return max(0.0, timestamp - time.time())
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -367,7 +367,7 @@ class HonchoHTTPClient:
|
|||
from datetime import datetime
|
||||
from email.utils import parsedate_to_datetime
|
||||
|
||||
dt: datetime = cast(datetime, parsedate_to_datetime(header))
|
||||
dt: datetime = parsedate_to_datetime(header)
|
||||
timestamp: float = dt.timestamp()
|
||||
return max(0.0, timestamp - time.time())
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,12 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
ModelTransport = Literal["anthropic", "openai", "gemini"]
|
||||
EmbeddingTransport = Literal["openai", "gemini"]
|
||||
EmbeddingDimensionsMode = Literal["auto", "always", "never"]
|
||||
|
||||
# OpenAI-compatible models that reject the `dimensions=` request parameter.
|
||||
_EMBEDDING_KNOWN_REJECTING_MODELS: frozenset[str] = frozenset(
|
||||
{"text-embedding-ada-002"}
|
||||
)
|
||||
|
||||
|
||||
def _default_embedding_model_for_transport(transport: EmbeddingTransport) -> str:
|
||||
|
|
@ -294,6 +300,7 @@ class ConfiguredEmbeddingModelSettings(BaseModel):
|
|||
model: str = "text-embedding-3-small"
|
||||
transport: EmbeddingTransport = "openai"
|
||||
overrides: ModelOverrideSettings = Field(default_factory=ModelOverrideSettings)
|
||||
dimensions_mode: EmbeddingDimensionsMode = "auto"
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
@ -701,6 +708,23 @@ class EmbeddingSettings(HonchoSettings):
|
|||
)
|
||||
return data # pyright: ignore[reportUnknownVariableType]
|
||||
|
||||
def resolve_send_dimensions(self) -> bool:
|
||||
"""Decide whether OpenAI embedding calls should forward ``dimensions=``.
|
||||
|
||||
Lives on the settings instance because ``auto`` mode needs access to
|
||||
``self.model_fields_set`` to tell whether the operator explicitly set
|
||||
``VECTOR_DIMENSIONS`` — a standalone resolver over
|
||||
``ConfiguredEmbeddingModelSettings`` cannot see that.
|
||||
"""
|
||||
mode = self.MODEL_CONFIG.dimensions_mode
|
||||
if mode == "always":
|
||||
return True
|
||||
if mode == "never":
|
||||
return False
|
||||
if self.MODEL_CONFIG.model in _EMBEDDING_KNOWN_REJECTING_MODELS:
|
||||
return False
|
||||
return "VECTOR_DIMENSIONS" in self.model_fields_set
|
||||
|
||||
|
||||
class DeriverSettings(HonchoSettings):
|
||||
model_config = SettingsConfigDict( # pyright: ignore
|
||||
|
|
@ -1272,25 +1296,27 @@ class AppSettings(HonchoSettings):
|
|||
self.CACHE.NAMESPACE = self.NAMESPACE
|
||||
if "NAMESPACE" not in self.VECTOR_STORE.model_fields_set:
|
||||
self.VECTOR_STORE.NAMESPACE = self.NAMESPACE
|
||||
if "DIMENSIONS" not in self.VECTOR_STORE.model_fields_set:
|
||||
self.VECTOR_STORE.DIMENSIONS = self.EMBEDDING.VECTOR_DIMENSIONS
|
||||
elif self.VECTOR_STORE.DIMENSIONS != self.EMBEDDING.VECTOR_DIMENSIONS:
|
||||
raise ValueError(
|
||||
"VECTOR_STORE.DIMENSIONS must match EMBEDDING.VECTOR_DIMENSIONS"
|
||||
if "DIMENSIONS" in self.VECTOR_STORE.model_fields_set:
|
||||
# VECTOR_STORE_DIMENSIONS is deprecated: EMBEDDING_VECTOR_DIMENSIONS
|
||||
# is the single source of truth. Log a runtime-visible warning
|
||||
# so operators see it (DeprecationWarning is filtered by Python's
|
||||
# default config outside __main__/tests) and also raise the stdlib
|
||||
# warning so tests can assert on it.
|
||||
import warnings
|
||||
|
||||
message = (
|
||||
"VECTOR_STORE_DIMENSIONS is deprecated; "
|
||||
"EMBEDDING_VECTOR_DIMENSIONS is authoritative. "
|
||||
"Drop VECTOR_STORE_DIMENSIONS from your .env."
|
||||
)
|
||||
logger.warning(message)
|
||||
warnings.warn(message, DeprecationWarning, stacklevel=2)
|
||||
self.VECTOR_STORE.DIMENSIONS = self.EMBEDDING.VECTOR_DIMENSIONS
|
||||
if "NAMESPACE" not in self.TELEMETRY.model_fields_set:
|
||||
self.TELEMETRY.NAMESPACE = self.NAMESPACE
|
||||
if "NAMESPACE" not in self.METRICS.model_fields_set:
|
||||
self.METRICS.NAMESPACE = self.NAMESPACE
|
||||
|
||||
if self.EMBEDDING.VECTOR_DIMENSIONS != 1536 and (
|
||||
self.VECTOR_STORE.TYPE == "pgvector" or not self.VECTOR_STORE.MIGRATED
|
||||
):
|
||||
raise ValueError(
|
||||
"EMBEDDING.VECTOR_DIMENSIONS must remain 1536 while pgvector is "
|
||||
+ "active or vector-store migration is incomplete"
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import uvloop
|
|||
from prometheus_client import start_http_server
|
||||
|
||||
from src.config import settings
|
||||
from src.db import engine
|
||||
from src.startup import validate_embedding_schema
|
||||
from src.telemetry import initialize_telemetry_async, shutdown_telemetry
|
||||
|
||||
from .queue_manager import main
|
||||
|
|
@ -56,7 +58,12 @@ async def run_deriver():
|
|||
"""Run the deriver with proper telemetry lifecycle management."""
|
||||
# Initialize async telemetry (CloudEvents emitter)
|
||||
await initialize_telemetry_async()
|
||||
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
import logging
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from typing import NamedTuple
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import tiktoken
|
||||
from google import genai
|
||||
|
|
@ -34,10 +34,12 @@ class _EmbeddingClient:
|
|||
vector_dimensions: int,
|
||||
max_input_tokens: int,
|
||||
max_tokens_per_request: int,
|
||||
send_dimensions: bool,
|
||||
):
|
||||
self.transport: str = config.transport
|
||||
self.model: str = config.model
|
||||
self.vector_dimensions: int = vector_dimensions
|
||||
self.send_dimensions: bool = send_dimensions
|
||||
|
||||
if self.transport == "gemini":
|
||||
if not config.api_key:
|
||||
|
|
@ -101,9 +103,10 @@ class _EmbeddingClient:
|
|||
raise ValueError("No embedding returned from Gemini API")
|
||||
return self._validate_embedding_dimensions(response.embeddings[0].values)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
model=self.model, input=[query]
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {"model": self.model, "input": [query]}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
return self._validate_embedding_dimensions(response.data[0].embedding)
|
||||
|
||||
async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
|
||||
|
|
@ -138,10 +141,13 @@ class _EmbeddingClient:
|
|||
self._validate_embedding_dimensions(emb.values)
|
||||
)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
input=batch,
|
||||
model=self.model,
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {
|
||||
"input": batch,
|
||||
"model": self.model,
|
||||
}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
embeddings.extend(
|
||||
[
|
||||
self._validate_embedding_dimensions(data.embedding)
|
||||
|
|
@ -287,9 +293,13 @@ class _EmbeddingClient:
|
|||
)
|
||||
)
|
||||
else: # openai
|
||||
response = await self.client.embeddings.create(
|
||||
model=self.model, input=[item.text for item in batch]
|
||||
)
|
||||
openai_kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"input": [item.text for item in batch],
|
||||
}
|
||||
if self.send_dimensions:
|
||||
openai_kwargs["dimensions"] = self.vector_dimensions
|
||||
response = await self.client.embeddings.create(**openai_kwargs)
|
||||
for item, embedding_data in zip(batch, response.data, strict=True):
|
||||
result[item.text_id][item.chunk_index] = (
|
||||
self._validate_embedding_dimensions(
|
||||
|
|
@ -411,6 +421,7 @@ class EmbeddingClient:
|
|||
vector_dimensions=settings.EMBEDDING.VECTOR_DIMENSIONS,
|
||||
max_input_tokens=settings.EMBEDDING.MAX_INPUT_TOKENS,
|
||||
max_tokens_per_request=settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
|
||||
send_dimensions=settings.EMBEDDING.resolve_send_dimensions(),
|
||||
)
|
||||
self._instance_signature = signature
|
||||
logger.debug(
|
||||
|
|
@ -434,6 +445,7 @@ class EmbeddingClient:
|
|||
settings.EMBEDDING.VECTOR_DIMENSIONS,
|
||||
settings.EMBEDDING.MAX_INPUT_TOKENS,
|
||||
settings.EMBEDDING.MAX_TOKENS_PER_REQUEST,
|
||||
settings.EMBEDDING.resolve_send_dimensions(),
|
||||
)
|
||||
|
||||
async def embed(self, query: str) -> list[float]:
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from src.routers import (
|
|||
webhooks,
|
||||
workspaces,
|
||||
)
|
||||
from src.startup import validate_embedding_schema
|
||||
from src.telemetry import (
|
||||
initialize_telemetry_async,
|
||||
metrics_endpoint,
|
||||
|
|
@ -125,6 +126,12 @@ async def lifespan(_: FastAPI):
|
|||
# Initialize CloudEvents telemetry
|
||||
await initialize_telemetry_async()
|
||||
|
||||
# Validate embedding schema before serving any traffic. Fails closed: if
|
||||
# the configured EMBEDDING_VECTOR_DIMENSIONS does not match the physical
|
||||
# pgvector columns, the process refuses to start rather than silently
|
||||
# writing wrong-dim vectors.
|
||||
await validate_embedding_schema(engine)
|
||||
|
||||
try:
|
||||
await init_cache()
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -25,12 +25,15 @@ from sqlalchemy.orm import Mapped, MappedColumn, mapped_column, relationship
|
|||
from sqlalchemy.sql import func
|
||||
from typing_extensions import override
|
||||
|
||||
from src.config import settings
|
||||
from src.utils.types import DocumentLevel, TaskType, VectorSyncState
|
||||
|
||||
from .db import Base
|
||||
|
||||
load_dotenv(override=True)
|
||||
|
||||
_VECTOR_DIM: int = settings.EMBEDDING.VECTOR_DIMENSIONS
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
|
||||
|
|
@ -278,7 +281,7 @@ class MessageEmbedding(Base):
|
|||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(1536), nullable=True)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True)
|
||||
message_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("messages.public_id", ondelete="CASCADE"), nullable=False, index=True
|
||||
)
|
||||
|
|
@ -386,7 +389,7 @@ class Document(Base):
|
|||
times_derived: Mapped[int] = mapped_column(
|
||||
Integer, nullable=False, server_default=text("1")
|
||||
)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(1536), nullable=True)
|
||||
embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True)
|
||||
source_ids: Mapped[list[str] | None] = mapped_column(
|
||||
JSONB, nullable=True, server_default=text("NULL")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,8 @@
|
|||
"""Startup-time validators that gate API/deriver boot."""
|
||||
|
||||
from src.startup.embedding_validator import (
|
||||
StartupValidationError,
|
||||
validate_embedding_schema,
|
||||
)
|
||||
|
||||
__all__ = ("StartupValidationError", "validate_embedding_schema")
|
||||
|
|
@ -0,0 +1,267 @@
|
|||
"""Startup validator for the embedding pipeline.
|
||||
|
||||
Crashes the process at boot if the configured EMBEDDING_VECTOR_DIMENSIONS does
|
||||
not match the physical pgvector schema. Replaces an earlier config-time guard
|
||||
that forbade non-1536 dims unless the operator asserted a VECTOR_STORE.MIGRATED
|
||||
flag — the schema introspection here is more accurate because it inspects
|
||||
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 scripts/configure_embeddings.py --report`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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.exceptions import HonchoException
|
||||
from src.models import Collection, Workspace
|
||||
from src.vector_store import VectorStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Embedding tables that must exist with matching dim.
|
||||
_EMBEDDING_TABLES: tuple[str, ...] = ("documents", "message_embeddings")
|
||||
|
||||
# Retry budget for transient introspection failures. Total wall time is
|
||||
# bounded so a sick DB does not hang readiness; fail-closed after exhaustion.
|
||||
_RETRY_ATTEMPTS = 3
|
||||
_RETRY_BACKOFF_SECONDS = 1.0
|
||||
|
||||
# Best-effort external sampler bounds.
|
||||
_EXTERNAL_SAMPLE_LIMIT = 10
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
|
||||
async def validate_embedding_schema(
|
||||
engine: AsyncEngine,
|
||||
*,
|
||||
app_settings: AppSettings | None = None,
|
||||
) -> None:
|
||||
"""Validate that the embedding schema matches the configured dimension.
|
||||
|
||||
Run after the DB pool is initialized and before the embedding client is
|
||||
constructed. Fails closed: any unrecoverable introspection error raises
|
||||
rather than letting the process serve traffic with an unknown state.
|
||||
"""
|
||||
s = app_settings if app_settings is not None else settings
|
||||
target_dim = s.EMBEDDING.VECTOR_DIMENSIONS
|
||||
schema = s.DB.SCHEMA
|
||||
|
||||
dims = await _introspect_pgvector_dims_with_retry(engine, schema)
|
||||
_assert_pgvector_dims_match(dims, schema=schema, target_dim=target_dim)
|
||||
|
||||
if s.VECTOR_STORE.TYPE in ("turbopuffer", "lancedb"):
|
||||
await _sample_external_namespaces(engine, target_dim=target_dim)
|
||||
|
||||
|
||||
async def _introspect_pgvector_dims_with_retry(
|
||||
engine: AsyncEngine, schema: str
|
||||
) -> dict[str, int]:
|
||||
"""Schema-qualified pg_attribute introspection with bounded retries.
|
||||
|
||||
Returns a mapping of table name -> raw ``atttypmod`` for the embedding
|
||||
columns. Fails closed on the last attempt — uncertainty is not a green
|
||||
light to serve traffic.
|
||||
"""
|
||||
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(
|
||||
engine: AsyncEngine, schema: str
|
||||
) -> dict[str, int]:
|
||||
"""Single-shot schema-qualified pg_attribute lookup.
|
||||
|
||||
The join through ``pg_class``/``pg_namespace`` lets us respect
|
||||
``DB.SCHEMA`` rather than relying on the ambient search_path.
|
||||
"""
|
||||
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'
|
||||
"""
|
||||
)
|
||||
async with engine.connect() as conn:
|
||||
result = await conn.execute(
|
||||
query,
|
||||
{"schema": schema, "tables": list(_EMBEDDING_TABLES)},
|
||||
)
|
||||
return {row.table_name: row.typmod for row in result}
|
||||
|
||||
|
||||
def _assert_pgvector_dims_match(
|
||||
dims: dict[str, int], *, schema: str, target_dim: int
|
||||
) -> None:
|
||||
expected = set(_EMBEDDING_TABLES)
|
||||
missing = expected - dims.keys()
|
||||
if missing:
|
||||
listing = ", ".join(sorted(f"{schema}.{t}.embedding" for t in missing))
|
||||
raise StartupValidationError(
|
||||
f"Required vector columns missing: {listing}."
|
||||
+ " Run `alembic upgrade head` first."
|
||||
)
|
||||
for table in sorted(expected):
|
||||
atttypmod = dims[table]
|
||||
if atttypmod == -1:
|
||||
raise StartupValidationError(
|
||||
f"{schema}.{table}.embedding has no declared vector dimension"
|
||||
+ " (unbounded typmod). Run"
|
||||
+ " `uv run python scripts/configure_embeddings.py`."
|
||||
)
|
||||
# pgvector stores the declared dim directly in atttypmod (no VARHDRSZ).
|
||||
actual = atttypmod
|
||||
if actual != target_dim:
|
||||
raise StartupValidationError(
|
||||
f"{schema}.{table}.embedding dim ({actual}) does not match"
|
||||
+ f" EMBEDDING_VECTOR_DIMENSIONS ({target_dim}). Run"
|
||||
+ " `uv run python scripts/configure_embeddings.py`"
|
||||
+ " or fix EMBEDDING_VECTOR_DIMENSIONS."
|
||||
)
|
||||
|
||||
|
||||
async def _sample_external_namespaces(engine: AsyncEngine, *, target_dim: int) -> None:
|
||||
"""Best-effort dim check across existing external-store namespaces.
|
||||
|
||||
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`` 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)
|
||||
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 or collections exist yet,"
|
||||
+ " skipping sample"
|
||||
)
|
||||
return
|
||||
|
||||
# Import lazily to avoid pulling in vector store deps when not configured.
|
||||
from src.vector_store import get_external_vector_store
|
||||
|
||||
store = get_external_vector_store()
|
||||
if store is None:
|
||||
# Settings said TYPE != pgvector but the store could not be created.
|
||||
# That is its own problem and not for this validator to swallow.
|
||||
return
|
||||
|
||||
candidates: list[str] = []
|
||||
for workspace_name in workspace_names:
|
||||
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))
|
||||
|
||||
if mismatches:
|
||||
formatted = ", ".join(f"{ns} (dim={d})" for ns, d in mismatches)
|
||||
raise StartupValidationError(
|
||||
f"Existing external-store namespaces have dim != {target_dim}:"
|
||||
+ f" {formatted}. Run"
|
||||
+ " `uv run python scripts/configure_embeddings.py --report`."
|
||||
)
|
||||
|
||||
|
||||
async def _sample_workspace_names(engine: AsyncEngine, limit: int) -> list[str]:
|
||||
"""Pull up to ``limit`` workspace names ordered by creation time.
|
||||
|
||||
Uses the ORM ``Workspace`` model so ``Base.metadata.schema`` (configured
|
||||
from ``settings.DB.SCHEMA`` in ``src/db.py``) is honored automatically —
|
||||
a non-public schema deployment must not silently sample the wrong table.
|
||||
"""
|
||||
stmt = select(Workspace.name).order_by(Workspace.created_at.desc()).limit(limit)
|
||||
async with engine.connect() as conn:
|
||||
result = await conn.execute(stmt)
|
||||
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.
|
||||
|
||||
Delegates to the store's own ``probe_namespace_dim`` implementation
|
||||
(lancedb opens the table, turbopuffer reads the schema). ``None`` means
|
||||
"lazy-create namespace, nothing to validate against."
|
||||
"""
|
||||
return await store.probe_namespace_dim(namespace)
|
||||
|
|
@ -183,6 +183,17 @@ class VectorStore(ABC):
|
|||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def probe_namespace_dim(self, namespace: str) -> int | None:
|
||||
"""
|
||||
Return the declared vector dimension of an existing namespace.
|
||||
|
||||
Returns ``None`` if the namespace does not exist yet (lazy-create
|
||||
model: not an error). Raises only when the SDK reports the
|
||||
namespace exists but its schema is unreadable.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
def _create_store_by_type(store_type: str) -> VectorStore:
|
||||
"""Create a vector store instance by type name."""
|
||||
|
|
|
|||
|
|
@ -99,7 +99,7 @@ class LanceDBVectorStore(VectorStore):
|
|||
fields.extend(self._metadata_fields_for_namespace(namespace))
|
||||
schema = pa.schema(fields)
|
||||
try:
|
||||
table = await db.create_table(namespace, schema=schema) # pyright: ignore[reportUnknownArgumentType]
|
||||
table = await db.create_table(namespace, schema=schema)
|
||||
return table
|
||||
except Exception:
|
||||
# Table may have been created by another worker, try to open it
|
||||
|
|
@ -378,3 +378,28 @@ class LanceDBVectorStore(VectorStore):
|
|||
self._db.close()
|
||||
self._db = None
|
||||
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.
|
||||
|
||||
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:
|
||||
return None
|
||||
table = await db.open_table(namespace)
|
||||
schema = await table.schema()
|
||||
for field in schema:
|
||||
if field.name == "vector" and hasattr(field.type, "list_size"):
|
||||
return int(field.type.list_size)
|
||||
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."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ This module provides a Turbopuffer-based implementation of the VectorStore inter
|
|||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
|
|
@ -87,7 +88,7 @@ class TurbopufferVectorStore(VectorStore):
|
|||
|
||||
try:
|
||||
await ns.write(
|
||||
upsert_rows=rows,
|
||||
upsert_rows=cast(Any, rows),
|
||||
distance_metric=DISTANCE_METRIC,
|
||||
)
|
||||
return
|
||||
|
|
@ -310,3 +311,46 @@ class TurbopufferVectorStore(VectorStore):
|
|||
"""Close the Turbopuffer client and release resources."""
|
||||
await self.tpuf.close()
|
||||
logger.debug("Turbopuffer client closed")
|
||||
|
||||
async def probe_namespace_dim(self, namespace: str) -> int | None:
|
||||
"""Inspect a Turbopuffer namespace schema to recover the vector dim.
|
||||
|
||||
Turbopuffer namespaces are lazy-created; ``namespace.exists()`` returns
|
||||
False before the first write. The schema response maps attribute name
|
||||
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:
|
||||
if not await ns.exists():
|
||||
return None
|
||||
except NotFoundError:
|
||||
return None
|
||||
|
||||
try:
|
||||
schema = await ns.schema()
|
||||
except NotFoundError:
|
||||
return None
|
||||
|
||||
vector_attr = schema.get("vector")
|
||||
if vector_attr is 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))
|
||||
|
|
|
|||
|
|
@ -321,6 +321,7 @@ async def fake_cache(fake_cache_session: FakeAsyncRedis):
|
|||
async def client(
|
||||
db_session: AsyncSession,
|
||||
fake_cache_session: FakeAsyncRedis, # pyright: ignore[reportUnusedParameter]
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> AsyncGenerator[TestClient, Any]:
|
||||
"""Create a FastAPI TestClient for the scope of a single test function"""
|
||||
|
||||
|
|
@ -338,6 +339,18 @@ async def client(
|
|||
yield db_session
|
||||
|
||||
app.dependency_overrides[get_db] = override_get_db
|
||||
|
||||
# No-op the startup embedding-schema validator inside the lifespan. The
|
||||
# global `engine` it would inspect points to a DB that isn't migrated in
|
||||
# CI (per-worker test DBs are migrated separately by db_engine), and we
|
||||
# don't want the validator to dispose the test engine via the lifespan
|
||||
# finally block either. The validator has its own dedicated coverage in
|
||||
# tests/startup/test_embedding_validator.py against db_engine directly.
|
||||
async def _skip_validate(_engine: object) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("src.main.validate_embedding_schema", _skip_validate)
|
||||
|
||||
with TestClient(app) as c:
|
||||
if settings.AUTH.USE_AUTH:
|
||||
# give the test client the admin JWT
|
||||
|
|
@ -509,7 +522,7 @@ def mock_vector_store(request: pytest.FixtureRequest):
|
|||
from src.vector_store import (
|
||||
VectorQueryResult,
|
||||
VectorRecord,
|
||||
_hash_namespace_components, # pyright: ignore[reportPrivateUsage]
|
||||
_hash_namespace_components,
|
||||
)
|
||||
|
||||
# Create a mock vector store that stores vectors in memory
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ from src.reconciler.sync_vectors import (
|
|||
from src.vector_store import (
|
||||
VectorRecord,
|
||||
VectorStore,
|
||||
_hash_namespace_components, # pyright: ignore[reportPrivateUsage]
|
||||
_hash_namespace_components,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -12,8 +12,16 @@ class FakeOpenAIEmbeddingsAPI:
|
|||
self.embedding: list[float] = embedding
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def create(self, *, model: str, input: str | list[str]) -> SimpleNamespace:
|
||||
self.calls.append({"model": model, "input": input})
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[str],
|
||||
**kwargs: Any,
|
||||
) -> SimpleNamespace:
|
||||
call: dict[str, Any] = {"model": model, "input": input}
|
||||
call.update(kwargs)
|
||||
self.calls.append(call)
|
||||
if isinstance(input, list):
|
||||
data = [SimpleNamespace(embedding=self.embedding) for _ in input]
|
||||
else:
|
||||
|
|
@ -45,6 +53,7 @@ async def test_openai_embedding_client_uses_configured_model_and_dimensions(
|
|||
vector_dimensions=8,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
embedding = await client.embed("hello world")
|
||||
|
|
@ -76,6 +85,7 @@ async def test_openai_embedding_client_rejects_dimension_mismatch(
|
|||
vector_dimensions=8,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Embedding dimension mismatch"):
|
||||
|
|
@ -125,6 +135,7 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
vector_dimensions=12,
|
||||
max_input_tokens=4096,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
embedding = await client.embed("hello world")
|
||||
|
|
@ -137,3 +148,194 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
"config": {"output_dimensionality": 12},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _build_openai_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
embedding: list[float],
|
||||
model: str,
|
||||
send_dimensions: bool,
|
||||
vector_dimensions: int,
|
||||
) -> tuple[_EmbeddingClient, FakeOpenAIEmbeddingsAPI]:
|
||||
fake_embeddings = FakeOpenAIEmbeddingsAPI(embedding)
|
||||
|
||||
class FakeOpenAIClient:
|
||||
def __init__(self, *, api_key: str | None, base_url: str | None) -> None:
|
||||
self.api_key: str | None = api_key
|
||||
self.base_url: str | None = base_url
|
||||
self.embeddings: FakeOpenAIEmbeddingsAPI = fake_embeddings
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.AsyncOpenAI", FakeOpenAIClient)
|
||||
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="openai",
|
||||
model=model,
|
||||
api_key="test-key",
|
||||
),
|
||||
vector_dimensions=vector_dimensions,
|
||||
max_input_tokens=8192,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=send_dimensions,
|
||||
)
|
||||
return client, fake_embeddings
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_embed_forwards_dimensions_when_send_dimensions_true(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.embed("hello")
|
||||
|
||||
assert fake.calls == [
|
||||
{
|
||||
"model": "text-embedding-3-small",
|
||||
"input": ["hello"],
|
||||
"dimensions": 768,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_embed_omits_dimensions_when_send_dimensions_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 1536,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=False,
|
||||
vector_dimensions=1536,
|
||||
)
|
||||
|
||||
await client.embed("hello")
|
||||
|
||||
assert fake.calls == [{"model": "text-embedding-3-small", "input": ["hello"]}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_simple_batch_embed_forwards_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.simple_batch_embed(["a", "b"])
|
||||
|
||||
assert len(fake.calls) == 1
|
||||
assert fake.calls[0]["dimensions"] == 768
|
||||
assert fake.calls[0]["input"] == ["a", "b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_batch_embed_forwards_dimensions(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client, fake = _build_openai_client(
|
||||
monkeypatch,
|
||||
embedding=[0.1] * 768,
|
||||
model="text-embedding-3-small",
|
||||
send_dimensions=True,
|
||||
vector_dimensions=768,
|
||||
)
|
||||
|
||||
await client.batch_embed({"a": "hello", "b": "world"})
|
||||
|
||||
assert len(fake.calls) == 1
|
||||
assert fake.calls[0]["dimensions"] == 768
|
||||
|
||||
|
||||
def _build_embedding_settings(
|
||||
env: dict[str, str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> Any:
|
||||
"""Construct a fresh EmbeddingSettings from the given env, isolated from os.environ."""
|
||||
from src.config import EmbeddingSettings
|
||||
|
||||
for key in (
|
||||
"EMBEDDING_VECTOR_DIMENSIONS",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL",
|
||||
"EMBEDDING_MODEL_CONFIG__TRANSPORT",
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE",
|
||||
):
|
||||
monkeypatch.delenv(key, raising=False)
|
||||
for key, value in env.items():
|
||||
monkeypatch.setenv(key, value)
|
||||
return EmbeddingSettings()
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_default_dim_returns_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings({}, monkeypatch)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_explicit_dim_returns_true(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings({"EMBEDDING_VECTOR_DIMENSIONS": "768"}, monkeypatch)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_auto_ada_002_returns_false(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "1536",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL": "text-embedding-ada-002",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_always_returns_true_regardless(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "always"},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_always_overrides_ada_rejecting_allowlist(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "always",
|
||||
"EMBEDDING_MODEL_CONFIG__MODEL": "text-embedding-ada-002",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is True
|
||||
|
||||
|
||||
def test_resolve_send_dimensions_never_returns_false_regardless(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
s = _build_embedding_settings(
|
||||
{
|
||||
"EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE": "never",
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "768",
|
||||
},
|
||||
monkeypatch,
|
||||
)
|
||||
assert s.resolve_send_dimensions() is False
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
|
|
@ -251,14 +250,14 @@ def test_app_settings_propagate_embedding_dimensions_to_vector_store() -> None:
|
|||
assert settings.VECTOR_STORE.DIMENSIONS == 2048
|
||||
|
||||
|
||||
def test_app_settings_require_matching_embedding_and_vector_store_dimensions() -> None:
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=re.escape(
|
||||
"VECTOR_STORE.DIMENSIONS must match EMBEDDING.VECTOR_DIMENSIONS"
|
||||
),
|
||||
):
|
||||
AppSettings(
|
||||
def test_app_settings_explicit_vector_store_dimensions_warns_and_overrides() -> None:
|
||||
"""VECTOR_STORE.DIMENSIONS is deprecated: EMBEDDING.VECTOR_DIMENSIONS wins
|
||||
and the operator gets a DeprecationWarning if they set it explicitly."""
|
||||
import warnings
|
||||
|
||||
with warnings.catch_warnings(record=True) as captured:
|
||||
warnings.simplefilter("always")
|
||||
settings = AppSettings(
|
||||
EMBEDDING=EmbeddingSettings(VECTOR_DIMENSIONS=2048),
|
||||
VECTOR_STORE=VectorStoreSettings(
|
||||
TYPE="lancedb",
|
||||
|
|
@ -266,28 +265,47 @@ def test_app_settings_require_matching_embedding_and_vector_store_dimensions() -
|
|||
DIMENSIONS=1536,
|
||||
),
|
||||
)
|
||||
messages = [
|
||||
str(w.message) for w in captured if issubclass(w.category, DeprecationWarning)
|
||||
]
|
||||
assert any(
|
||||
"VECTOR_STORE_DIMENSIONS is deprecated" in m for m in messages
|
||||
), f"expected deprecation warning, got {messages!r}"
|
||||
assert settings.EMBEDDING.VECTOR_DIMENSIONS == 2048
|
||||
assert settings.VECTOR_STORE.DIMENSIONS == 2048, (
|
||||
"EMBEDDING.VECTOR_DIMENSIONS should always overwrite the operator-supplied "
|
||||
"VECTOR_STORE.DIMENSIONS value"
|
||||
)
|
||||
|
||||
|
||||
def test_app_settings_reject_non_1536_dimensions_while_pgvector_or_dual_write_active() -> (
|
||||
None
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=re.escape("EMBEDDING.VECTOR_DIMENSIONS must remain 1536"),
|
||||
):
|
||||
AppSettings(
|
||||
EMBEDDING=EmbeddingSettings(VECTOR_DIMENSIONS=2048),
|
||||
VECTOR_STORE=VectorStoreSettings(TYPE="pgvector", MIGRATED=True),
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=re.escape("EMBEDDING.VECTOR_DIMENSIONS must remain 1536"),
|
||||
):
|
||||
AppSettings(
|
||||
EMBEDDING=EmbeddingSettings(VECTOR_DIMENSIONS=2048),
|
||||
VECTOR_STORE=VectorStoreSettings(TYPE="lancedb", MIGRATED=False),
|
||||
def test_app_settings_accepts_non_1536_with_any_vector_store_configuration() -> None:
|
||||
"""The dim-vs-MIGRATED guard was removed; the runtime startup schema
|
||||
validator (src/startup/embedding_validator.py) is the new safety net.
|
||||
Construction must succeed for every combination at config time."""
|
||||
from typing import Literal
|
||||
|
||||
combos: list[tuple[Literal["pgvector", "turbopuffer", "lancedb"], bool]] = [
|
||||
("pgvector", True),
|
||||
("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(**vs_kwargs),
|
||||
)
|
||||
assert settings.EMBEDDING.VECTOR_DIMENSIONS == 768
|
||||
assert store_type == settings.VECTOR_STORE.TYPE
|
||||
assert settings.VECTOR_STORE.MIGRATED is migrated
|
||||
|
||||
|
||||
def test_config_toml_example_uses_nested_model_config_sections() -> None:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -96,16 +96,29 @@ def ts_test_server(
|
|||
|
||||
app.dependency_overrides[get_db] = override_get_db
|
||||
|
||||
# No-op the lifespan's startup embedding-schema validator — same
|
||||
# reasoning as the `client` fixture in tests/conftest.py: the module-
|
||||
# level `engine` it inspects points to an unmigrated DB in CI, and the
|
||||
# validator has dedicated coverage in tests/startup/. Use a manual
|
||||
# MonkeyPatch since this fixture is module-scoped (the built-in
|
||||
# `monkeypatch` fixture is function-scoped only).
|
||||
async def _skip_validate(_engine: object) -> None:
|
||||
return None
|
||||
|
||||
mp = pytest.MonkeyPatch()
|
||||
mp.setattr("src.main.validate_embedding_schema", _skip_validate)
|
||||
|
||||
# Start the server
|
||||
server = TestServer(app, port)
|
||||
server.start()
|
||||
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
|
||||
# Cleanup
|
||||
server.stop()
|
||||
app.dependency_overrides.clear()
|
||||
_ts_session_factory = None
|
||||
try:
|
||||
server.start()
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
# Cleanup
|
||||
server.stop()
|
||||
app.dependency_overrides.clear()
|
||||
_ts_session_factory = None
|
||||
mp.undo()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,205 @@
|
|||
"""Phase 2: startup embedding-schema validator + VECTOR_STORE_DIMENSIONS deprecation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from src.startup.embedding_validator import (
|
||||
StartupValidationError,
|
||||
_assert_pgvector_dims_match, # pyright: ignore[reportPrivateUsage]
|
||||
validate_embedding_schema,
|
||||
)
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
|
||||
|
||||
# pgvector stores the declared dim directly in atttypmod (no VARHDRSZ offset).
|
||||
def _typmod(dim: int) -> int:
|
||||
return dim
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure-function unit tests for the dim assertion
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_assert_pgvector_dims_match_passes_when_all_dims_align() -> None:
|
||||
_assert_pgvector_dims_match(
|
||||
{"documents": _typmod(1536), "message_embeddings": _typmod(1536)},
|
||||
schema="public",
|
||||
target_dim=1536,
|
||||
)
|
||||
|
||||
|
||||
def test_assert_pgvector_dims_match_raises_on_dim_mismatch() -> None:
|
||||
with pytest.raises(StartupValidationError, match="dim .* does not match"):
|
||||
_assert_pgvector_dims_match(
|
||||
{"documents": _typmod(1536), "message_embeddings": _typmod(768)},
|
||||
schema="public",
|
||||
target_dim=1536,
|
||||
)
|
||||
|
||||
|
||||
def test_assert_pgvector_dims_match_lists_all_missing_columns() -> None:
|
||||
with pytest.raises(StartupValidationError) as excinfo:
|
||||
_assert_pgvector_dims_match(
|
||||
{"documents": _typmod(1536)},
|
||||
schema="public",
|
||||
target_dim=1536,
|
||||
)
|
||||
msg = str(excinfo.value)
|
||||
assert "message_embeddings" in msg
|
||||
assert "alembic upgrade head" in msg
|
||||
|
||||
|
||||
def test_assert_pgvector_dims_match_raises_on_unbounded_typmod() -> None:
|
||||
with pytest.raises(StartupValidationError, match="unbounded typmod"):
|
||||
_assert_pgvector_dims_match(
|
||||
{"documents": -1, "message_embeddings": _typmod(1536)},
|
||||
schema="public",
|
||||
target_dim=1536,
|
||||
)
|
||||
|
||||
|
||||
def test_assert_pgvector_dims_match_respects_non_public_schema() -> None:
|
||||
with pytest.raises(StartupValidationError, match="my_schema.documents"):
|
||||
_assert_pgvector_dims_match(
|
||||
{},
|
||||
schema="my_schema",
|
||||
target_dim=1536,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fail-closed retry behavior
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validator_fails_closed_when_introspection_keeps_failing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""After the retry budget exhausts, the validator crashes — uncertainty
|
||||
is not a green light to serve traffic."""
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def always_raise(_engine: AsyncEngine, _schema: str) -> dict[str, int]:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
raise OperationalError("SELECT 1", {}, Exception("DB unreachable"))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"src.startup.embedding_validator._introspect_pgvector_dims_once",
|
||||
always_raise,
|
||||
)
|
||||
# Make backoff effectively instant for the test.
|
||||
monkeypatch.setattr("src.startup.embedding_validator._RETRY_BACKOFF_SECONDS", 0.0)
|
||||
|
||||
with pytest.raises(StartupValidationError, match="could not validate"):
|
||||
await validate_embedding_schema(engine=AsyncMock())
|
||||
|
||||
assert call_count == 3, "should exhaust the retry budget before failing"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration: real test DB
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validator_passes_against_test_database(
|
||||
db_engine: AsyncEngine,
|
||||
) -> None:
|
||||
"""The test DB is provisioned at the default dim (1536); the validator
|
||||
should accept it without raising."""
|
||||
await validate_embedding_schema(db_engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validator_raises_when_schema_dim_diverges_from_settings(
|
||||
db_engine: AsyncEngine,
|
||||
) -> None:
|
||||
"""ALTER one of the embedding columns to a non-1536 dim and confirm the
|
||||
validator raises with an actionable message."""
|
||||
async with db_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
"ALTER TABLE documents"
|
||||
+ " ALTER COLUMN embedding TYPE vector(768) USING NULL"
|
||||
)
|
||||
)
|
||||
try:
|
||||
with pytest.raises(StartupValidationError, match="dim .* does not match"):
|
||||
await validate_embedding_schema(db_engine)
|
||||
finally:
|
||||
async with db_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
"ALTER TABLE documents"
|
||||
+ " ALTER COLUMN embedding TYPE vector(1536) USING NULL"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# VECTOR_STORE_DIMENSIONS deprecation + dim-vs-MIGRATED guard removal
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_vector_store_dimensions_explicit_set_warns(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Setting VECTOR_STORE_DIMENSIONS explicitly should trigger a deprecation
|
||||
warning. EMBEDDING_VECTOR_DIMENSIONS remains authoritative."""
|
||||
monkeypatch.setenv("PYTHON_DOTENV_DISABLED", "1")
|
||||
monkeypatch.setenv("VECTOR_STORE_DIMENSIONS", "1536")
|
||||
with warnings.catch_warnings(record=True) as captured:
|
||||
warnings.simplefilter("always")
|
||||
from src.config import AppSettings
|
||||
|
||||
AppSettings()
|
||||
messages = [
|
||||
str(w.message) for w in captured if issubclass(w.category, DeprecationWarning)
|
||||
]
|
||||
assert any(
|
||||
"VECTOR_STORE_DIMENSIONS is deprecated" in m for m in messages
|
||||
), f"expected deprecation warning, got {messages!r}"
|
||||
|
||||
|
||||
def test_non_1536_pgvector_without_migrated_no_longer_raises_at_config_time() -> None:
|
||||
"""Phase 2 removed the dim-vs-MIGRATED guard. Constructing AppSettings
|
||||
with non-1536 + default pgvector + MIGRATED=false should now succeed
|
||||
(the runtime schema validator at startup is the new safety net)."""
|
||||
env = {
|
||||
**os.environ,
|
||||
"PYTHON_DOTENV_DISABLED": "1",
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "768",
|
||||
}
|
||||
# Use a subprocess so the global settings singleton in this test
|
||||
# process is not perturbed and is re-evaluated freshly in the child.
|
||||
snippet = (
|
||||
"from src.config import AppSettings\n"
|
||||
"s = AppSettings()\n"
|
||||
"print(s.EMBEDDING.VECTOR_DIMENSIONS, s.VECTOR_STORE.TYPE, s.VECTOR_STORE.MIGRATED)\n"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", snippet],
|
||||
env=env,
|
||||
cwd=str(_PROJECT_ROOT),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
last_line = result.stdout.strip().splitlines()[-1]
|
||||
assert last_line == "768 pgvector False"
|
||||
|
|
@ -497,7 +497,7 @@ class TestErrorHandlingAndRecovery:
|
|||
# parse_datetime_iso should raise appropriate errors
|
||||
if error_case is None or not isinstance(error_case, str):
|
||||
with pytest.raises(ValueError):
|
||||
parse_datetime_iso(error_case)
|
||||
parse_datetime_iso(error_case) # pyright: ignore[reportArgumentType]
|
||||
|
||||
# Filter validation should return None for invalid types
|
||||
if isinstance(error_case, str) or error_case is None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,58 @@
|
|||
"""Phase 1: verify src/models.py honors EMBEDDING_VECTOR_DIMENSIONS at import time."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
|
||||
|
||||
|
||||
def _run_in_fresh_interpreter(env_overrides: dict[str, str]) -> dict[str, int]:
|
||||
"""Import src.models in a fresh interpreter and return the vector dims.
|
||||
|
||||
A subprocess is required because src.models reads
|
||||
settings.EMBEDDING.VECTOR_DIMENSIONS at module import time to construct
|
||||
SQLAlchemy column types — reloading the module in-process would conflict
|
||||
with the existing Base.registry from earlier imports.
|
||||
|
||||
PYTHON_DOTENV_DISABLED=1 prevents config.py:20 from reloading the
|
||||
developer's .env file (which calls load_dotenv with override=True)
|
||||
and clobbering our test overrides.
|
||||
"""
|
||||
env: dict[str, str] = {
|
||||
**os.environ,
|
||||
"PYTHON_DOTENV_DISABLED": "1",
|
||||
**env_overrides,
|
||||
}
|
||||
snippet = (
|
||||
"import json\n"
|
||||
"from src.models import Document, MessageEmbedding\n"
|
||||
"print(json.dumps({\n"
|
||||
" 'message_embedding_dim': MessageEmbedding.__table__.c.embedding.type.dim,\n"
|
||||
" 'document_dim': Document.__table__.c.embedding.type.dim,\n"
|
||||
"}))\n"
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", snippet],
|
||||
env=env,
|
||||
cwd=str(_PROJECT_ROOT),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
parsed: dict[str, int] = json.loads(result.stdout.strip().splitlines()[-1])
|
||||
return parsed
|
||||
|
||||
|
||||
def test_models_uses_default_1536_when_no_env_override() -> None:
|
||||
dims = _run_in_fresh_interpreter({})
|
||||
assert dims == {"message_embedding_dim": 1536, "document_dim": 1536}
|
||||
|
||||
|
||||
def test_models_honors_explicit_embedding_vector_dimensions() -> None:
|
||||
dims = _run_in_fresh_interpreter({"EMBEDDING_VECTOR_DIMENSIONS": "768"})
|
||||
assert dims == {"message_embedding_dim": 768, "document_dim": 768}
|
||||
|
|
@ -0,0 +1,72 @@
|
|||
"""Per-store namespace dim probe tests.
|
||||
|
||||
LanceDB has an embedded driver we can spin up in a tmp dir, so we exercise
|
||||
the real probe end-to-end. Turbopuffer needs a network + API key, so it is
|
||||
covered only by static analysis + the parsing test below.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
from src.vector_store.lancedb import LanceDBVectorStore
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lancedb_probe_returns_declared_dim(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: object
|
||||
) -> None:
|
||||
"""Create a real LanceDB table at dim 768, confirm the probe recovers it."""
|
||||
monkeypatch.setattr("src.config.settings.VECTOR_STORE.LANCEDB_PATH", str(tmp_path))
|
||||
|
||||
store = LanceDBVectorStore()
|
||||
try:
|
||||
db = await store._get_db() # pyright: ignore[reportPrivateUsage]
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.string()),
|
||||
pa.field("vector", pa.list_(pa.float32(), 768)),
|
||||
]
|
||||
)
|
||||
await db.create_table("probe_test", schema=schema)
|
||||
|
||||
dim = await store.probe_namespace_dim("probe_test")
|
||||
assert dim == 768
|
||||
finally:
|
||||
await store.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lancedb_probe_returns_none_for_missing_namespace(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: object
|
||||
) -> None:
|
||||
"""Lazy-create model: probing a nonexistent table is not an error."""
|
||||
monkeypatch.setattr("src.config.settings.VECTOR_STORE.LANCEDB_PATH", str(tmp_path))
|
||||
|
||||
store = LanceDBVectorStore()
|
||||
try:
|
||||
dim = await store.probe_namespace_dim("does_not_exist")
|
||||
assert dim is None
|
||||
finally:
|
||||
await store.close()
|
||||
|
||||
|
||||
def test_turbopuffer_vector_dim_regex_extracts_dim_from_type_string() -> None:
|
||||
"""Turbopuffer's attribute type for a vector column is a bracket-prefixed
|
||||
dim with a width suffix: ``[768]f32``, ``[1536]f16``, ``[256]i8``. The
|
||||
probe extracts the integer inside the brackets. Lock the format here so
|
||||
an SDK change is loud."""
|
||||
pattern = re.compile(r"\[(\d+)\]")
|
||||
cases = {
|
||||
"[768]f32": "768",
|
||||
"[1536]f16": "1536",
|
||||
"[256]i8": "256",
|
||||
}
|
||||
for type_str, expected in cases.items():
|
||||
match = pattern.search(type_str)
|
||||
assert match is not None, f"failed to match {type_str!r}"
|
||||
assert match.group(1) == expected
|
||||
assert pattern.search("string") is None
|
||||
Loading…
Reference in New Issue