feat: Support async embedding/rerank/completion (#204)

Author: jiayin-nvidia

libs/ai-endpoints/langchain_nvidia_ai_endpoints/_common.py

@@ -810,7 +810,8 @@ async def aget_req_stream(
                 line = await reader.readline()
                 if not line:  # EOF
                     break
-                if line and line.strip() != b"data: [DONE]":
+                line = line.strip()
+                if line and line != b"data: [DONE]":
                     line_str = line.decode("utf-8")
                     msg, final_line = call.postprocess(line_str)
                     yield msg

libs/ai-endpoints/langchain_nvidia_ai_endpoints/embeddings.py

@@ -1,3 +1,4 @@
+import json
 from typing import Any, Dict, List, Literal, Optional
 
 from langchain_core.embeddings import Embeddings
@@ -164,6 +165,46 @@ def get_available_models(
         """Get a list of available models that work with `NVIDIAEmbeddings`."""
         return cls(**kwargs).available_models
 
+    def _prepare_payload(
+        self, texts: List[str], model_type: Literal["passage", "query"]
+    ) -> Dict[str, Any]:
+        """Prepare payload for both sync and async methods.
+
+        Args:
+            texts: List of texts to embed
+            model_type: Type of embedding ("passage" or "query")
+
+        Returns:
+            Payload dictionary
+        """
+        payload: Dict[str, Any] = {
+            "input": texts,
+            "model": self.model,
+            "encoding_format": "float",
+            "input_type": model_type,
+        }
+        if self.truncate:
+            payload["truncate"] = self.truncate
+        if self.dimensions:
+            payload["dimensions"] = self.dimensions
+        return payload
+
+    def _process_response(self, result: Dict[str, Any]) -> List[List[float]]:
+        """Process response for both sync and async methods.
+
+        Args:
+            result: Parsed JSON response from the API
+
+        Returns:
+            List of embeddings sorted by index
+        """
+        data = result.get("data", result)
+        if not isinstance(data, list):
+            raise ValueError(f"Expected data with a list of embeddings. Got: {data}")
+        embedding_list = [(res["embedding"], res["index"]) for res in data]
+        self._invoke_callback_vars(result)
+        return [x[0] for x in sorted(embedding_list, key=lambda x: x[1])]
+
     def _embed(
         self, texts: List[str], model_type: Literal["passage", "query"]
     ) -> List[List[float]]:
@@ -177,47 +218,69 @@ def _embed(
         #  truncate: "NONE" | "START" | "END"  -- default "NONE", error raised if
         #                                         an input is too long
         #  dimensions: int                     -- not supported by all models
-        payload: Dict[str, Any] = {
-            "input": texts,
-            "model": self.model,
-            "encoding_format": "float",
-            "input_type": model_type,
-        }
-        if self.truncate:
-            payload["truncate"] = self.truncate
-        if self.dimensions:
-            payload["dimensions"] = self.dimensions
-
+        payload = self._prepare_payload(texts, model_type)
         response = self._client.get_req(
             payload=payload,
             extra_headers=self.default_headers,
         )
         response.raise_for_status()
         result = response.json()
-        data = result.get("data", result)
-        if not isinstance(data, list):
-            raise ValueError(f"Expected data with a list of embeddings. Got: {data}")
-        embedding_list = [(res["embedding"], res["index"]) for res in data]
-        self._invoke_callback_vars(result)
-        return [x[0] for x in sorted(embedding_list, key=lambda x: x[1])]
+        return self._process_response(result)
+
+    def _validate_texts(self, texts: List[str]) -> None:
+        """Validate that texts is a list of strings.
+
+        Args:
+            texts: List to validate
+
+        Raises:
+            ValueError: If texts is not a list of strings
+        """
+        if not isinstance(texts, list) or not all(
+            isinstance(text, str) for text in texts
+        ):
+            raise ValueError(f"`texts` must be a list of strings, given: {repr(texts)}")
 
     def embed_query(self, text: str) -> List[float]:
         """Input pathway for query embeddings."""
         return self._embed([text], model_type="query")[0]
 
     def embed_documents(self, texts: List[str]) -> List[List[float]]:
         """Input pathway for document embeddings."""
-        if not isinstance(texts, list) or not all(
-            isinstance(text, str) for text in texts
-        ):
-            raise ValueError(f"`texts` must be a list of strings, given: {repr(texts)}")
+        self._validate_texts(texts)
 
         all_embeddings = []
         for i in range(0, len(texts), self.max_batch_size):
             batch = texts[i : i + self.max_batch_size]
             all_embeddings.extend(self._embed(batch, model_type="passage"))
         return all_embeddings
 
+    async def _aembed(
+        self, texts: List[str], model_type: Literal["passage", "query"]
+    ) -> List[List[float]]:
+        """Async version of _embed."""
+        payload = self._prepare_payload(texts, model_type)
+        response_text = await self._client.aget_req(
+            payload=payload,
+            extra_headers=self.default_headers,
+        )
+        result = json.loads(response_text)
+        return self._process_response(result)
+
+    async def aembed_query(self, text: str) -> List[float]:
+        """Async input pathway for query embeddings."""
+        return (await self._aembed([text], model_type="query"))[0]
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Async input pathway for document embeddings."""
+        self._validate_texts(texts)
+
+        all_embeddings: List[List[float]] = []
+        for i in range(0, len(texts), self.max_batch_size):
+            batch = texts[i : i + self.max_batch_size]
+            all_embeddings.extend(await self._aembed(batch, model_type="passage"))
+        return all_embeddings
+
     def _invoke_callback_vars(self, response: dict) -> None:
         """Invoke the callback context variables if there are any."""
         callback_vars = [

libs/ai-endpoints/langchain_nvidia_ai_endpoints/llm.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import json
 import warnings
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
 
-from langchain_core.callbacks.manager import CallbackManagerForLLMRun
+from langchain_core.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
 from langchain_core.language_models.llms import LLM
 from langchain_core.outputs import GenerationChunk
 from pydantic import ConfigDict, Field, PrivateAttr
@@ -168,13 +172,22 @@ def _identifying_params(self) -> Dict[str, Any]:
             "base_url": self.base_url,
         }
 
-    def _call(
+    def _prepare_call_payload(
         self,
         prompt: str,
         stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
-    ) -> str:
+    ) -> Dict[str, Any]:
+        """Prepare payload for non-streaming calls (both sync and async).
+
+        Args:
+            prompt: The prompt to send
+            stop: Stop words
+            kwargs: Additional keyword arguments
+
+        Returns:
+            Payload dictionary
+        """
         payload: Dict[str, Any] = {
             "model": self.model,
             "prompt": prompt,
@@ -188,28 +201,24 @@ def _call(
             warnings.warn("stream set to true for non-streaming call, ignoring")
             del payload["stream"]
 
-        response = self._client.get_req(payload=payload)
-        response.raise_for_status()
+        return payload
 
-        # todo: handle response's usage and system_fingerprint
-
-        choices = response.json()["choices"]
-        # todo: write a test for this by setting n > 1 on the request
-        #       aug 2024: n > 1 is not supported by endpoints
-        if len(choices) > 1:
-            warnings.warn(
-                f"Multiple choices in response, returning only the first: {choices}"
-            )
-
-        return choices[0]["text"]
-
-    def _stream(
+    def _prepare_stream_payload(
         self,
         prompt: str,
         stop: Optional[List[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
         **kwargs: Any,
-    ) -> Iterator[GenerationChunk]:
+    ) -> Dict[str, Any]:
+        """Prepare payload for streaming calls (both sync and async).
+
+        Args:
+            prompt: The prompt to send
+            stop: Stop words
+            kwargs: Additional keyword arguments
+
+        Returns:
+            Payload dictionary
+        """
         payload: Dict[str, Any] = {
             "model": self.model,
             "prompt": prompt,
@@ -226,9 +235,79 @@ def _stream(
             warnings.warn("stream set to false for streaming call, ignoring")
             payload["stream"] = True
 
+        return payload
+
+    def _process_result(self, result: Dict[str, Any]) -> str:
+        """Process parsed JSON result from both sync and async call methods.
+
+        Args:
+            result: Parsed JSON response
+
+        Returns:
+            Generated text
+        """
+        # todo: handle response's usage and system_fingerprint
+        choices = result["choices"]
+        # todo: write a test for this by setting n > 1 on the request
+        #       aug 2024: n > 1 is not supported by endpoints
+        if len(choices) > 1:
+            warnings.warn(
+                f"Multiple choices in response, returning only the first: {choices}"
+            )
+
+        return choices[0]["text"]
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        payload = self._prepare_call_payload(prompt, stop, **kwargs)
+        response = self._client.get_req(payload=payload)
+        response.raise_for_status()
+        result = response.json()
+        return self._process_result(result)
+
+    def _stream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        payload = self._prepare_stream_payload(prompt, stop, **kwargs)
         for chunk in self._client.get_req_stream(payload=payload):
             content = chunk["content"]
             generation = GenerationChunk(text=content)
             if run_manager:  # todo: add tests for run_manager
                 run_manager.on_llm_new_token(content, chunk=generation)
             yield generation
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        payload = self._prepare_call_payload(prompt, stop, **kwargs)
+        response_text = await self._client.aget_req(payload=payload)
+        result = json.loads(response_text)
+        return self._process_result(result)
+
+    async def _astream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[GenerationChunk]:
+        payload = self._prepare_stream_payload(prompt, stop, **kwargs)
+        async for chunk in self._client.aget_req_stream(payload=payload):
+            content = chunk["content"]
+            generation = GenerationChunk(text=content)
+            if run_manager:
+                await run_manager.on_llm_new_token(content, chunk=generation)
+            yield generation