Merge pull request wagtail#75 from wagtail/fix/fix-input-arg-to-async…

…-litellm-embedding-api

Fix/fix input arg to async LiteLLM embedding API

src/wagtail_vector_index/ai_utils/backends/litellm.py

@@ -179,7 +179,7 @@ def embed(self, inputs: Iterable[str], **kwargs) -> Iterator[list[float]]:
 
     async def aembed(self, inputs: Iterable[str], **kwargs) -> list[float]:
         response = await litellm.aembedding(
-            model=self.config.model_id, inputs=inputs, **kwargs
+            model=self.config.model_id, input=inputs, **kwargs
         )
 
         return [data["embedding"] for data in response["data"]]

tests/test_ai_utils/test_backends/test_litellm.py

@@ -1,5 +1,6 @@
 import re
 from typing import List
+from unittest.mock import AsyncMock
 
 import pytest
 from wagtail_vector_index.ai_utils.backends import (
@@ -241,3 +242,21 @@ def test_litellm_embed(make_embedding_backend, mocker):
     ]
     list(backend.embed(input_text))
     embed_mock.assert_called_once_with(input=input_text, model=backend.config.model_id)
+
+
+@if_litellm_installed
+async def test_litellm_aembed(make_embedding_backend, mocker):
+    backend = make_embedding_backend()
+    mock = AsyncMock()
+    embed_mock = mocker.patch(
+        "wagtail_vector_index.ai_utils.backends.litellm.litellm.aembedding",
+        side_effect=mock,
+    )
+    input_text = [
+        "Little trotty wagtail, he waddled in the mud,",
+        "And left his little footmarks, trample where he would.",
+        "He waddled in the water-pudge, and waggle went his tail,",
+        "And chirrupt up his wings to dry upon the garden rail.",
+    ]
+    list(await backend.aembed(input_text))
+    embed_mock.assert_called_once_with(input=input_text, model=backend.config.model_id)