Skip to content

fix: validate embedding dimensions before pgvector writes - #1670

Merged
nicoloboschi merged 4 commits into
vectorize-io:mainfrom
ai-ag2026:fix/validate-retain-embedding-dimensions
Jun 5, 2026
Merged

fix: validate embedding dimensions before pgvector writes#1670
nicoloboschi merged 4 commits into
vectorize-io:mainfrom
ai-ag2026:fix/validate-retain-embedding-dimensions

Conversation

@ai-ag2026

@ai-ag2026 ai-ag2026 commented May 20, 2026

Copy link
Copy Markdown
Contributor

Summary

  • validate embedding vectors before they can reach pgvector writes
  • reject empty vectors and vectors whose length does not match the backend's declared dimension
  • cover both retain and consolidation observation creation paths with regression tests

Root cause

Retain embedding generation already guaranteed that the number of returned vectors matched the number of input texts, but it did not validate per-vector dimensions. If an embedding backend returned [] or another wrong-sized vector, the invalid value could flow downstream until PostgreSQL/pgvector raised different vector dimensions 384 and 0.

The same shared embedding utility is also used by consolidation when creating observations, so the validation boundary protects both retain writes and consolidation-created observations.

Related reconnaissance

Test plan

  • RED on current origin/main: uv run pytest tests/test_consolidation_embedding_validation.py::test_create_observation_rejects_zero_length_embedding_before_insert -q -o addopts= failed because the zero-length vector reached the DB insert path.
  • uv run pytest tests/test_consolidation_embedding_validation.py::test_create_observation_rejects_zero_length_embedding_before_insert tests/test_retain_orchestrator_mapping.py -q -o addopts= → 11 passed
  • uv run ruff check tests/test_consolidation_embedding_validation.py hindsight_api/engine/retain/embedding_utils.py tests/test_retain_orchestrator_mapping.py → passed
  • Added-line secret/private marker scan → 0 findings

@ai-ag2026
ai-ag2026 marked this pull request as ready for review May 20, 2026 14:22
@ai-ag2026 ai-ag2026 changed the title fix: validate retain embedding dimensions fix: validate embedding dimensions before pgvector writes May 20, 2026

@nicoloboschi nicoloboschi left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A couple of must-fix items before this lands:

1. The new RuntimeError is swallowed in generate_embedding(). The single-text path still has except Exception as e: raise Exception(f"Failed to generate embedding: {str(e)}") wrapping the body, and _validate_embedding_vector is called inside that try. So a precise RuntimeError("embedding 0 has dimension 0; expected 384") becomes a plain Exception("Failed to generate embedding: embedding 0 …") — callers can't except RuntimeError to distinguish validation failures from backend transport errors. The batch path is fine because the list comprehension runs after the try/except. Easiest fix: move validation outside the try, or re-raise with raise / raise … from e without reclassifying.

2. No test coverage for generate_embedding() (single-text). All three new unit tests hit generate_embeddings_batch. The function you just modified has zero coverage of the new validation branch — which is also why issue #1 is invisible right now.

A few non-blocking nits I'd consider while you're in here:

  • embeddings_backend is untyped and defensively probed. The Embeddings ABC in engine/embeddings.py already guarantees a dimension: int property; the try/except around getattr(..., None) plus isinstance(..., int) and > 0 is paper over the missing type annotation. Typing the parameter and reading .dimension directly would be the right shape.
  • The consolidation test (test_consolidation_embedding_validation.py) calls the private _create_observation_directly, monkey-patches the private _filter_live_source_memories, fakes MemoryEngine with a bare class, and uses a _FailingConn that asserts on fetchrow. The validation itself is already covered by the unit tests, and "consolidation uses the validating helper" is visible by inspection — this test will break on any consolidator refactor without catching anything the unit tests miss.
  • Error message lacks input context: "embedding 1 has dimension 2; expected 3" tells you the batch offset but not the text, call site, or bank — correlating back in prod will require log fishing. A logger.error(...) with bank/operation/text length before raising would help.

@ai-ag2026
ai-ag2026 force-pushed the fix/validate-retain-embedding-dimensions branch from e99aa4d to ba03c44 Compare May 28, 2026 17:53
@ai-ag2026

Copy link
Copy Markdown
Contributor Author

Addressed the review feedback on the latest head (496ea737):

  • Moved single-text generate_embedding() validation outside the backend-call try/except, so validation failures now preserve their RuntimeError type instead of being reclassified as a generic Exception.
  • Added single-text coverage for:
    • preserving validation RuntimeError on zero-length vectors
    • backend/vector count mismatch
    • input_type="query" using the query encoder path
  • Rebased the branch onto current origin/main; the PR now reads back as mergeable.
  • Typed the retain embedding backend protocol with the dimension property and validate against that directly.

Validation:

  • uv run pytest tests/test_retain_orchestrator_mapping.py tests/test_consolidation_embedding_validation.py -q → 14 passed
  • uv run ruff check hindsight_api/engine/retain/embedding_utils.py tests/test_retain_orchestrator_mapping.py tests/test_consolidation_embedding_validation.py → passed
  • uv run ruff format --check hindsight_api/engine/retain/embedding_utils.py tests/test_retain_orchestrator_mapping.py tests/test_consolidation_embedding_validation.py → passed
  • Hosted CI on the updated head → all required non-skipped jobs pass, including verify-generated-files, build-api-python-versions 3.11/3.12/3.13/3.14, check-openapi-compatibility, and check-cli-coverage.

Note: the full ./scripts/hooks/lint.sh path attempted to run ESLint without the control-plane npm dependencies installed in this fresh worktree (@eslint/js missing). The touched files are Python-only, so I ran the focused Ruff checks above and let hosted CI verify generated-file cleanliness.

@nicoloboschi
nicoloboschi merged commit 06c88e0 into vectorize-io:main Jun 5, 2026
61 checks passed
nicoloboschi added a commit that referenced this pull request Jun 5, 2026
PR #1670 added post-encode dimension validation to generate_embeddings_batch
— it now reads embeddings_backend.dimension, which the EmbeddingsBackend
Protocol already requires. The pre-existing QueryAwareEmbeddings/
DocumentAwareEmbeddings routing mocks (#1770) omit it, so the two routing
tests started failing with AttributeError on main.

The mocks return single-element vectors, so declare dimension = 1 to satisfy
the Protocol and let validation pass. Pure test fix; no behavior change.
nicoloboschi added a commit that referenced this pull request Jun 5, 2026
…ients (#2009)

* fix(openapi): keep binary upload fields as format:binary; regen spec+clients

The #1982 dep bump (FastAPI 0.136 / Pydantic 2.12) serializes binary upload
fields as OpenAPI-3.1 {"type":"string","contentMediaType":"application/
octet-stream"}. openapi-generator v7.10.0 (generate-clients.sh) does NOT
treat contentMediaType as a file upload, so it regenerated the Files `files`
and document-transfer `file` params as plain strings — silently breaking
multipart upload in the Go/Python/TypeScript clients ([]*os.File -> []string,
StrictBytes -> StrictStr, Blob|File -> string).

generate_openapi.py now post-processes the exported schema to restore the
prior `format: binary` representation (still valid under openapi 3.1.0, and
what the generator understands) for application/octet-stream string fields,
scoped to binary uploads only. Regenerated the spec and clients: the upload
signatures are back to the file-upload form (identical to main); the only
remaining delta vs main is ValidationError dropping its `url` field, a real
Pydantic 2.12 change (error metadata, harmless).

* test(embeddings): give zeroentropy routing mocks a dimension attribute

PR #1670 added post-encode dimension validation to generate_embeddings_batch
— it now reads embeddings_backend.dimension, which the EmbeddingsBackend
Protocol already requires. The pre-existing QueryAwareEmbeddings/
DocumentAwareEmbeddings routing mocks (#1770) omit it, so the two routing
tests started failing with AttributeError on main.

The mocks return single-element vectors, so declare dimension = 1 to satisfy
the Protocol and let validation pass. Pure test fix; no behavior change.

* test(openapi): lock _restore_binary_format binary-upload rewrite

Regression guard for the file-upload break: asserts octet-stream string
fields are rewritten to format:binary (incl. nested/array-item schemas) and
that other content media types are left untouched.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

batch_retain can pass zero-length embeddings to pgvector (vector(384) vs vector(0))

2 participants