diff --git a/src/basic_memory/repository/litellm_provider.py b/src/basic_memory/repository/litellm_provider.py index e7eb8b5b..a0735bb5 100644 --- a/src/basic_memory/repository/litellm_provider.py +++ b/src/basic_memory/repository/litellm_provider.py @@ -155,12 +155,17 @@ class LiteLLMEmbeddingProvider(EmbeddingProvider): vectors_by_index: dict[int, list[float]] = {} for item in response.data: - response_index = int(item.index) + if isinstance(item, dict): + response_index = int(item["index"]) + embedding = item["embedding"] + else: + response_index = int(item.index) + embedding = item.embedding 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] + vectors_by_index[response_index] = [float(v) for v in embedding] ordered_vectors: list[list[float]] = [] for index in range(len(batch)): diff --git a/tests/repository/test_litellm_provider.py b/tests/repository/test_litellm_provider.py index 93d83f39..d21d97c5 100644 --- a/tests/repository/test_litellm_provider.py +++ b/tests/repository/test_litellm_provider.py @@ -444,6 +444,24 @@ async def test_litellm_provider_zero_vector_does_not_raise(monkeypatch): assert result == [[0.0, 0.0, 0.0]] +@pytest.mark.asyncio +async def test_litellm_provider_accepts_dict_response_items(monkeypatch): + """LiteLLM providers may return embedding data as dict items.""" + + async def _aembedding(**kwargs): + data = [{"index": i, "embedding": [1.0, 0.0, 0.0]} for i in range(len(kwargs["input"]))] + return SimpleNamespace(data=data) + + module = type(sys)("litellm") + setattr(module, "aembedding", _aembedding) + monkeypatch.setitem(sys.modules, "litellm", module) + + provider = LiteLLMEmbeddingProvider(dimensions=3) + result = await provider.embed_documents(["first", "second"]) + + assert result == [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]] + + @pytest.mark.asyncio async def test_litellm_provider_duplicate_index_raises_error(monkeypatch): """A backend returning duplicate indexes is malformed and must fail fast."""