diff --git a/libs/vertexai/Makefile b/libs/vertexai/Makefile index 9ee09b32b..6c7e948d1 100644 --- a/libs/vertexai/Makefile +++ b/libs/vertexai/Makefile @@ -9,7 +9,7 @@ TEST_FILE ?= tests/unit_tests/ integration_test integration_tests: TEST_FILE = tests/integration_tests/ test tests integration_test integration_tests: - poetry run pytest $(TEST_FILE) + poetry run pytest --release $(TEST_FILE) ###################### diff --git a/libs/vertexai/tests/conftest.py b/libs/vertexai/tests/conftest.py new file mode 100644 index 000000000..7aee5c088 --- /dev/null +++ b/libs/vertexai/tests/conftest.py @@ -0,0 +1,52 @@ +""" +Tests configuration to be executed before tests execution. +""" + +from typing import List +import pytest + + +_RELEASE_FLAG = "release" +_GPU_FLAG = "gpu" +_LONG_FLAG = "long" +_EXTENDED_FLAG = "extended" + +_PYTEST_FLAGS = [ + _RELEASE_FLAG, + _GPU_FLAG, + _LONG_FLAG, + _EXTENDED_FLAG +] + + +def pytest_addoption(parser: pytest.Parser) -> None: + """ + Add flags accepted by our pytest CLI, + :param parser: The pytest parser object. + :return: + """ + for flag in _PYTEST_FLAGS: + parser.addoption(f"--{flag}", action="store_true", default=False, help=f"run {flag} tests") + + +def pytest_configure(config: pytest.Config) -> None: + """ + Add our custom configuration, + :param config: The pytest config object. + :return: + """ + for flag in _PYTEST_FLAGS: + config.addinivalue_line("markers", f"{flag}: mark test to run as {flag} only test") + + +def pytest_collection_modifyitems(config: pytest.Config, items: List[pytest.Item]) -> None: + """ + Skip tests with a marker from our list that were not explicitly invoked. + :param config: The pytest config object. + :param items: The list of tests to be executed. + :return: + """ + for item in items: + keywords = list(set(item.keywords).intersection(_PYTEST_FLAGS)) + if keywords and not any((config.getoption(f"--{keyword}") for keyword in keywords)): + item.add_marker(pytest.mark.skip(reason=f"need --{keywords[0]} option to run")) diff --git a/libs/vertexai/tests/integration_tests/test_callbacks.py b/libs/vertexai/tests/integration_tests/test_callbacks.py index f0b9e84cd..910cb1561 100644 --- a/libs/vertexai/tests/integration_tests/test_callbacks.py +++ b/libs/vertexai/tests/integration_tests/test_callbacks.py @@ -6,6 +6,7 @@ from langchain_google_vertexai.llms import VertexAI +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "text-bison@001", "code-bison@001"], @@ -25,6 +26,7 @@ def test_llm_invoke(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "chat-bison@001", "codechat-bison@001"], @@ -45,6 +47,7 @@ def test_chat_call(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release @pytest.mark.parametrize( "model_name", ["gemini-pro", "text-bison@001", "code-bison@001"], @@ -64,6 +67,7 @@ def test_invoke_config(model_name: str) -> None: assert vb.completion_tokens > completion_tokens +@pytest.mark.release def test_llm_stream() -> None: vb = VertexAICallbackHandler() llm = VertexAI(model_name="gemini-pro", temperature=0.0, callbacks=[vb]) @@ -74,6 +78,7 @@ def test_llm_stream() -> None: assert vb.completion_tokens > 0 +@pytest.mark.release def test_chat_stream() -> None: vb = VertexAICallbackHandler() llm = ChatVertexAI(model_name="gemini-pro", temperature=0.0, callbacks=[vb]) diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py index 365c7a213..0357bbb0b 100644 --- a/libs/vertexai/tests/integration_tests/test_chat_models.py +++ b/libs/vertexai/tests/integration_tests/test_chat_models.py @@ -18,6 +18,7 @@ model_names_to_test = [None, "codechat-bison", "chat-bison", "gemini-pro"] +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_initialization(model_name: Optional[str]) -> None: """Test chat model initialization.""" @@ -32,6 +33,7 @@ def test_initialization(model_name: Optional[str]) -> None: assert model.model_name == model.client._model_name.split("/")[-1] +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_vertexai_single_call(model_name: Optional[str]) -> None: if model_name: @@ -44,6 +46,7 @@ def test_vertexai_single_call(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release # mark xfail because Vertex API randomly doesn't respect # the n/candidate_count parameter @pytest.mark.xfail @@ -56,6 +59,7 @@ def test_candidates() -> None: assert len(response.generations[0]) == 2 +@pytest.mark.release @pytest.mark.parametrize("model_name", ["chat-bison@001", "gemini-pro"]) async def test_vertexai_agenerate(model_name: str) -> None: model = ChatVertexAI(temperature=0, model_name=model_name) @@ -76,6 +80,7 @@ async def test_vertexai_agenerate(model_name: str) -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.parametrize("model_name", ["chat-bison@001", "gemini-pro"]) def test_vertexai_stream(model_name: str) -> None: model = ChatVertexAI(temperature=0, model_name=model_name) @@ -86,6 +91,7 @@ def test_vertexai_stream(model_name: str) -> None: assert isinstance(chunk, AIMessageChunk) +@pytest.mark.release async def test_vertexai_astream() -> None: model = ChatVertexAI(temperature=0, model_name="gemini-pro") message = HumanMessage(content="Hello") @@ -94,6 +100,7 @@ async def test_vertexai_astream() -> None: assert isinstance(chunk, AIMessageChunk) +@pytest.mark.release def test_vertexai_single_call_with_context() -> None: model = ChatVertexAI() raw_context = ( @@ -110,6 +117,7 @@ def test_vertexai_single_call_with_context() -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_multimodal() -> None: llm = ChatVertexAI(model_name="gemini-pro-vision") gcs_url = ( @@ -129,6 +137,7 @@ def test_multimodal() -> None: assert isinstance(output.content, str) +@pytest.mark.release @pytest.mark.xfail(reason="problem on vertex side") def test_multimodal_history() -> None: llm = ChatVertexAI(model_name="gemini-pro-vision") @@ -158,6 +167,7 @@ def test_multimodal_history() -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_vertexai_single_call_with_examples() -> None: model = ChatVertexAI() raw_context = "My name is Peter. You are my personal assistant." @@ -172,6 +182,7 @@ def test_vertexai_single_call_with_examples() -> None: assert isinstance(response.content, str) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_vertexai_single_call_with_history(model_name: Optional[str]) -> None: if model_name: @@ -188,6 +199,7 @@ def test_vertexai_single_call_with_history(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release def test_vertexai_single_call_fails_no_message() -> None: chat = ChatVertexAI() with pytest.raises(ValueError) as exc_info: @@ -198,6 +210,7 @@ def test_vertexai_single_call_fails_no_message() -> None: ) +@pytest.mark.release @pytest.mark.parametrize("model_name", ["gemini-pro"]) def test_chat_vertexai_gemini_system_message_error(model_name: str) -> None: model = ChatVertexAI(model_name=model_name) @@ -211,6 +224,7 @@ def test_chat_vertexai_gemini_system_message_error(model_name: str) -> None: model([system_message, message1, message2, message3]) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_chat_vertexai_system_message(model_name: Optional[str]) -> None: if model_name: @@ -231,6 +245,7 @@ def test_chat_vertexai_system_message(model_name: Optional[str]) -> None: assert isinstance(response.content, str) +@pytest.mark.release @pytest.mark.parametrize("model_name", model_names_to_test) def test_get_num_tokens_from_messages(model_name: str) -> None: if model_name: @@ -243,6 +258,7 @@ def test_get_num_tokens_from_messages(model_name: str) -> None: assert token == 3 +@pytest.mark.release def test_chat_vertexai_gemini_function_calling() -> None: class MyModel(BaseModel): name: str diff --git a/libs/vertexai/tests/integration_tests/test_embeddings.py b/libs/vertexai/tests/integration_tests/test_embeddings.py index 42495629b..6772274ac 100644 --- a/libs/vertexai/tests/integration_tests/test_embeddings.py +++ b/libs/vertexai/tests/integration_tests/test_embeddings.py @@ -8,11 +8,13 @@ from langchain_google_vertexai.embeddings import VertexAIEmbeddings +@pytest.mark.release def test_initialization() -> None: """Test embedding model initialization.""" VertexAIEmbeddings() +@pytest.mark.release def test_langchain_google_vertexai_embedding_documents() -> None: documents = ["foo bar"] model = VertexAIEmbeddings() @@ -23,6 +25,7 @@ def test_langchain_google_vertexai_embedding_documents() -> None: assert model.model_name == "textembedding-gecko@001" +@pytest.mark.release def test_langchain_google_vertexai_embedding_query() -> None: document = "foo bar" model = VertexAIEmbeddings() @@ -30,6 +33,7 @@ def test_langchain_google_vertexai_embedding_query() -> None: assert len(output) == 768 +@pytest.mark.release def test_langchain_google_vertexai_large_batches() -> None: documents = ["foo bar" for _ in range(0, 251)] model_uscentral1 = VertexAIEmbeddings(location="us-central1") @@ -40,6 +44,7 @@ def test_langchain_google_vertexai_large_batches() -> None: assert model_asianortheast1.instance["batch_size"] < 50 +@pytest.mark.release def test_langchain_google_vertexai_paginated_texts() -> None: documents = [ "foo bar", @@ -58,6 +63,7 @@ def test_langchain_google_vertexai_paginated_texts() -> None: assert model.model_name == model.client._model_id +@pytest.mark.release def test_warning(caplog: pytest.LogCaptureFixture) -> None: _ = VertexAIEmbeddings() assert len(caplog.records) == 1 diff --git a/libs/vertexai/tests/integration_tests/test_llms.py b/libs/vertexai/tests/integration_tests/test_llms.py index f9525c8b2..cd5e08970 100644 --- a/libs/vertexai/tests/integration_tests/test_llms.py +++ b/libs/vertexai/tests/integration_tests/test_llms.py @@ -13,6 +13,7 @@ model_names_to_test_with_default = [None] + model_names_to_test +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -26,6 +27,7 @@ def test_vertex_initialization(model_name: str) -> None: assert llm.model_name == llm.client._model_name.split("/")[-1] +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -40,6 +42,7 @@ def test_vertex_invoke(model_name: str) -> None: assert isinstance(output, str) +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -58,6 +61,7 @@ def test_vertex_generate(model_name: str) -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates") def test_vertex_generate_multiple_candidates() -> None: llm = VertexAI(temperature=0.3, n=2, model_name="text-bison@001") @@ -67,6 +71,7 @@ def test_vertex_generate_multiple_candidates() -> None: assert len(output.generations[0]) == 2 +@pytest.mark.release @pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates") def test_vertex_generate_code() -> None: llm = VertexAI(temperature=0.3, n=2, model_name="code-bison@001") @@ -79,6 +84,7 @@ def test_vertex_generate_code() -> None: assert int(usage_metadata["candidates_token_count"]) > 1 +@pytest.mark.release async def test_vertex_agenerate() -> None: llm = VertexAI(temperature=0) output = await llm.agenerate(["Please say foo:"]) @@ -88,6 +94,7 @@ async def test_vertex_agenerate() -> None: assert int(usage_metadata["candidates_token_count"]) > 0 +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test_with_default, @@ -102,6 +109,7 @@ def test_stream(model_name: str) -> None: assert isinstance(token, str) +@pytest.mark.release async def test_vertex_consistency() -> None: llm = VertexAI(temperature=0) output = llm.generate(["Please say foo:"]) @@ -111,12 +119,14 @@ async def test_vertex_consistency() -> None: assert output.generations[0][0].text == async_output.generations[0][0].text +@pytest.mark.release async def test_astream() -> None: llm = VertexAI(temperature=0, model_name="gemini-pro") async for token in llm.astream("I'm Pickle Rick"): assert isinstance(token, str) +@pytest.mark.release @pytest.mark.parametrize( "model_name", model_names_to_test, diff --git a/libs/vertexai/tests/integration_tests/test_llms_safety.py b/libs/vertexai/tests/integration_tests/test_llms_safety.py index 1b02eb5a5..0b3ef041c 100644 --- a/libs/vertexai/tests/integration_tests/test_llms_safety.py +++ b/libs/vertexai/tests/integration_tests/test_llms_safety.py @@ -41,7 +41,7 @@ """ -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended def test_gemini_safety_settings_generate() -> None: llm = VertexAI(model_name="gemini-pro", safety_settings=SAFETY_SETTINGS) output = llm.generate(["What do you think about child abuse:"]) @@ -70,7 +70,7 @@ def test_gemini_safety_settings_generate() -> None: assert not generation_info.get("is_blocked") -@pytest.mark.skip("CI testing not set up") +@pytest.mark.extended async def test_gemini_safety_settings_agenerate() -> None: llm = VertexAI(model_name="gemini-pro", safety_settings=SAFETY_SETTINGS) output = await llm.agenerate(["What do you think about child abuse:"]) diff --git a/libs/vertexai/tests/integration_tests/test_tools.py b/libs/vertexai/tests/integration_tests/test_tools.py index ef172a1bb..9fcd95879 100644 --- a/libs/vertexai/tests/integration_tests/test_tools.py +++ b/libs/vertexai/tests/integration_tests/test_tools.py @@ -112,6 +112,7 @@ def test_stream() -> None: assert "function_call" in response[0].additional_kwargs +@pytest.mark.release def test_multiple_tools() -> None: from langchain.agents import AgentExecutor from langchain.agents.format_scratchpad import format_to_openai_function_messages diff --git a/libs/vertexai/tests/integration_tests/test_vision_models.py b/libs/vertexai/tests/integration_tests/test_vision_models.py index a3edaa047..7714373ad 100644 --- a/libs/vertexai/tests/integration_tests/test_vision_models.py +++ b/libs/vertexai/tests/integration_tests/test_vision_models.py @@ -10,6 +10,7 @@ ) +@pytest.mark.release def test_vertex_ai_image_captioning_chat(base64_image: str): # This should work model = VertexAIImageCaptioningChat() @@ -43,12 +44,14 @@ def test_vertex_ai_image_captioning_chat(base64_image: str): ) +@pytest.mark.release def test_vertex_ai_image_captioning(base64_image: str): model = VertexAIImageCaptioning() response = model.invoke(base64_image) assert isinstance(response, str) +@pytest.mark.release def test_vertex_ai_visual_qna_chat(base64_image: str): model = VertexAIVisualQnAChat() @@ -112,6 +115,7 @@ def test_vertex_ai_visual_qna_chat(base64_image: str): ) +@pytest.mark.release def test_vertex_ai_image_generation_and_edition(): generator = VertexAIImageGeneratorChat()