fix: validate embedding dimensions before pgvector writes#1670
fix: validate embedding dimensions before pgvector writes#1670ai-ag2026 wants to merge 2 commits into
Conversation
nicoloboschi
left a comment
There was a problem hiding this comment.
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_backendis untyped and defensively probed. TheEmbeddingsABC inengine/embeddings.pyalready guarantees adimension: intproperty; thetry/exceptaroundgetattr(..., None)plusisinstance(..., int) and > 0is paper over the missing type annotation. Typing the parameter and reading.dimensiondirectly 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, fakesMemoryEnginewith a bare class, and uses a_FailingConnthat asserts onfetchrow. 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. Alogger.error(...)with bank/operation/text length before raising would help.
Summary
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 raiseddifferent 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
/healthis green.Test plan
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 passeduv run ruff check tests/test_consolidation_embedding_validation.py hindsight_api/engine/retain/embedding_utils.py tests/test_retain_orchestrator_mapping.py→ passed