From 3ee05e81c05268dbaa785cc391be4e014fb145cd Mon Sep 17 00:00:00 2001 From: MbappeWU Date: Sun, 25 Oct 2026 12:13:35 -0700 Subject: [PATCH] fix(core): L2-normalize OpenAI-compatible embedding vectors OpenAIEmbeddingProvider now unit-normalizes embeddings so non-OpenAI OpenAI-compatible endpoints (e.g. Ark/doubao) satisfy the sqlite distance-to-similarity contract (same as FastEmbed/LiteLLM). Fixes #1862 Signed-off-by: MbappeWU Co-authored-by: Cursor --- .../repository/openai_provider.py | 19 ++++- tests/repository/test_openai_provider.py | 83 +++++++++++++++++-- 2 files changed, 93 insertions(+), 9 deletions(-) diff --git a/src/basic_memory/repository/openai_provider.py b/src/basic_memory/repository/openai_provider.py index 426e1a12d..82e409aa4 100644 --- a/src/basic_memory/repository/openai_provider.py +++ b/src/basic_memory/repository/openai_provider.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import math import os from typing import Any @@ -119,12 +120,24 @@ async def embed_batch(batch_index: int, batch: list[str]) -> None: raise RuntimeError("OpenAI embedding batch did not produce vectors.") all_vectors.extend(vectors) - if all_vectors and len(all_vectors[0]) != self.dimensions: + # sqlite_search_repository maps L2 distance to cosine similarity via + # `1 - L²/2`, which is correct only for unit vectors. FastEmbed (#1023) + # and LiteLLM already normalize. OpenAI-compatible endpoints such as + # Ark/doubao may return unnormalized vectors, so normalize here too. + normalized: list[list[float]] = [] + for vector in all_vectors: + norm = math.sqrt(sum(x * x for x in vector)) + if norm > 0: + normalized.append([x / norm for x in vector]) + else: + normalized.append(vector) + + if normalized and len(normalized[0]) != self.dimensions: raise RuntimeError( - f"Embedding model returned {len(all_vectors[0])}-dimensional vectors " + f"Embedding model returned {len(normalized[0])}-dimensional vectors " f"but provider was configured for {self.dimensions} dimensions." ) - return all_vectors + return normalized async def embed_query(self, text: str) -> list[float]: vectors = await self.embed_documents([text]) diff --git a/tests/repository/test_openai_provider.py b/tests/repository/test_openai_provider.py index 0a47ca045..b8e82bc33 100644 --- a/tests/repository/test_openai_provider.py +++ b/tests/repository/test_openai_provider.py @@ -2,6 +2,7 @@ import asyncio import builtins +import math import sys from types import SimpleNamespace @@ -72,6 +73,34 @@ async def create(self, *, model: str, input: list[str]): return SimpleNamespace(data=[SimpleNamespace(index=0, embedding=[1.0, 2.0, 3.0])]) +class _FixedVectorEmbeddingsApi: + def __init__(self, embedding: list[float]): + self._embedding = embedding + + async def create(self, *, model: str, input: list[str]): + return SimpleNamespace( + data=[ + SimpleNamespace(index=index, embedding=list(self._embedding)) + for index, _text in enumerate(input) + ] + ) + + +def _install_fixed_openai(monkeypatch: pytest.MonkeyPatch, embedding: list[float]) -> None: + class _FixedAsyncOpenAI: + def __init__(self, *, api_key: str, base_url=None, timeout=30.0): + self.embeddings = _FixedVectorEmbeddingsApi(embedding) + + module = type(sys)("openai") + setattr(module, "AsyncOpenAI", _FixedAsyncOpenAI) + monkeypatch.setitem(sys.modules, "openai", module) + monkeypatch.setenv("OPENAI_API_KEY", "test-key") + + +def _l2_norm(vector: list[float]) -> float: + return math.sqrt(sum(value * value for value in vector)) + + @pytest.fixture(autouse=True) def _reset_embedding_provider_cache_fixture(): reset_embedding_provider_cache() @@ -116,6 +145,41 @@ async def test_openai_provider_dimension_mismatch_raises_error(monkeypatch): await provider.embed_documents(["semantic note"]) +@pytest.mark.asyncio +async def test_openai_provider_l2_normalizes_vectors_for_sqlite_similarity(monkeypatch): + """OpenAI-compatible endpoints may return non-unit vectors. + + sqlite maps L2 distance to cosine similarity with ``1 - L2^2/2``, which + holds only for unit vectors. The provider must return a unit vector in the + same direction, and leave an already-unit vector on the unit sphere. + A zero vector stays zero so normalization does not divide by zero. + """ + _install_fixed_openai(monkeypatch, [3.0, 4.0, 0.0]) + provider = OpenAIEmbeddingProvider(dimensions=3) + + documents = await provider.embed_documents(["ark note"]) + query = await provider.embed_query("ark note") + + assert documents[0] == pytest.approx([0.6, 0.8, 0.0]) + assert query == pytest.approx([0.6, 0.8, 0.0]) + assert _l2_norm(documents[0]) == pytest.approx(1.0) + assert _l2_norm(query) == pytest.approx(1.0) + + _install_fixed_openai(monkeypatch, [1.0, 0.0, 0.0]) + unit_provider = OpenAIEmbeddingProvider(dimensions=3) + unit_documents = await unit_provider.embed_documents(["openai note"]) + unit_query = await unit_provider.embed_query("openai note") + + assert unit_documents[0] == pytest.approx([1.0, 0.0, 0.0]) + assert unit_query == pytest.approx([1.0, 0.0, 0.0]) + assert _l2_norm(unit_documents[0]) == pytest.approx(1.0) + assert _l2_norm(unit_query) == pytest.approx(1.0) + + _install_fixed_openai(monkeypatch, [0.0, 0.0, 0.0]) + zero_provider = OpenAIEmbeddingProvider(dimensions=3) + assert await zero_provider.embed_documents(["empty"]) == [[0.0, 0.0, 0.0]] + + @pytest.mark.asyncio async def test_openai_provider_missing_dependency_raises_actionable_error(monkeypatch): """Missing openai package should raise SemanticDependenciesMissingError.""" @@ -565,13 +629,20 @@ def __init__(self, *, api_key: str, base_url=None, timeout=30.0): vectors = await provider.embed_documents(["a", "bbbb", "ccc", "dd"]) - assert shared_api.max_in_flight >= 2 - assert vectors == [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [3.0, 4.0, 5.0], - [2.0, 3.0, 4.0], + def _unit(raw: list[float]) -> list[float]: + norm = math.sqrt(sum(component * component for component in raw)) + return [component / norm for component in raw] + + expected = [ + _unit([1.0, 2.0, 3.0]), + _unit([4.0, 5.0, 6.0]), + _unit([3.0, 4.0, 5.0]), + _unit([2.0, 3.0, 4.0]), ] + assert shared_api.max_in_flight >= 2 + assert len(vectors) == len(expected) + for actual, wanted in zip(vectors, expected, strict=True): + assert actual == pytest.approx(wanted) @pytest.mark.asyncio