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

PiperOrigin-RevId: 836360793

Author: jsondai

vertexai/_genai/_evals_common.py

@@ -56,6 +56,22 @@
 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("Overriding location from %s to %s", api_client.location, 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 +731,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 +747,17 @@ 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.")
+
+    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 +1078,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 +1089,16 @@ 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.
     """
 
+    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

@@ -911,6 +911,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.
@@ -935,6 +936,8 @@ 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.
           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.
@@ -962,8 +965,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(
@@ -975,6 +979,7 @@ def evaluate(
             list[types.EvaluationDatasetOrDict],
         ],
         metrics: list[types.MetricOrDict] = None,
+        location: Optional[str] = None,
         config: Optional[types.EvaluateMethodConfigOrDict] = None,
         **kwargs,
     ) -> types.EvaluationResult:
@@ -984,6 +989,8 @@ 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.
           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
@@ -1029,6 +1036,7 @@ def evaluate(
             metrics=metrics,
             dataset_schema=config.dataset_schema,
             dest=config.dest,
+            location=location,
             **kwargs,
         )