From 17f9769eb47f685475a61f23b6e7c722a9b9d372 Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Wed, 20 May 2026 16:15:51 +0200 Subject: [PATCH 1/4] fix: validate retain embedding dimensions --- .../engine/retain/embedding_utils.py | 29 +++++++++++- .../tests/test_retain_orchestrator_mapping.py | 46 +++++++++++++++++++ 2 files changed, 73 insertions(+), 2 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py index 174598b3cd..fc794218ae 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py @@ -15,11 +15,23 @@ class EmbeddingsBackend(Protocol): """Minimal duck-typed surface used by retain/recall — the concrete `Embeddings` ABC supplies default implementations that delegate to `encode()`.""" + @property + def dimension(self) -> int: ... + def encode_query(self, texts: list[str]) -> list[list[float]]: ... def encode_documents(self, texts: list[str]) -> list[list[float]]: ... +def _validate_embedding_vector(vector: list[float], *, index: int, expected_dimension: int) -> list[float]: + actual_dimension = len(vector) + if actual_dimension == 0: + raise RuntimeError(f"embedding {index} has dimension 0; expected {expected_dimension}") + if actual_dimension != expected_dimension: + raise RuntimeError(f"embedding {index} has dimension {actual_dimension}; expected {expected_dimension}") + return vector + + def generate_embedding( embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document" ) -> list[float]: @@ -36,10 +48,20 @@ def generate_embedding( """ try: embeddings = _encode_with_input_type(embeddings_backend, [text], input_type) - return embeddings[0] except Exception as e: raise Exception(f"Failed to generate embedding: {str(e)}") + if len(embeddings) != 1: + raise RuntimeError( + f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; " + "expected exact 1:1 alignment" + ) + return _validate_embedding_vector( + embeddings[0], + index=0, + expected_dimension=embeddings_backend.dimension, + ) + def _encode_with_input_type( embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType @@ -81,4 +103,7 @@ async def generate_embeddings_batch( "expected exact 1:1 alignment" ) - return embeddings + return [ + _validate_embedding_vector(embedding, index=index, expected_dimension=embeddings_backend.dimension) + for index, embedding in enumerate(embeddings) + ] diff --git a/hindsight-api-slim/tests/test_retain_orchestrator_mapping.py b/hindsight-api-slim/tests/test_retain_orchestrator_mapping.py index e6a51cecdf..63898ae1e0 100644 --- a/hindsight-api-slim/tests/test_retain_orchestrator_mapping.py +++ b/hindsight-api-slim/tests/test_retain_orchestrator_mapping.py @@ -90,6 +90,35 @@ def test_unit_ids_assigned_by_processed_fact_position(self): assert result == [["u-a1"], ["u-b1", "u-b2"]] +class TestEmbeddingSingleValidation: + def test_generate_embedding_preserves_validation_runtime_error(self): + backend = MagicMock() + backend.dimension = 3 + backend.encode_documents.return_value = [[]] + + with pytest.raises(RuntimeError, match="embedding 0 has dimension 0; expected 3"): + embedding_utils.generate_embedding(backend, "a") + + def test_generate_embedding_raises_when_backend_returns_wrong_count(self): + backend = MagicMock() + backend.dimension = 3 + backend.encode_documents.return_value = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]] + + with pytest.raises(RuntimeError, match="returned 2 vectors for 1 input text"): + embedding_utils.generate_embedding(backend, "a") + + def test_generate_embedding_uses_query_encoder_when_requested(self): + backend = MagicMock() + backend.dimension = 2 + backend.encode_query.return_value = [[0.1, 0.2]] + + result = embedding_utils.generate_embedding(backend, "a", input_type="query") + + assert result == [0.1, 0.2] + backend.encode_query.assert_called_once_with(["a"]) + backend.encode_documents.assert_not_called() + + class TestEmbeddingsBatchLengthGuarantee: def test_raises_when_backend_returns_fewer_embeddings(self): # Regression for #1037: backends that silently truncate must not pass @@ -110,8 +139,25 @@ def test_raises_when_backend_returns_more_embeddings(self): def test_passes_through_aligned_embeddings(self): backend = MagicMock() + backend.dimension = 1 backend.encode_documents.return_value = [[0.1], [0.2]] result = asyncio.run(embedding_utils.generate_embeddings_batch(backend, ["a", "b"])) assert result == [[0.1], [0.2]] + + def test_raises_when_backend_returns_empty_embedding_vector(self): + backend = MagicMock() + backend.dimension = 3 + backend.encode_documents.return_value = [[0.1, 0.2, 0.3], []] + + with pytest.raises(RuntimeError, match="embedding 1 has dimension 0; expected 3"): + asyncio.run(embedding_utils.generate_embeddings_batch(backend, ["a", "b"])) + + def test_raises_when_backend_returns_wrong_embedding_dimension(self): + backend = MagicMock() + backend.dimension = 3 + backend.encode_documents.return_value = [[0.1, 0.2, 0.3], [0.4, 0.5]] + + with pytest.raises(RuntimeError, match="embedding 1 has dimension 2; expected 3"): + asyncio.run(embedding_utils.generate_embeddings_batch(backend, ["a", "b"])) From 69da21e8bb000b8be34d647da01cb87b71318eda Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Wed, 20 May 2026 17:32:06 +0200 Subject: [PATCH 2/4] test: cover consolidation embedding dimension validation --- ...test_consolidation_embedding_validation.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 hindsight-api-slim/tests/test_consolidation_embedding_validation.py diff --git a/hindsight-api-slim/tests/test_consolidation_embedding_validation.py b/hindsight-api-slim/tests/test_consolidation_embedding_validation.py new file mode 100644 index 0000000000..ec4d54d7fc --- /dev/null +++ b/hindsight-api-slim/tests/test_consolidation_embedding_validation.py @@ -0,0 +1,41 @@ +import uuid + +import pytest + +from hindsight_api.engine.consolidation import consolidator + + +class _ZeroLengthEmbeddings: + dimension = 384 + + def encode(self, texts): + assert texts == ["Consolidated observation text."] + return [[]] + + +class _FakeMemoryEngine: + embeddings = _ZeroLengthEmbeddings() + + +class _FailingConn: + async def fetchrow(self, *args, **kwargs): + raise AssertionError("zero-length embedding should be rejected before database insert") + + +@pytest.mark.asyncio +async def test_create_observation_rejects_zero_length_embedding_before_insert(monkeypatch): + source_id = uuid.uuid4() + + async def fake_filter_live_source_memories(conn, bank_id, source_memory_ids): + return source_memory_ids + + monkeypatch.setattr(consolidator, "_filter_live_source_memories", fake_filter_live_source_memories) + + with pytest.raises(RuntimeError, match="embedding 0 has dimension 0; expected 384"): + await consolidator._create_observation_directly( + conn=_FailingConn(), + memory_engine=_FakeMemoryEngine(), + bank_id="test-bank", + source_memory_ids=[source_id], + observation_text="Consolidated observation text.", + ) From ba03c449df203554b25ee96f2b2a47e21c0fd1cc Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Thu, 28 May 2026 19:53:17 +0200 Subject: [PATCH 3/4] test: align consolidation embedding fake with document encoder --- .../tests/test_consolidation_embedding_validation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hindsight-api-slim/tests/test_consolidation_embedding_validation.py b/hindsight-api-slim/tests/test_consolidation_embedding_validation.py index ec4d54d7fc..0b7c13e800 100644 --- a/hindsight-api-slim/tests/test_consolidation_embedding_validation.py +++ b/hindsight-api-slim/tests/test_consolidation_embedding_validation.py @@ -8,7 +8,7 @@ class _ZeroLengthEmbeddings: dimension = 384 - def encode(self, texts): + def encode_documents(self, texts): assert texts == ["Consolidated observation text."] return [[]] From 496ea737f1897be7c49b797ee5088fe02fc3ee51 Mon Sep 17 00:00:00 2001 From: ai-ag2026 <261867348+ai-ag2026@users.noreply.github.com> Date: Thu, 28 May 2026 20:01:37 +0200 Subject: [PATCH 4/4] style: format embedding validation error message --- .../hindsight_api/engine/retain/embedding_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py index fc794218ae..55ff2d9048 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/embedding_utils.py @@ -53,8 +53,7 @@ def generate_embedding( if len(embeddings) != 1: raise RuntimeError( - f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; " - "expected exact 1:1 alignment" + f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; expected exact 1:1 alignment" ) return _validate_embedding_vector( embeddings[0],