Add pagination for Vertex AI embeddings

langchain/embeddings/vertexai.py

@@ -31,8 +31,13 @@ def embed_documents(self, texts: List[str]) -> List[List[float]]:
         Returns:
             List of embeddings, one for each text.
         """
-        embeddings = self.client.get_embeddings(texts)
-        return [el.values for el in embeddings]
+        batch_size = 5
+        embeddings = []
+        for batch in range(0, len(texts), batch_size):
+            text_batch = texts[batch : batch + batch_size]
+            embeddings_batch = self.client.get_embeddings(text_batch)
+            embeddings.extend([el.values for el in embeddings_batch])
+        return embeddings
 
     def embed_query(self, text: str) -> List[float]:
         """Embed a text.

tests/integration_tests/embeddings/test_vertexai.py

@@ -23,3 +23,22 @@ def test_embedding_query() -> None:
     model = VertexAIEmbeddings()
     output = model.embed_query(document)
     assert len(output) == 768
+
+
+def test_paginated_texts() -> None:
+    documents = [
+        "foo bar",
+        "foo baz",
+        "bar foo",
+        "baz foo",
+        "bar bar",
+        "foo foo",
+        "baz baz",
+        "baz bar",
+    ]
+    model = VertexAIEmbeddings()
+    output = model.embed_documents(documents)
+    assert len(output) == 8
+    assert len(output[0]) == 768
+    assert model._llm_type == "vertexai"
+    assert model.model_name == model.client._model_id