Add pagination for Vertex AI embeddings (#5325)

Fixes #5316

---------

Co-authored-by: Justin Flick <jflick@homesite.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
This commit is contained in:
Justin Flick
2023-05-29 06:57:41 -07:00
committed by GitHub
parent 3e16468423
commit c09f8e4ddc
2 changed files with 31 additions and 4 deletions

View File

@@ -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