feat: add run_async support for OllamaTextEmbedder (#1877)

* feat: add run_async support for OllamaTextEmbedder

* fix: address review comments (use one async client per embedder obj)

---------

Co-authored-by: Michele Pangrazzi <xmikex83@gmail.com>

Author: srini047

integrations/ollama/pyproject.toml

@@ -50,6 +50,7 @@ installer = "uv"
 dependencies = [
     "coverage[toml]>=6.5",
     "pytest",
+    "pytest-asyncio",
     "pytest-rerunfailures",
     "haystack-pydoc-tools",
 ]
@@ -158,7 +159,10 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 markers = [
     "integration: marks tests as slow (deselect with '-m \"not integration\"')",
 ]
+log_cli = true
 addopts = ["--import-mode=importlib"]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "class"
 
 [[tool.mypy.overrides]]
 module = ["haystack.*", "haystack_integrations.*", "pytest.*", "ollama.*", "pydantic.*"]

integrations/ollama/src/haystack_integrations/components/embedders/ollama/text_embedder.py

@@ -2,7 +2,7 @@
 
 from haystack import component
 
-from ollama import Client
+from ollama import AsyncClient, Client
 
 
 @component
@@ -46,6 +46,7 @@ def __init__(
         self.model = model
 
         self._client = Client(host=self.url, timeout=self.timeout)
+        self._async_client = AsyncClient(host=self.url, timeout=self.timeout)
 
     @component.output_types(embedding=List[float], meta=Dict[str, Any])
     def run(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
@@ -66,3 +67,24 @@ def run(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
         result["meta"] = {"model": self.model}
 
         return result
+
+    @component.output_types(embedding=List[float], meta=Dict[str, Any])
+    async def run_async(self, text: str, generation_kwargs: Optional[Dict[str, Any]] = None):
+        """
+        Asynchronously run an Ollama Model to compute embeddings of the provided text.
+
+        :param text:
+            Text to be converted to an embedding.
+        :param generation_kwargs:
+            Optional arguments to pass to the Ollama generation endpoint, such as temperature,
+            top_p, etc. See the
+            [Ollama docs](https://github.com/jmorganca/ollama/blob/main/docs/modelfile.md#valid-parameters-and-values).
+        :returns: A dictionary with the following keys:
+            - `embedding`: The computed embeddings
+            - `meta`: The metadata collected during the embedding process
+        """
+        response = await self._async_client.embeddings(model=self.model, prompt=text, options=generation_kwargs)
+        result = response.model_dump()
+        result["meta"] = {"model": self.model}
+
+        return result

integrations/ollama/tests/test_text_embedder.py

@@ -37,7 +37,20 @@ def test_model_not_found(self):
     def test_run(self):
         embedder = OllamaTextEmbedder(model="nomic-embed-text")
 
-        reply = embedder.run("hello")
+        text = "hello"
+        reply = embedder.run(text=text)
+
+        assert isinstance(reply, dict)
+        assert all(isinstance(element, float) for element in reply["embedding"])
+        assert reply["meta"]["model"] == "nomic-embed-text"
+
+    @pytest.mark.asyncio
+    @pytest.mark.integration
+    async def test_run_async(self):
+        embedder = OllamaTextEmbedder(model="nomic-embed-text")
+
+        text = "hello"
+        reply = await embedder.run_async(text=text)
 
         assert isinstance(reply, dict)
         assert all(isinstance(element, float) for element in reply["embedding"])