From 78f9643b20aacc23eb157a5d87607cee55ecea2b Mon Sep 17 00:00:00 2001 From: Leonid Kuligin Date: Tue, 29 Oct 2024 21:27:47 +0100 Subject: [PATCH] fix integration tests (#581) --- .../integration_tests/test_embeddings.py | 26 ++++--------------- 1 file changed, 5 insertions(+), 21 deletions(-) diff --git a/libs/vertexai/tests/integration_tests/test_embeddings.py b/libs/vertexai/tests/integration_tests/test_embeddings.py index fa1c518b..a8c2b6cf 100644 --- a/libs/vertexai/tests/integration_tests/test_embeddings.py +++ b/libs/vertexai/tests/integration_tests/test_embeddings.py @@ -14,7 +14,7 @@ ) _EMBEDDING_MODELS = [ - ("textembedding-gecko@004", 768), + ("text-embedding-004", 768), ("multimodalembedding@001", 1408), ] @@ -24,11 +24,11 @@ def test_initialization() -> None: """Test embedding model initialization.""" for embeddings in [ VertexAIEmbeddings( - model_name="textembedding-gecko", + model_name="text-embedding-004", ), - VertexAIEmbeddings(model="textembedding-gecko"), + VertexAIEmbeddings(model="text-embedding-004"), ]: - assert embeddings.model_name == "textembedding-gecko" + assert embeddings.model_name == "text-embedding-004" @pytest.mark.release @@ -65,22 +65,6 @@ def test_langchain_google_vertexai_embedding_query(model_name, embeddings_dim) - assert len(output) == embeddings_dim -@pytest.mark.release -def test_langchain_google_vertexai_large_batches() -> None: - batch_size = 32 - documents = ["foo bar" for _ in range(batch_size)] - model_uscentral1 = VertexAIEmbeddings( - model_name="textembedding-gecko@001", location="us-central1" - ) - # model_asianortheast1 = VertexAIEmbeddings( - # model_name="textembedding-gecko@001", location="asia-northeast1" - # ) - model_uscentral1.embed_documents(documents) - # model_asianortheast1.embed_documents(documents) - assert model_uscentral1.instance["batch_size"] >= batch_size - # assert model_asianortheast1.instance["batch_size"] < 50 - - @pytest.mark.release @pytest.mark.parametrize( "dim, expected_dim", @@ -100,7 +84,7 @@ def test_langchain_google_vertexai_image_embeddings( @pytest.mark.release def test_langchain_google_vertexai_text_model() -> None: - embeddings_model = VertexAIEmbeddings(model_name="textembedding-gecko@004") + embeddings_model = VertexAIEmbeddings(model_name="text-embedding-004") assert isinstance(embeddings_model.client, TextEmbeddingModel) assert embeddings_model.model_type == GoogleEmbeddingModelType.TEXT