diff --git a/libs/langchain/langchain/embeddings/voyageai.py b/libs/langchain/langchain/embeddings/voyageai.py index 85f2b5d4137..8d39b448c45 100644 --- a/libs/langchain/langchain/embeddings/voyageai.py +++ b/libs/langchain/langchain/embeddings/voyageai.py @@ -115,10 +115,16 @@ class VoyageEmbeddings(BaseModel, Embeddings): return params def _get_embeddings( - self, texts: List[str], batch_size: int, input_type: Optional[str] = None + self, + texts: List[str], + batch_size: Optional[int] = None, + input_type: Optional[str] = None, ) -> List[List[float]]: embeddings: List[List[float]] = [] + if batch_size is None: + batch_size = self.batch_size + if self.show_progress_bar: try: from tqdm.auto import tqdm @@ -168,9 +174,24 @@ class VoyageEmbeddings(BaseModel, Embeddings): Args: text: The text to embed. + Returns: + Embedding for the text. + """ + return self._get_embeddings([text], input_type="query")[0] + + def embed_general_texts( + self, texts: List[str], *, input_type: Optional[str] = None + ) -> List[List[float]]: + """Call out to Voyage Embedding endpoint for embedding general text. + + Args: + texts: The list of texts to embed. + input_type: Type of the input text. Default to None, meaning the type is + unspecified. Other options: query, document. + Returns: Embedding for the text. """ return self._get_embeddings( - [text], batch_size=self.batch_size, input_type="query" - )[0] + texts, batch_size=self.batch_size, input_type=input_type + )