diff --git a/src/basic_memory/repository/fastembed_provider.py b/src/basic_memory/repository/fastembed_provider.py index 7e14cb3f0..f9804a02b 100644 --- a/src/basic_memory/repository/fastembed_provider.py +++ b/src/basic_memory/repository/fastembed_provider.py @@ -276,6 +276,20 @@ async def _load_model(self) -> "TextEmbedding": ) return self._model + def _normalize_vectors(self, vectors) -> list[list[float]]: + # sqlite_search_repository.py uses a distance-to-similarity formula that assumes + # unit-normalized vectors (see the comment on line 65-67 of that file). + # Some models (e.g. multilingual ones) return vectors with norm > 1, so we + # L2-normalize here to satisfy that contract regardless of the chosen model. + normalized: list[list[float]] = [] + for vector in vectors: + values = vector.tolist() if hasattr(vector, "tolist") else list(vector) + norm = math.sqrt(sum(x * x for x in values)) + if norm > 0: + values = [x / norm for x in values] + normalized.append([float(v) for v in values]) + return normalized + async def embed_documents(self, texts: list[str]) -> list[list[float]]: if not texts: return [] @@ -297,19 +311,7 @@ def _embed_batch() -> list[list[float]]: embed_kwargs: dict[str, int] = {"batch_size": self.batch_size} if effective_parallel is not None: embed_kwargs["parallel"] = effective_parallel - vectors = list(model.embed(texts, **embed_kwargs)) - # sqlite_search_repository.py uses a distance-to-similarity formula that assumes - # unit-normalized vectors (see the comment on line 65-67 of that file). - # Some models (e.g. multilingual ones) return vectors with norm > 1, so we - # L2-normalize here to satisfy that contract regardless of the chosen model. - normalized: list[list[float]] = [] - for vector in vectors: - values = vector.tolist() if hasattr(vector, "tolist") else list(vector) - norm = math.sqrt(sum(x * x for x in values)) - if norm > 0: - values = [x / norm for x in values] - normalized.append([float(v) for v in values]) - return normalized + return self._normalize_vectors(model.embed(texts, **embed_kwargs)) vectors = await asyncio.to_thread(_embed_batch) if vectors and len(vectors[0]) != self.dimensions: @@ -320,5 +322,19 @@ def _embed_batch() -> list[list[float]]: return vectors async def embed_query(self, text: str) -> list[float]: - vectors = await self.embed_documents([text]) + model = await self._load_model() + + # Asymmetric models (e.g. bge-small-en-v1.5) apply a query-specific + # instruction inside query_embed; embed() is the passage/document path and + # silently loses that asymmetry for queries (#1264). query_embed takes no + # batch/parallel kwargs — it embeds a single query string. + def _embed_query() -> list[list[float]]: + return self._normalize_vectors(model.query_embed(text)) + + vectors = await asyncio.to_thread(_embed_query) + if vectors and len(vectors[0]) != self.dimensions: + raise RuntimeError( + f"Embedding model returned {len(vectors[0])}-dimensional vectors " + f"but provider was configured for {self.dimensions} dimensions." + ) return vectors[0] if vectors else [0.0] * self.dimensions diff --git a/tests/repository/test_fastembed_provider.py b/tests/repository/test_fastembed_provider.py index 703e32d5e..65101457a 100644 --- a/tests/repository/test_fastembed_provider.py +++ b/tests/repository/test_fastembed_provider.py @@ -24,6 +24,7 @@ class _StubTextEmbedding: init_count = 0 last_init_kwargs: dict[str, Any] = {} last_embed_kwargs: dict[str, Any] = {} + last_query_embed_inputs: list[str] = [] def __init__( self, @@ -34,6 +35,7 @@ def __init__( ): self.model_name = model_name self.embed_calls = 0 + self.query_embed_calls = 0 _StubTextEmbedding.last_init_kwargs = { "model_name": model_name, "cache_dir": cache_dir, @@ -51,6 +53,15 @@ def embed(self, texts: list[str], batch_size: int = 64, **kwargs): else: yield _StubVector([1.0, 0.0, 0.0, 0.0]) + def query_embed(self, query): + self.query_embed_calls += 1 + inputs = [query] if isinstance(query, str) else list(query) + _StubTextEmbedding.last_query_embed_inputs = inputs + for _ in inputs: + # Distinct from the embed() passage vector so tests can tell the + # query path apart from the document path. + yield _StubVector([0.0, 1.0, 0.0, 0.0]) + @pytest.mark.asyncio async def test_fastembed_provider_lazy_loads_and_reuses_model(monkeypatch): @@ -205,6 +216,10 @@ def embed(self, texts: list[str], **_kwargs): for _ in texts: yield _UnormalizedVector([1.5, 2.0, 1.0, 0.5]) + def query_embed(self, query): + for _ in [query] if isinstance(query, str) else list(query): + yield _UnormalizedVector([1.5, 2.0, 1.0, 0.5]) + @pytest.mark.asyncio async def test_fastembed_provider_l2_normalizes_output_vectors(monkeypatch): @@ -226,6 +241,58 @@ async def test_fastembed_provider_l2_normalizes_output_vectors(monkeypatch): assert abs(norm - 1.0) < 1e-6, f"Expected unit norm, got {norm}" +@pytest.mark.asyncio +async def test_fastembed_provider_embed_query_uses_query_embed(monkeypatch): + """embed_query must take FastEmbed's query_embed path, not the passage embed path (#1264). + + Asymmetric models (e.g. bge-small-en-v1.5) apply a query-specific instruction + inside query_embed; embedding a query through the document path silently loses + that asymmetry and degrades retrieval quality. + """ + module = type(sys)("fastembed") + setattr(module, "TextEmbedding", _StubTextEmbedding) + monkeypatch.setitem(sys.modules, "fastembed", module) + _StubTextEmbedding.last_embed_kwargs = {} + _StubTextEmbedding.last_query_embed_inputs = [] + + provider = FastEmbedEmbeddingProvider(model_name="stub-model", dimensions=4) + vector = await provider.embed_query("auth query") + + assert provider._model is not None + assert _StubTextEmbedding.last_query_embed_inputs == ["auth query"] + assert _StubTextEmbedding.last_embed_kwargs == {} + # The stub's query vector is [0, 1, 0, 0], distinct from its passage vector, + # so this assertion fails if the query is embedded via embed(). + assert vector == [0.0, 1.0, 0.0, 0.0] + + +@pytest.mark.asyncio +async def test_fastembed_provider_embed_query_normalizes_and_checks_dimensions(monkeypatch): + """Query vectors must get the same L2 normalization and dimension checks as documents.""" + module = type(sys)("fastembed") + setattr(module, "TextEmbedding", _UnnormalizedTextEmbedding) + monkeypatch.setitem(sys.modules, "fastembed", module) + + provider = FastEmbedEmbeddingProvider(model_name="stub-multilingual", dimensions=4) + result = await provider.embed_query("some query") + + norm = math.sqrt(sum(x * x for x in result)) + assert abs(norm - 1.0) < 1e-6, f"Expected unit norm, got {norm}" + + class _WideQueryEmbedding: + def __init__(self, model_name: str, **_kwargs): + pass + + def query_embed(self, query): + for _ in [query] if isinstance(query, str) else list(query): + yield _UnormalizedVector([1.0, 0.0, 0.0, 0.0, 0.5]) + + setattr(module, "TextEmbedding", _WideQueryEmbedding) + provider = FastEmbedEmbeddingProvider(model_name="stub-wide", dimensions=4) + with pytest.raises(RuntimeError, match="5-dimensional vectors"): + await provider.embed_query("wide query") + + @pytest.mark.asyncio async def test_fastembed_provider_zero_vector_does_not_raise(monkeypatch): """A zero vector from the model must be returned as-is without a division error."""