diff --git a/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py b/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py index e9e9b861..cb12c633 100644 --- a/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py +++ b/libs/vertexai/langchain_google_vertexai/_anthropic_utils.py @@ -31,7 +31,9 @@ from pydantic import BaseModel if TYPE_CHECKING: - from anthropic.types import RawMessageStreamEvent # type: ignore[import-not-found] + from anthropic.types import ( + RawMessageStreamEvent, # type: ignore[unused-ignore, import-not-found] + ) _message_type_lookups = { "human": "user", diff --git a/libs/vertexai/langchain_google_vertexai/model_garden.py b/libs/vertexai/langchain_google_vertexai/model_garden.py index 67b30e0c..9070ac05 100644 --- a/libs/vertexai/langchain_google_vertexai/model_garden.py +++ b/libs/vertexai/langchain_google_vertexai/model_garden.py @@ -148,7 +148,7 @@ def __init__(self, **kwargs: Any) -> None: @model_validator(mode="after") def validate_environment(self) -> Self: - from anthropic import ( # type: ignore[import-not-found] + from anthropic import ( # type: ignore[unused-ignore, import-not-found] AnthropicVertex, AsyncAnthropicVertex, ) diff --git a/libs/vertexai/langchain_google_vertexai/vectorstores/document_storage.py b/libs/vertexai/langchain_google_vertexai/vectorstores/document_storage.py index bfaa5cf8..2d6147b8 100644 --- a/libs/vertexai/langchain_google_vertexai/vectorstores/document_storage.py +++ b/libs/vertexai/langchain_google_vertexai/vectorstores/document_storage.py @@ -292,10 +292,15 @@ def mset(self, key_value_pairs: Sequence[Tuple[str, Document]]) -> None: entities = [] for key, document in zip(keys, documents): entity = self._client.entity( - key=key, exclude_from_indexes=self.exclude_from_indexes + key=key, + exclude_from_indexes=self.exclude_from_indexes + if self.exclude_from_indexes + else [], ) metadata_entity = self._client.entity( exclude_from_indexes=self.exclude_from_indexes + if self.exclude_from_indexes + else [] ) metadata_entity.update(document.metadata) entity[self._text_property_name] = document.page_content diff --git a/libs/vertexai/poetry.lock b/libs/vertexai/poetry.lock index e3697d77..d9e6edc2 100644 --- a/libs/vertexai/poetry.lock +++ b/libs/vertexai/poetry.lock @@ -2231,6 +2231,23 @@ pytest = [ {version = ">=6.2.4", markers = "python_version >= \"3.10\""}, ] +[[package]] +name = "pytest-retry" +version = "1.6.3" +description = "Adds the ability to retry flaky tests in CI environments" +optional = false +python-versions = ">=3.9" +files = [ + {file = "pytest_retry-1.6.3-py3-none-any.whl", hash = "sha256:e96f7df77ee70b0838d1085f9c3b8b5b7d74bf8947a0baf32e2b8c71b27683c8"}, + {file = "pytest_retry-1.6.3.tar.gz", hash = "sha256:36ccfa11c8c8f9ddad5e20375182146d040c20c4a791745139c5a99ddf1b557d"}, +] + +[package.dependencies] +pytest = ">=7.0.0" + +[package.extras] +dev = ["black", "flake8", "isort", "mypy"] + [[package]] name = "pytest-watcher" version = "0.3.5" @@ -2905,4 +2922,4 @@ mistral = ["langchain-mistralai"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "b047522f04bcdfd163bbe017fcb1f616b81f51dc342e61edc14e8f70d1a1bc2f" +content-hash = "ac58beb4bb70401121f19c64af1782eee2aebfc0d17d9fe46463559adceb066f" diff --git a/libs/vertexai/pyproject.toml b/libs/vertexai/pyproject.toml index d3e494b2..4213370c 100644 --- a/libs/vertexai/pyproject.toml +++ b/libs/vertexai/pyproject.toml @@ -29,6 +29,7 @@ optional = true pytest = "^7.3.0" pytest-xdist = "^3.0" pytest-order = "^1.3" +pytest-retry = "^1.6.3" freezegun = "^1.2.2" pytest-mock = "^3.10.0" syrupy = "^4.0.2" diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py index fc33d363..f9c860a7 100644 --- a/libs/vertexai/tests/integration_tests/test_chat_models.py +++ b/libs/vertexai/tests/integration_tests/test_chat_models.py @@ -693,6 +693,7 @@ class Response(BaseModel): assert isinstance(response, Response) +@pytest.mark.flaky(retries=3) @pytest.mark.release def test_chat_vertexai_gemini_function_calling_with_multiple_parts() -> None: @tool diff --git a/libs/vertexai/tests/integration_tests/test_evaluation.py b/libs/vertexai/tests/integration_tests/test_evaluation.py index e5757e34..0c866feb 100644 --- a/libs/vertexai/tests/integration_tests/test_evaluation.py +++ b/libs/vertexai/tests/integration_tests/test_evaluation.py @@ -29,6 +29,7 @@ def test_evaluate() -> None: @pytest.mark.extended +@pytest.mark.flaky(retries=3) def test_evaluate_strings() -> None: evaluator = VertexStringEvaluator( metric="safety", project_id=os.environ["PROJECT_ID"] @@ -40,6 +41,7 @@ def test_evaluate_strings() -> None: @pytest.mark.extended +@pytest.mark.flaky(retries=3) async def test_aevaluate_strings() -> None: evaluator = VertexStringEvaluator( metric="question_answering_quality", project_id=os.environ["PROJECT_ID"] diff --git a/libs/vertexai/tests/integration_tests/test_maas.py b/libs/vertexai/tests/integration_tests/test_maas.py index b5419f8f..046bd94f 100644 --- a/libs/vertexai/tests/integration_tests/test_maas.py +++ b/libs/vertexai/tests/integration_tests/test_maas.py @@ -18,10 +18,7 @@ ) model_names = _LLAMA_MODELS + _MISTRAL_MODELS -model_names_with_tools_support = _MISTRAL_MODELS + [ - "meta/llama3-405b-instruct-maas", - "meta/llama-3.2-90b-vision-instruct-maas", -] +model_names_with_tools_support = _MISTRAL_MODELS + ["meta/llama3-405b-instruct-maas"] @pytest.mark.extended @@ -45,8 +42,8 @@ async def test_agenerate(model_name: str) -> None: @pytest.mark.extended @pytest.mark.parametrize("model_name", model_names) def test_stream(model_name: str) -> None: - # b/371631576 - if model_name == "meta/llama3-8b-instruct-maas": + # streaming currently fails with mistral-nemo@2407 + if model_name == "mistral-nemo@2407": return llm = get_vertex_maas_model(model_name=model_name, location="us-central1") output = llm.stream("What is the meaning of life?") @@ -57,8 +54,8 @@ def test_stream(model_name: str) -> None: @pytest.mark.extended @pytest.mark.parametrize("model_name", model_names) async def test_astream(model_name: str) -> None: - # b/371631576 - if model_name == "meta/llama3-8b-instruct-maas": + # streaming currently fails with mistral-nemo@2407 + if model_name == "mistral-nemo@2407": return llm = get_vertex_maas_model(model_name=model_name, location="us-central1") output = llm.astream("What is the meaning of life?") @@ -68,6 +65,7 @@ async def test_astream(model_name: str) -> None: @pytest.mark.extended @pytest.mark.parametrize("model_name", model_names_with_tools_support) +@pytest.mark.flaky(retries=3) async def test_tools(model_name: str) -> None: @tool def search( diff --git a/libs/vertexai/tests/integration_tests/test_model_garden.py b/libs/vertexai/tests/integration_tests/test_model_garden.py index 832cfb6a..f84e11b4 100644 --- a/libs/vertexai/tests/integration_tests/test_model_garden.py +++ b/libs/vertexai/tests/integration_tests/test_model_garden.py @@ -178,6 +178,7 @@ def _check_tool_calls(response: BaseMessage, expected_name: str) -> None: @pytest.mark.extended +@pytest.mark.flaky(retries=3) def test_anthropic_tool_calling() -> None: project = os.environ["PROJECT_ID"] location = "us-east5" diff --git a/libs/vertexai/tests/integration_tests/test_vision_models.py b/libs/vertexai/tests/integration_tests/test_vision_models.py index c5670ded..ee4cfede 100644 --- a/libs/vertexai/tests/integration_tests/test_vision_models.py +++ b/libs/vertexai/tests/integration_tests/test_vision_models.py @@ -120,6 +120,7 @@ def test_vertex_ai_visual_qna_chat(base64_image: str): @pytest.mark.release +@pytest.mark.flaky(retries=3) def test_vertex_ai_image_generation_and_edition(): generator = VertexAIImageGeneratorChat()