feat: add run_async support for CohereTextEmbedder (#1873)

* feat: add run_async support for CohereTextEmbedder

* fix: review comments

Author: srini047

integrations/cohere/pyproject.toml

@@ -168,3 +168,5 @@ markers = [
   "integration: integration tests",
 ]
 log_cli = true
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "class"

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

@@ -75,6 +75,16 @@ def __init__(
         self.timeout = timeout
         self.embedding_type = embedding_type or EmbeddingTypes.FLOAT
 
+    def _prepare_input(self, text: str) -> str:
+        if not isinstance(text, str):
+            msg = (
+                "CohereTextEmbedder expects a string as input."
+                "In case you want to embed a list of Documents, please use the CohereDocumentEmbedder."
+            )
+            raise TypeError(msg)
+
+        return text
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serializes the component to a dictionary.
@@ -114,20 +124,19 @@ def from_dict(cls, data: Dict[str, Any]) -> "CohereTextEmbedder":
 
     @component.output_types(embedding=List[float], meta=Dict[str, Any])
     def run(self, text: str):
-        """Embed text.
+        """
+        Embed text.
 
-        :param text: the text to embed.
-        :returns: A dictionary with the following keys:
-            - `embedding`: the embedding of the text.
-            - `meta`: metadata about the request.
-        :raises TypeError: If the input is not a string.
+        :param text:
+            the text to embed.
+        :returns:
+            A dictionary with the following keys:
+                - `embedding`: the embedding of the text.
+                - `meta`: metadata about the request.
+        :raises TypeError:
+            If the input is not a string.
         """
-        if not isinstance(text, str):
-            msg = (
-                "CohereTextEmbedder expects a string as input."
-                "In case you want to embed a list of Documents, please use the CohereDocumentEmbedder."
-            )
-            raise TypeError(msg)
+        text = self._prepare_input(text=text)
 
         # Establish connection to API
 
@@ -158,3 +167,40 @@ def run(self, text: str):
             )
 
         return {"embedding": embedding[0], "meta": metadata}
+
+    @component.output_types(embedding=List[float], meta=Dict[str, Any])
+    async def run_async(self, text: str):
+        """
+        Asynchronously embed text.
+
+        This is the asynchronous version of the `run` method. It has the same parameters and return values
+        but can be used with `await` in async code.
+
+         :param text:
+            Text to embed.
+
+        :returns:
+            A dictionary with the following keys:
+            - `embedding`: the embedding of the text.
+            - `meta`: metadata about the request.
+
+        :raises TypeError:
+            If the input is not a string.
+        """
+        text = self._prepare_input(text=text)
+
+        api_key = self.api_key.resolve_value()
+        assert api_key is not None
+
+        cohere_client = AsyncClientV2(
+            api_key,
+            base_url=self.api_base_url,
+            timeout=self.timeout,
+            client_name="haystack",
+        )
+
+        embedding, metadata = await get_async_response(
+            cohere_client, [text], self.model, self.input_type, self.truncate, self.embedding_type
+        )
+
+        return {"embedding": embedding[0], "meta": metadata}

integrations/cohere/tests/test_text_embedder.py

@@ -116,5 +116,21 @@ def test_run_wrong_input_format(self):
     def test_run(self):
         embedder = CohereTextEmbedder()
         text = "The food was delicious"
-        result = embedder.run(text)
+        result = embedder.run(text=text)
+
+        assert len(result["embedding"]) == 4096
+        assert all(isinstance(x, float) for x in result["embedding"])
+
+    @pytest.mark.asyncio
+    @pytest.mark.skipif(
+        not os.environ.get("COHERE_API_KEY", None) and not os.environ.get("CO_API_KEY", None),
+        reason="Export an env var called COHERE_API_KEY/CO_API_KEY containing the Cohere API key to run this test.",
+    )
+    @pytest.mark.integration
+    async def test_run_async(self):
+        embedder = CohereTextEmbedder()
+        text = "The food was delicious"
+        result = await embedder.run_async(text=text)
+
+        assert len(result["embedding"]) == 4096
         assert all(isinstance(x, float) for x in result["embedding"])