diff --git a/CHANGELOG.md b/CHANGELOG.md index 49662fe..a5d047b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **Parameterized LIMIT clauses**: `get_reads`, `get_actions`, and `get_activity` now use bind parameters (`%s`) for LIMIT values instead of f-string interpolation, eliminating a fragile SQL construction pattern (MEDIUM #3) ### Fixed +- **Ollama response validation**: `OllamaEmbedding.embed()` now validates that the number of returned embeddings matches the number of input texts, raising `ValueError` on partial responses instead of silently dropping entries via `zip(strict=False)` (MEDIUM #21) - **Alert expiry filter**: `get_active_alerts` and `get_all_active_alerts` now filter out expired alerts (`expires > NOW()`), matching the behavior of `get_active_suppressions` (MEDIUM #18) - **Intention lifecycle**: `generate_briefing` now transitions fired intentions from "pending" to "fired" state, preventing them from firing on every subsequent briefing read - **Custom prompt sync uses DEFAULT_OWNER**: `_sync_custom_prompts` now queries `DEFAULT_OWNER` instead of the request-scoped `_owner_id()`, preventing User A's prompt sync from leaking into User B's prompt registry in multi-tenant deployments (MEDIUM #14) diff --git a/src/mcp_awareness/embeddings.py b/src/mcp_awareness/embeddings.py index a14ee85..8032bee 100644 --- a/src/mcp_awareness/embeddings.py +++ b/src/mcp_awareness/embeddings.py @@ -118,6 +118,15 @@ def embed(self, texts: list[str]) -> list[list[float]]: with urllib.request.urlopen(req, timeout=self._timeout) as resp: data = json.loads(resp.read()) result: list[list[float]] = data["embeddings"] + if len(result) != len(texts): + logger.warning( + "Ollama returned %d embeddings for %d texts — partial response", + len(result), + len(texts), + ) + raise ValueError( + f"Ollama returned {len(result)} embeddings for {len(texts)} input texts" + ) return result def is_available(self) -> bool: diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py index 090e626..72fed23 100644 --- a/tests/test_embeddings.py +++ b/tests/test_embeddings.py @@ -18,7 +18,9 @@ from __future__ import annotations +import json import os +from unittest.mock import MagicMock, patch import pytest @@ -203,6 +205,35 @@ def test_unreachable_is_not_available(self): p = OllamaEmbedding(base_url="http://localhost:19999") assert p.is_available() is False + def test_partial_response_raises(self): + """embed() raises ValueError when Ollama returns fewer embeddings than inputs.""" + p = OllamaEmbedding(base_url="http://localhost:19999") + # Simulate Ollama returning only 1 embedding for 3 input texts + partial_body = json.dumps({"embeddings": [[0.1] * 768]}).encode() + mock_resp = MagicMock() + mock_resp.read.return_value = partial_body + mock_resp.__enter__ = lambda s: s + mock_resp.__exit__ = MagicMock(return_value=False) + + with ( + patch("urllib.request.urlopen", return_value=mock_resp), + pytest.raises(ValueError, match=r"1 embeddings for 3 input texts"), + ): + p.embed(["text one", "text two", "text three"]) + + def test_correct_response_count_succeeds(self): + """embed() returns normally when embedding count matches input count.""" + p = OllamaEmbedding(base_url="http://localhost:19999") + good_body = json.dumps({"embeddings": [[0.1] * 768, [0.2] * 768]}).encode() + mock_resp = MagicMock() + mock_resp.read.return_value = good_body + mock_resp.__enter__ = lambda s: s + mock_resp.__exit__ = MagicMock(return_value=False) + + with patch("urllib.request.urlopen", return_value=mock_resp): + result = p.embed(["text one", "text two"]) + assert len(result) == 2 + # --------------------------------------------------------------------------- # create_provider