Merge pull request #3 from voyage-ai/voyage_2_model

community: model attribute is required, set default batch size
langchain-ai · Feb 16, 2024 · 6e53bf9 · 6e53bf9
2 parents 1113700 + 4de3553
commit 6e53bf9
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/docs/docs/integrations/text_embedding/voyageai.ipynb b/docs/docs/integrations/text_embedding/voyageai.ipynb
@@ -27,7 +27,7 @@
    "id": "137cfde9-b88c-409a-9394-a9e31a6bf30d",
    "metadata": {},
    "source": [
-    "Voyage AI utilizes API keys to monitor usage and manage permissions. To obtain your key, create an account on our [homepage](https://www.voyageai.com). Then, create a VoyageEmbeddings model with your API key."
+    "Voyage AI utilizes API keys to monitor usage and manage permissions. To obtain your key, create an account on our [homepage](https://www.voyageai.com). Then, create a VoyageEmbeddings model with your API key. Please refer to the documentation for further details on the available models: https://docs.voyageai.com/embeddings/"
    ]
   },
   {
@@ -37,7 +37,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "embeddings = VoyageEmbeddings(voyage_api_key=\"[ Your Voyage API key ]\")"
+    "embeddings = VoyageEmbeddings(voyage_api_key=\"[ Your Voyage API key ]\", model=\"voyage-2\")"
    ]
   },
   {

diff --git a/libs/community/langchain_community/embeddings/voyageai.py b/libs/community/langchain_community/embeddings/voyageai.py
@@ -69,15 +69,15 @@ class VoyageEmbeddings(BaseModel, Embeddings):
 
             from langchain_community.embeddings import VoyageEmbeddings
 
-            voyage = VoyageEmbeddings(voyage_api_key="your-api-key")
+            voyage = VoyageEmbeddings(voyage_api_key="your-api-key", model="voyage-2")
             text = "This is a test query."
             query_result = voyage.embed_query(text)
     """
 
-    model: str = "voyage-01"
+    model: str
     voyage_api_base: str = "https://api.voyageai.com/v1/embeddings"
     voyage_api_key: Optional[SecretStr] = None
-    batch_size: int = 8
+    batch_size: Optional[int] = None
     """Maximum number of texts to embed in each API request."""
     max_retries: int = 6
     """Maximum number of retries to make when generating."""
@@ -121,7 +121,7 @@ def _get_embeddings(
         embeddings: List[List[float]] = []
 
         if batch_size is None:
-            batch_size = self.batch_size
+            batch_size = 72 if self.model in ["voyage-2", "voyage-02"] else 7
 
         if self.show_progress_bar:
             try:

diff --git a/libs/community/tests/integration_tests/embeddings/test_voyageai.py b/libs/community/tests/integration_tests/embeddings/test_voyageai.py
@@ -2,7 +2,7 @@
 from langchain_community.embeddings.voyageai import VoyageEmbeddings
 
 # Please set VOYAGE_API_KEY in the environment variables
-MODEL = "voyage-01"
+MODEL = "voyage-2"
 
 
 def test_voyagi_embedding_documents() -> None: