chore(vertex): add migration TODOs, suppress warning (#1196)

mdrxy · web-flow · commit 9ea5b893ea22 · 2025-09-18T15:03:21.000-04:00
diff --git a/libs/vertexai/langchain_google_vertexai/_base.py b/libs/vertexai/langchain_google_vertexai/_base.py
@@ -15,7 +15,7 @@
     cast,
 )
 
-import vertexai  # type: ignore[import-untyped]
+import vertexai
 from google.api_core.client_options import ClientOptions
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform.constants import base as constants
@@ -42,8 +42,8 @@
 from langchain_core.outputs import Generation, LLMResult
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 from typing_extensions import Literal, Self
-from vertexai.generative_models._generative_models import (  # type: ignore
-    SafetySettingsType,
+from vertexai.generative_models._generative_models import (
+    SafetySettingsType,  # TODO: migrate to google-genai since this is deprecated
 )
 
 from langchain_google_vertexai._client_utils import (
diff --git a/libs/vertexai/langchain_google_vertexai/_enums.py b/libs/vertexai/langchain_google_vertexai/_enums.py
@@ -1,6 +1,6 @@
 from google.cloud.aiplatform_v1beta1.types.content import Modality
-from vertexai.generative_models import (  # type: ignore
-    HarmBlockThreshold,
+from vertexai.generative_models import (
+    HarmBlockThreshold,  # TODO: migrate to google-genai since this is deprecated
     HarmCategory,
     SafetySetting,
 )
diff --git a/libs/vertexai/langchain_google_vertexai/_image_utils.py b/libs/vertexai/langchain_google_vertexai/_image_utils.py
@@ -11,7 +11,9 @@
 import requests
 from google.cloud import storage
 from google.cloud.aiplatform_v1beta1.types.content import Part as GapicPart
-from vertexai.generative_models import Image, Part  # type: ignore
+from vertexai.generative_models import Image, Part
+
+# TODO: migrate to google-genai since vertexai.generative_models is deprecated
 
 
 class Route(Enum):
diff --git a/libs/vertexai/langchain_google_vertexai/_utils.py b/libs/vertexai/langchain_google_vertexai/_utils.py
@@ -14,12 +14,12 @@
     AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
-from vertexai.generative_models import (  # type: ignore[import-untyped]
-    Candidate,
+from vertexai.generative_models import (
+    Candidate,  # TODO: migrate to google-genai since this is deprecated
     Image,
 )
-from vertexai.language_models import (  # type: ignore[import-untyped]
-    TextGenerationResponse,
+from vertexai.language_models import (
+    TextGenerationResponse,  # TODO: migrate to google-genai since this is deprecated
 )
 
 from langchain_google_vertexai._retry import create_base_retry_decorator
diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py
@@ -74,18 +74,18 @@
 )
 from langchain_core.utils.pydantic import is_basemodel_subclass
 from langchain_core.utils.utils import _build_model_kwargs
-from vertexai.generative_models import (  # type: ignore
-    Tool as VertexTool,
+from vertexai.generative_models import (
+    Tool as VertexTool,  # TODO: migrate to google-genai since this is deprecated
 )
-from vertexai.generative_models._generative_models import (  # type: ignore
-    ToolConfig,
+from vertexai.generative_models._generative_models import (
+    ToolConfig,  # TODO: migrate to google-genai since this is deprecated
     SafetySettingsType,
     GenerationConfigType,
     GenerationResponse,
     _convert_schema_dict_to_gapic,
 )
-from vertexai.language_models import (  # type: ignore
-    ChatMessage,
+from vertexai.language_models import (
+    ChatMessage,  # TODO: migrate to google-genai since this is deprecated
     InputOutputTextPair,
 )
 from google.cloud.aiplatform_v1.types import (
diff --git a/libs/vertexai/langchain_google_vertexai/embeddings.py b/libs/vertexai/langchain_google_vertexai/embeddings.py
@@ -22,12 +22,12 @@
 from langchain_core.language_models.llms import create_base_retry_decorator
 from pydantic import ConfigDict, model_validator
 from typing_extensions import Self
-from vertexai.language_models import (  # type: ignore
-    TextEmbeddingInput,
+from vertexai.language_models import (
+    TextEmbeddingInput,  # TODO: migrate to google-genai since this is deprecated
     TextEmbeddingModel,
 )
-from vertexai.vision_models import (  # type: ignore
-    Image,
+from vertexai.vision_models import (
+    Image,  # TODO: migrate to google-genai since this is deprecated
     MultiModalEmbeddingModel,
     MultiModalEmbeddingResponse,
 )
diff --git a/libs/vertexai/langchain_google_vertexai/functions_utils.py b/libs/vertexai/langchain_google_vertexai/functions_utils.py
@@ -17,7 +17,7 @@
 )
 
 import google.cloud.aiplatform_v1beta1.types as gapic
-import vertexai.generative_models as vertexai  # type: ignore
+import vertexai.generative_models as vertexai  # TODO: migrate to google-genai
 from google.cloud.aiplatform_v1beta1.types import (
     ToolConfig as GapicToolConfig,
 )
diff --git a/libs/vertexai/langchain_google_vertexai/utils.py b/libs/vertexai/langchain_google_vertexai/utils.py
@@ -2,7 +2,7 @@
 from typing import List, Optional
 
 from langchain_core.messages import BaseMessage
-from vertexai.preview import caching  # type: ignore
+from vertexai.preview import caching
 
 from langchain_google_vertexai._image_utils import ImageBytesLoader
 from langchain_google_vertexai.chat_models import (
diff --git a/libs/vertexai/langchain_google_vertexai/vision_models.py b/libs/vertexai/langchain_google_vertexai/vision_models.py
@@ -11,7 +11,7 @@
 from langchain_core.outputs.chat_generation import ChatGeneration
 from langchain_core.outputs.generation import Generation
 from pydantic import BaseModel, ConfigDict, Field
-from vertexai.vision_models import (  # type: ignore[import-untyped]
+from vertexai.vision_models import (  # TODO: migrate to google-genai
     GeneratedImage,
     Image,
     ImageGenerationModel,
diff --git a/libs/vertexai/pyproject.toml b/libs/vertexai/pyproject.toml
@@ -146,6 +146,11 @@ warn_unused_ignores = true
 # strict = true
 # disallow_untyped_defs = true
 
+[[tool.mypy.overrides]]
+# TODO: remove after migrating to google-genai
+module = "vertexai.*"
+ignore_missing_imports = true
+
 [tool.coverage.run]
 omit = ["tests/*"]
 
diff --git a/libs/vertexai/tests/integration_tests/test_embeddings.py b/libs/vertexai/tests/integration_tests/test_embeddings.py
@@ -5,9 +5,9 @@
 """
 
 import pytest
-from vertexai.language_models import TextEmbeddingModel  # type: ignore[import-untyped]
-from vertexai.vision_models import (  # type: ignore[import-untyped]
-    MultiModalEmbeddingModel,
+from vertexai.language_models import TextEmbeddingModel  # TODO: migrate to google-genai
+from vertexai.vision_models import (
+    MultiModalEmbeddingModel,  # TODO: migrate to google-genai since this is deprecated
 )
 
 from langchain_google_vertexai.embeddings import (
diff --git a/libs/vertexai/tests/integration_tests/test_vectorstores.py b/libs/vertexai/tests/integration_tests/test_vectorstores.py
@@ -15,7 +15,7 @@
 from uuid import uuid4
 
 import pytest
-import vertexai  # type: ignore[import-untyped]
+import vertexai
 from google.cloud import storage  # type: ignore[attr-defined, unused-ignore]
 from google.cloud.aiplatform.matching_engine import (
     MatchingEngineIndex,
diff --git a/libs/vertexai/tests/unit_tests/test_chat_models.py b/libs/vertexai/tests/unit_tests/test_chat_models.py
@@ -33,11 +33,11 @@
     PydanticToolsParser,
 )
 from pydantic import BaseModel
-from vertexai.generative_models import (  # type: ignore
-    SafetySetting as VertexSafetySetting,
+from vertexai.generative_models import (
+    SafetySetting as VertexSafetySetting,  # TODO: migrate to google-genai
 )
-from vertexai.language_models import (  # type: ignore
-    ChatMessage,
+from vertexai.language_models import (
+    ChatMessage,  # TODO: migrate to google-genai since this is deprecated
     InputOutputTextPair,
 )
 
diff --git a/libs/vertexai/tests/unit_tests/test_function_utils.py b/libs/vertexai/tests/unit_tests/test_function_utils.py
@@ -9,7 +9,7 @@
 
 import google.cloud.aiplatform_v1beta1.types as gapic
 import pytest
-import vertexai.generative_models as vertexai  # type: ignore
+import vertexai.generative_models as vertexai  # TODO: migrate to google-genai
 from google.cloud.aiplatform_v1beta1.types import (
     FunctionCallingConfig as GapicFunctionCallingConfig,
 )
diff --git a/libs/vertexai/tests/unit_tests/test_vision_models.py b/libs/vertexai/tests/unit_tests/test_vision_models.py
@@ -1,26 +1,35 @@
+import warnings
+
 import pytest
-from vertexai.vision_models import Image  # type: ignore[import-untyped]
+from vertexai.vision_models import (
+    Image,  # TODO: migrate to google-genai since this is deprecated
+)
 
 from langchain_google_vertexai.vision_models import _BaseImageTextModel
 
 
 def test_get_image_from_message_part(base64_image: str) -> None:
-    model = _BaseImageTextModel()
+    # TODO: Remove this warning suppression when migrating to google-genai
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message=".*deprecated.*", category=UserWarning
+        )
+        model = _BaseImageTextModel()
 
-    # Should work with a well formatted dictionary:
-    message = {"type": "image_url", "image_url": {"url": base64_image}}
-    image = model._get_image_from_message_part(message)
-    assert isinstance(image, Image)
+        # Should work with a well formatted dictionary:
+        message = {"type": "image_url", "image_url": {"url": base64_image}}
+        image = model._get_image_from_message_part(message)
+        assert isinstance(image, Image)
 
-    # Should not work with a simple string
-    simple_string = base64_image
-    image = model._get_image_from_message_part(simple_string)
-    assert image is None
+        # Should not work with a simple string
+        simple_string = base64_image
+        image = model._get_image_from_message_part(simple_string)
+        assert image is None
 
-    # Should not work with a string message
-    message = {"type": "text", "text": "I'm a text message"}
-    image = model._get_image_from_message_part(message)
-    assert image is None
+        # Should not work with a string message
+        message = {"type": "text", "text": "I'm a text message"}
+        image = model._get_image_from_message_part(message)
+        assert image is None
 
 
 def test_get_text_from_message_part() -> None:

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`from google.cloud.aiplatform_v1beta1.types.content import Modality`
`2`		`-from vertexai.generative_models import ( # type: ignore`
`3`		`- HarmBlockThreshold,`
	`2`	`+from vertexai.generative_models import (`
	`3`	`+ HarmBlockThreshold, # TODO: migrate to google-genai since this is deprecated`
`4`	`4`	`HarmCategory,`
`5`	`5`	`SafetySetting,`
`6`	`6`	`)`
Original file line number	Diff line number	Diff line change
`@@ -14,12 +14,12 @@`
`14`	`14`	`AsyncCallbackManagerForLLMRun,`
`15`	`15`	`CallbackManagerForLLMRun,`
`16`	`16`	`)`
`17`		`-from vertexai.generative_models import ( # type: ignore[import-untyped]`
`18`		`- Candidate,`
	`17`	`+from vertexai.generative_models import (`
	`18`	`+ Candidate, # TODO: migrate to google-genai since this is deprecated`
`19`	`19`	`Image,`
`20`	`20`	`)`
`21`		`-from vertexai.language_models import ( # type: ignore[import-untyped]`
`22`		`- TextGenerationResponse,`
	`21`	`+from vertexai.language_models import (`
	`22`	`+ TextGenerationResponse, # TODO: migrate to google-genai since this is deprecated`
`23`	`23`	`)`
`24`	`24`
`25`	`25`	`from langchain_google_vertexai._retry import create_base_retry_decorator`
Original file line number	Diff line number	Diff line change
`@@ -17,7 +17,7 @@`
`17`	`17`	`)`
`18`	`18`
`19`	`19`	`import google.cloud.aiplatform_v1beta1.types as gapic`
`20`		`-import vertexai.generative_models as vertexai # type: ignore`
	`20`	`+import vertexai.generative_models as vertexai # TODO: migrate to google-genai`
`21`	`21`	`from google.cloud.aiplatform_v1beta1.types import (`
`22`	`22`	`ToolConfig as GapicToolConfig,`
`23`	`23`	`)`