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..55ff2d9048 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,19 @@ 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 +102,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_consolidation_embedding_validation.py b/hindsight-api-slim/tests/test_consolidation_embedding_validation.py new file mode 100644 index 0000000000..0b7c13e800 --- /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_documents(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.", + ) 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"]))