fix(core): L2-normalize LiteLLM vectors and mirror OpenAI provider shape

Bring the LiteLLM provider in line with the unit-norm contract from
sqlite_search_repository.py (lines 65-67): the cosine-similarity formula
`1 - L²/2` is correct only for unit-normalized vectors. LiteLLM routes to
many backends (Cohere, Vertex, Bedrock, etc.) that do not return normalized
embeddings, so normalize at the provider boundary — same fix shape as the
parallel FastEmbed change in #843.

Also align the response handling with OpenAIEmbeddingProvider:
- attribute access on response items (item.index / item.embedding)
- explicit duplicate-index guard

Tests cover the three behaviors directly (unit norm, zero-vector pass-through,
duplicate-index error) and the existing ordering test now reconstructs the
expected normalized vectors so a normalization regression would be caught.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: phernandez <paul@basicmachines.co>
This commit is contained in:
phernandez
2026-05-26 13:29:20 -05:00
parent fce7a67611
commit f9e7029ae7
2 changed files with 116 additions and 13 deletions
@@ -14,6 +14,7 @@ supported embedding models.
from __future__ import annotations
import asyncio
import math
from typing import Any
from basic_memory.repository.embedding_provider import EmbeddingProvider
@@ -80,8 +81,12 @@ class LiteLLMEmbeddingProvider(EmbeddingProvider):
vectors_by_index: dict[int, list[float]] = {}
for item in response.data:
response_index = int(item["index"])
vectors_by_index[response_index] = [float(v) for v in item["embedding"]]
response_index = int(item.index)
if response_index in vectors_by_index:
raise RuntimeError(
"LiteLLM embedding response returned duplicate vector indexes."
)
vectors_by_index[response_index] = [float(v) for v in item.embedding]
ordered_vectors: list[list[float]] = []
for index in range(len(batch)):
@@ -104,12 +109,25 @@ class LiteLLMEmbeddingProvider(EmbeddingProvider):
raise RuntimeError("LiteLLM embedding batch did not produce vectors.")
all_vectors.extend(vectors)
if all_vectors and len(all_vectors[0]) != self.dimensions:
# sqlite_search_repository.py maps L2 distance to cosine similarity via
# `1 - L²/2`, which is correct only for unit-normalized vectors. LiteLLM
# routes to many backends (Cohere, Vertex, Bedrock, etc.); not all of
# them return normalized embeddings, so we normalize here to honor the
# provider contract regardless of the underlying model.
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])