fix: GenAI - Fixed the GenerativeModel's handling of tuned models from different region

PiperOrigin-RevId: 622336161

Author: Ark-kun

tests/unit/vertexai/test_generative_models.py

@@ -36,7 +36,9 @@
 
 
 _TEST_PROJECT = "test-project"
+_TEST_PROJECT2 = "test-project2"
 _TEST_LOCATION = "us-central1"
+_TEST_LOCATION2 = "europe-west4"
 
 
 _RESPONSE_TEXT_PART_STRUCT = {
@@ -283,6 +285,50 @@ def setup_method(self):
     def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)
 
+    @mock.patch.object(
+        target=prediction_service.PredictionServiceClient,
+        attribute="generate_content",
+        new=mock_generate_content,
+    )
+    @pytest.mark.parametrize(
+        "generative_models",
+        [generative_models, preview_generative_models],
+    )
+    def test_generative_model_constructor_model_name(
+        self, generative_models: generative_models
+    ):
+        project_location_prefix = (
+            f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/"
+        )
+
+        model_name1 = "gemini-pro"
+        model1 = generative_models.GenerativeModel(model_name1)
+        assert (
+            model1._prediction_resource_name
+            == project_location_prefix + "publishers/google/models/" + model_name1
+        )
+
+        model_name2 = "models/gemini-pro"
+        model2 = generative_models.GenerativeModel(model_name2)
+        assert (
+            model2._prediction_resource_name
+            == project_location_prefix + "publishers/google/" + model_name2
+        )
+
+        model_name3 = "publishers/some_publisher/models/some_model"
+        model3 = generative_models.GenerativeModel(model_name3)
+        assert model3._prediction_resource_name == project_location_prefix + model_name3
+
+        model_name4 = (
+            f"projects/{_TEST_PROJECT2}/locations/{_TEST_LOCATION2}/endpoints/endpoint1"
+        )
+        model4 = generative_models.GenerativeModel(model_name4)
+        assert model4._prediction_resource_name == model_name4
+        assert _TEST_LOCATION2 in model4._prediction_client._api_endpoint
+
+        with pytest.raises(ValueError):
+            generative_models.GenerativeModel("foo/bar/models/gemini-pro")
+
     @mock.patch.object(
         target=prediction_service.PredictionServiceClient,
         attribute="generate_content",

vertexai/generative_models/_generative_models.py

@@ -32,6 +32,7 @@
 )
 
 from google.cloud.aiplatform import initializer as aiplatform_initializer
+from google.cloud.aiplatform import utils as aiplatform_utils
 from google.cloud.aiplatform_v1beta1 import types as aiplatform_types
 from google.cloud.aiplatform_v1beta1.services import prediction_service
 from google.cloud.aiplatform_v1beta1.types import (
@@ -169,11 +170,20 @@ def __init__(
             prediction_resource_name = (
                 f"projects/{project}/locations/{location}/{model_name}"
             )
-        else:
+        elif model_name.startswith("projects/"):
             prediction_resource_name = model_name
+        else:
+            raise ValueError(
+                "model_name must be either a Model Garden model ID or a full resource name."
+            )
+
+        location = aiplatform_utils.extract_project_and_location_from_parent(
+            prediction_resource_name
+        )["location"]
 
         self._model_name = model_name
         self._prediction_resource_name = prediction_resource_name
+        self._location = location
         self._generation_config = generation_config
         self._safety_settings = safety_settings
         self._tools = tools
@@ -197,6 +207,7 @@ def _prediction_client(self) -> prediction_service.PredictionServiceClient:
             self._prediction_client_value = (
                 aiplatform_initializer.global_config.create_client(
                     client_class=prediction_service.PredictionServiceClient,
+                    location_override=self._location,
                     prediction_client=True,
                 )
             )
@@ -211,6 +222,7 @@ def _prediction_async_client(
             self._prediction_async_client_value = (
                 aiplatform_initializer.global_config.create_client(
                     client_class=prediction_service.PredictionServiceAsyncClient,
+                    location_override=self._location,
                     prediction_client=True,
                 )
             )