feat: GenAI Client(evals) - Add location override parameter to run_inference and evaluate methods

PiperOrigin-RevId: 839481916

Author: jsondai

tests/unit/vertexai/genai/test_evals.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+# pylint: disable=protected-access,bad-continuation,
 import importlib
 import json
 import os
@@ -69,6 +70,7 @@ def mock_api_client_fixture():
     )
     mock_client._credentials.universe_domain = "googleapis.com"
     mock_client._evals_client = mock.Mock(spec=evals.Evals)
+    mock_client._http_options = None
     return mock_client
 
 
@@ -139,6 +141,46 @@ def mock_evaluate_instances_side_effect(*args, **kwargs):
         }
 
 
+class TestGetApiClientWithLocation:
+    @mock.patch("vertexai._genai._evals_common.vertexai.Client")
+    def test_get_api_client_with_location_override(
+        self, mock_vertexai_client, mock_api_client_fixture
+    ):
+        mock_api_client_fixture.location = "us-central1"
+        new_location = "europe-west1"
+        _evals_common._get_api_client_with_location(
+            mock_api_client_fixture, new_location
+        )
+        mock_vertexai_client.assert_called_once_with(
+            project=mock_api_client_fixture.project,
+            location=new_location,
+            credentials=mock_api_client_fixture._credentials,
+            http_options=mock_api_client_fixture._http_options,
+        )
+
+    @mock.patch("vertexai._genai._evals_common.vertexai.Client")
+    def test_get_api_client_with_same_location(
+        self, mock_vertexai_client, mock_api_client_fixture
+    ):
+        mock_api_client_fixture.location = "us-central1"
+        new_location = "us-central1"
+        _evals_common._get_api_client_with_location(
+            mock_api_client_fixture, new_location
+        )
+        mock_vertexai_client.assert_not_called()
+
+    @mock.patch("vertexai._genai._evals_common.vertexai.Client")
+    def test_get_api_client_with_none_location(
+        self, mock_vertexai_client, mock_api_client_fixture
+    ):
+        mock_api_client_fixture.location = "us-central1"
+        new_location = None
+        _evals_common._get_api_client_with_location(
+            mock_api_client_fixture, new_location
+        )
+        mock_vertexai_client.assert_not_called()
+
+
 class TestEvals:
     """Unit tests for the GenAI client."""
 
@@ -4984,7 +5026,7 @@ def test_execute_evaluation_adds_creation_timestamp(
         frozenset(["summarization_quality"]),
     )
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")
+    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")  # fmt: skip
     def test_predefined_metric_retry_on_resource_exhausted(
         self,
         mock_private_evaluate_instances,
@@ -5037,7 +5079,7 @@ def test_predefined_metric_retry_on_resource_exhausted(
         frozenset(["summarization_quality"]),
     )
     @mock.patch("time.sleep", return_value=None)
-    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")
+    @mock.patch("vertexai._genai.evals.Evals._evaluate_instances")  # fmt: skip
     def test_predefined_metric_retry_fail_on_resource_exhausted(
         self,
         mock_private_evaluate_instances,

vertexai/_genai/_evals_common.py

@@ -56,6 +56,26 @@
 AGENT_MAX_WORKERS = 10
 
 
+def _get_api_client_with_location(
+    api_client: BaseApiClient, location: Optional[str]
+) -> BaseApiClient:
+    """Returns a new API client with the specified location."""
+    if not location or location == api_client.location:
+        return api_client
+
+    logger.info(
+        "Model endpoint location set to %s, overriding client location %s for this API call.",
+        location,
+        api_client.location,
+    )
+    return vertexai.Client(
+        project=api_client.project,
+        location=location,
+        credentials=api_client._credentials,
+        http_options=api_client._http_options,
+    )._api_client
+
+
 def _get_agent_engine_instance(
     agent_name: str, api_client: BaseApiClient
 ) -> Union[types.AgentEngine, Any]:
@@ -715,6 +735,7 @@ def _execute_inference(
     dest: Optional[str] = None,
     config: Optional[genai_types.GenerateContentConfig] = None,
     prompt_template: Optional[Union[str, types.PromptTemplateOrDict]] = None,
+    location: Optional[str] = None,
 ) -> pd.DataFrame:
     """Executes inference on a given dataset using the specified model.
 
@@ -730,12 +751,18 @@ def _execute_inference(
           representing a file path or a GCS URI.
         config: The generation configuration for the model.
         prompt_template: The prompt template to use for inference.
+        location: The location to use for the inference. If not specified, the
+          location configured in the client will be used.
 
     Returns:
         A pandas DataFrame containing the inference results.
     """
     if not api_client:
         raise ValueError("'api_client' instance must be provided.")
+
+    if location:
+        api_client = _get_api_client_with_location(api_client, location)
+
     prompt_dataset = _load_dataframe(api_client, src)
     if prompt_template:
         logger.info("Applying prompt template...")
@@ -1056,6 +1083,7 @@ def _execute_evaluation(  # type: ignore[no-untyped-def]
     metrics: list[types.Metric],
     dataset_schema: Optional[Literal["GEMINI", "FLATTEN", "OPENAI"]] = None,
     dest: Optional[str] = None,
+    location: Optional[str] = None,
     **kwargs,
 ) -> types.EvaluationResult:
     """Evaluates a dataset using the provided metrics.
@@ -1066,12 +1094,17 @@ def _execute_evaluation(  # type: ignore[no-untyped-def]
         metrics: The metrics to evaluate the dataset against.
         dataset_schema: The schema of the dataset.
         dest: The destination to save the evaluation results.
+        location: The location to use for the evaluation. If not specified, the
+          location configured in the client will be used.
         **kwargs: Extra arguments to pass to evaluation, such as `agent_info`.
 
     Returns:
         The evaluation result.
     """
 
+    if location:
+        api_client = _get_api_client_with_location(api_client, location)
+
     logger.info("Preparing dataset(s) and metrics...")
     if isinstance(dataset, types.EvaluationDataset):
         dataset_list = [dataset]

vertexai/_genai/evals.py

@@ -904,6 +904,7 @@ def run_inference(
         src: Union[str, pd.DataFrame, types.EvaluationDataset],
         model: Optional[Union[str, Callable[[Any], Any]]] = None,
         agent: Optional[Union[str, types.AgentEngine]] = None,
+        location: Optional[str] = None,
         config: Optional[types.EvalRunInferenceConfigOrDict] = None,
     ) -> types.EvaluationDataset:
         """Runs inference on a dataset for evaluation.
@@ -928,6 +929,10 @@ def run_inference(
                 `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine_id}`,
                 run_inference will fetch the agent engine from the resource name.
               - Or `types.AgentEngine` object.
+          location: The location to use for the inference. If not specified, the
+                location configured in the client will be used. If specified,
+                this will override the location set in `vertexai.Client` only
+                for this API call.
           config: The optional configuration for the inference run. Must be a dict or
               `types.EvalRunInferenceConfig` type.
                 - dest: The destination path for storage of the inference results.
@@ -955,8 +960,9 @@ def run_inference(
             agent_engine=agent,
             src=src,
             dest=config.dest,
-            config=config.generate_content_config,
             prompt_template=config.prompt_template,
+            location=location,
+            config=config.generate_content_config,
         )
 
     def evaluate(
@@ -968,6 +974,7 @@ def evaluate(
             list[types.EvaluationDatasetOrDict],
         ],
         metrics: list[types.MetricOrDict] = None,
+        location: Optional[str] = None,
         config: Optional[types.EvaluateMethodConfigOrDict] = None,
         **kwargs,
     ) -> types.EvaluationResult:
@@ -977,6 +984,10 @@ def evaluate(
           dataset: The dataset(s) to evaluate. Can be a pandas DataFrame, a single
             `types.EvaluationDataset` or a list of `types.EvaluationDataset`.
           metrics: The list of metrics to use for evaluation.
+          location: The location to use for the evaluation service. If not specified,
+             the location configured in the client will be used. If specified,
+             this will override the location set in `vertexai.Client` only for
+             this API call.
           config: Optional configuration for the evaluation. Can be a dictionary or a
             `types.EvaluateMethodConfig` object.
             - dataset_schema: Schema to use for the dataset. If not specified, the
@@ -1022,6 +1033,7 @@ def evaluate(
             metrics=metrics,
             dataset_schema=config.dataset_schema,
             dest=config.dest,
+            location=location,
             **kwargs,
         )