feat: Add typing alias for models and their parameters

Testing for #4701

PiperOrigin-RevId: 704506046

Author: yeesian

noxfile.py

@@ -52,6 +52,7 @@
 )
 
 UNIT_TEST_PYTHON_VERSIONS = ["3.8", "3.9", "3.10", "3.11", "3.12"]
+UNIT_TEST_LANGCHAIN_PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
 UNIT_TEST_STANDARD_DEPENDENCIES = [
     "mock",
     "asyncmock",
@@ -275,7 +276,7 @@ def unit_ray(session, ray):
     )
 
 
-@nox.session(python=UNIT_TEST_PYTHON_VERSIONS)
+@nox.session(python=UNIT_TEST_LANGCHAIN_PYTHON_VERSIONS)
 def unit_langchain(session):
     # Install all test dependencies, then install this package in-place.
 

setup.py

@@ -123,6 +123,7 @@
 
 genai_requires = (
     "pydantic < 3",
+    "typing_extensions",
     "docstring_parser < 1",
 )
 
@@ -143,7 +144,8 @@
     "google-cloud-trace < 2",
     "opentelemetry-sdk < 2",
     "opentelemetry-exporter-gcp-trace < 2",
-    "pydantic >= 2.6.3, < 2.10",
+    "pydantic >= 2.6.3, < 3",
+    "typing_extensions",
 ]
 
 evaluation_extra_require = [

testing/constraints-langchain.txt

@@ -1,4 +1,3 @@
 langchain
 langchain-core
-langchain-google-vertexai
-pydantic<2.10
+langchain-google-vertexai

vertexai/generative_models/_generative_models.py

@@ -117,6 +117,51 @@
     ],
 ]
 
+try:
+    # For Pydantic to resolve the forward references inside these aliases.
+    from typing_extensions import TypeAliasType
+
+    PartsType = TypeAliasType(
+        "PartsType",
+        Union[
+            str,
+            "Image",
+            "Part",
+            List[Union[str, "Image", "Part"]],
+        ],
+    )
+    ContentsType = TypeAliasType(
+        "ContentsType",
+        Union[
+            List["Content"],
+            List[ContentDict],
+            str,
+            "Image",
+            "Part",
+            List[Union[str, "Image", "Part"]],
+        ],
+    )
+    GenerationConfigType = TypeAliasType(
+        "GenerationConfigType",
+        Union[
+            "GenerationConfig",
+            GenerationConfigDict,
+        ],
+    )
+    SafetySettingsType = TypeAliasType(
+        "SafetySettingsType",
+        Union[
+            List["SafetySetting"],
+            Dict[
+                gapic_content_types.HarmCategory,
+                gapic_content_types.SafetySetting.HarmBlockThreshold,
+            ],
+        ],
+    )
+except ImportError:
+    # Use existing definitions if typing_extensions is not available.
+    pass
+
 
 def _reconcile_model_name(model_name: str, project: str, location: str) -> str:
     """Returns a model name that's one of the following: