tests: add pytest flags

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)
 
 
 ######################

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"))

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])

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

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,13 +25,15 @@ 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()
     output = model.embed_query(document)
     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

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,

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:"])

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