feat: Add code_execution_region to SDK custom code execution metric.

agentplatform/_genai/_transformers.py

@@ -31,114 +31,118 @@ def t_metrics(
     metrics: "list[types.MetricSubclass]",
     set_default_aggregation_metrics: bool = False,
 ) -> list[dict[str, Any]]:
-    """Prepares the metric payload for the evaluation request.
+  """Prepares the metric payload for the evaluation request.
 
     Args:
         metrics: A list of metrics used for evaluation.
         set_default_aggregation_metrics: Whether to set default aggregation metrics.
     Returns:
         A list of resolved metric payloads for the evaluation request.
     """
-    metrics_payload = []
+  metrics_payload = []
 
-    for metric in metrics:
-        metric_payload_item: dict[str, Any] = {}
+  for metric in metrics:
+    metric_payload_item: dict[str, Any] = {}
 
-        metric_id = getv(metric, ["metric"]) or getv(metric, ["name"])
-        metric_name = metric_id.lower() if metric_id else None
+    metric_id = getv(metric, ["metric"]) or getv(metric, ["name"])
+    metric_name = metric_id.lower() if metric_id else None
 
-        if set_default_aggregation_metrics:
-            metric_payload_item["aggregation_metrics"] = [
+    if set_default_aggregation_metrics:
+      metric_payload_item["aggregation_metrics"] = [
                 "AVERAGE",
                 "STANDARD_DEVIATION",
             ]
 
-        if metric_name == "exact_match":
-            metric_payload_item["exact_match_spec"] = {}
-        elif metric_name == "bleu":
-            metric_payload_item["bleu_spec"] = {}
-        elif metric_name and metric_name.startswith("rouge"):
-            rouge_type = metric_name.replace("_", "")
-            metric_payload_item["rouge_spec"] = {"rouge_type": rouge_type}
-        # API Pre-defined metrics
-        elif (
+    if metric_name == "exact_match":
+      metric_payload_item["exact_match_spec"] = {}
+    elif metric_name == "bleu":
+      metric_payload_item["bleu_spec"] = {}
+    elif metric_name and metric_name.startswith("rouge"):
+      rouge_type = metric_name.replace("_", "")
+      metric_payload_item["rouge_spec"] = {"rouge_type": rouge_type}
+    # API Pre-defined metrics
+    elif (
             metric_name and metric_name in _evals_constant.SUPPORTED_PREDEFINED_METRICS
         ):
-            metric_payload_item["predefined_metric_spec"] = {
+      metric_payload_item["predefined_metric_spec"] = {
                 "metric_spec_name": metric_name,
                 "metric_spec_parameters": metric.metric_spec_parameters,
             }
-        # Custom Code Execution Metric
-        elif (
+    # Custom Code Execution Metric
+    elif (
             hasattr(metric, "remote_custom_function") and metric.remote_custom_function
         ):
-            metric_payload_item["custom_code_execution_spec"] = {
-                "evaluation_function": metric.remote_custom_function
-            }
-        elif (
-            isinstance(metric, types.CodeExecutionMetric)
-            or (
-                isinstance(metric, types.Metric)
-                and isinstance(getattr(metric, "custom_function", None), str)
-            )
-        ) and getattr(metric, "custom_function", None):
-            metric_payload_item["custom_code_execution_spec"] = {
-                "evaluation_function": metric.custom_function
-            }
-        # LLM-based metrics
-        elif hasattr(metric, "prompt_template") and metric.prompt_template:
-            llm_based_spec: dict[str, Any] = {
-                "metric_prompt_template": metric.prompt_template
-            }
-            system_instruction = getv(metric, ["judge_model_system_instruction"])
-            if system_instruction:
-                llm_based_spec["system_instruction"] = system_instruction
-            rubric_group_name = getv(metric, ["rubric_group_name"])
-            if rubric_group_name:
-                llm_based_spec["rubric_group_key"] = rubric_group_name
-            return_raw_output = getv(metric, ["return_raw_output"])
-            if return_raw_output:
-                llm_based_spec["custom_output_format_config"] = {
+      spec: dict[str, Any] = {
+          "evaluation_function": metric.remote_custom_function
+      }
+      if getattr(metric, "code_execution_region", None):
+        spec["code_execution_region"] = metric.code_execution_region
+      metric_payload_item["custom_code_execution_spec"] = spec
+    elif (
+        isinstance(metric, types.CodeExecutionMetric)
+        or (
+            isinstance(metric, types.Metric)
+            and isinstance(getattr(metric, "custom_function", None), str)
+        )
+    ) and getattr(metric, "custom_function", None):
+      spec = {"evaluation_function": metric.custom_function}
+      if getattr(metric, "code_execution_region", None):
+        spec["code_execution_region"] = metric.code_execution_region
+      metric_payload_item["custom_code_execution_spec"] = spec
+    # LLM-based metrics
+    elif hasattr(metric, "prompt_template") and metric.prompt_template:
+      llm_based_spec: dict[str, Any] = {
+          "metric_prompt_template": metric.prompt_template
+      }
+      system_instruction = getv(metric, ["judge_model_system_instruction"])
+      if system_instruction:
+        llm_based_spec["system_instruction"] = system_instruction
+      rubric_group_name = getv(metric, ["rubric_group_name"])
+      if rubric_group_name:
+        llm_based_spec["rubric_group_key"] = rubric_group_name
+      return_raw_output = getv(metric, ["return_raw_output"])
+      if return_raw_output:
+        llm_based_spec["custom_output_format_config"] = {
                     "return_raw_output": return_raw_output
                 }
 
-            autorater_config: dict[str, Any] = {}
-            if hasattr(metric, "judge_model") and metric.judge_model:
-                autorater_config["autorater_model"] = metric.judge_model
-            if (
+      autorater_config: dict[str, Any] = {}
+      if hasattr(metric, "judge_model") and metric.judge_model:
+        autorater_config["autorater_model"] = metric.judge_model
+      if (
                 hasattr(metric, "judge_model_generation_config")
                 and metric.judge_model_generation_config
             ):
-                autorater_config["generation_config"] = (
+        autorater_config["generation_config"] = (
                     metric.judge_model_generation_config
                 )
-            if (
+      if (
                 hasattr(metric, "judge_model_sampling_count")
                 and metric.judge_model_sampling_count
             ):
-                autorater_config["sampling_count"] = metric.judge_model_sampling_count
+        autorater_config["sampling_count"] = metric.judge_model_sampling_count
 
-            if autorater_config:
-                llm_based_spec["judge_autorater_config"] = autorater_config
+      if autorater_config:
+        llm_based_spec["judge_autorater_config"] = autorater_config
 
-            result_parsing_function = getv(metric, ["result_parsing_function"])
-            if result_parsing_function:
-                llm_based_spec["result_parser_config"] = {
+      result_parsing_function = getv(metric, ["result_parsing_function"])
+      if result_parsing_function:
+        llm_based_spec["result_parser_config"] = {
                     "custom_code_parser_config": {
                         "parsing_function": result_parsing_function
                     }
                 }
 
-            metric_payload_item["llm_based_metric_spec"] = llm_based_spec
-        elif getattr(metric, "metric_resource_name", None) is not None:
-            # Safe pass
-            pass
-        else:
-            raise ValueError(
+      metric_payload_item["llm_based_metric_spec"] = llm_based_spec
+    elif getattr(metric, "metric_resource_name", None) is not None:
+      # Safe pass
+      pass
+    else:
+      raise ValueError(
                 f"Unsupported metric type or invalid metric name: {metric_name}"
             )
-        metrics_payload.append(metric_payload_item)
-    return metrics_payload
+    metrics_payload.append(metric_payload_item)
+  return metrics_payload
 
 
 def t_metric_sources(metrics: list[Any]) -> list[dict[str, Any]]:

agentplatform/_genai/evals.py

@@ -191,44 +191,58 @@ def _CustomCodeExecutionSpec_from_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
-    to_object: dict[str, Any] = {}
-    if getv(from_object, ["evaluationFunction"]) is not None:
-        setv(
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ["evaluationFunction"]) is not None:
+    setv(
             to_object,
             ["evaluation_function"],
             getv(from_object, ["evaluationFunction"]),
         )
 
-    if getv(from_object, ["evaluation_function"]) is not None:
-        setv(
+  if getv(from_object, ["evaluation_function"]) is not None:
+    setv(
             to_object,
             ["remote_custom_function"],
             getv(from_object, ["evaluation_function"]),
         )
 
-    return to_object
+  if getv(from_object, ["codeExecutionRegion"]) is not None:
+    setv(
+        to_object,
+        ["code_execution_region"],
+        getv(from_object, ["codeExecutionRegion"]),
+    )
+
+  return to_object
 
 
 def _CustomCodeExecutionSpec_to_vertex(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
 ) -> dict[str, Any]:
-    to_object: dict[str, Any] = {}
-    if getv(from_object, ["evaluation_function"]) is not None:
-        setv(
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ["evaluation_function"]) is not None:
+    setv(
             to_object,
             ["evaluationFunction"],
             getv(from_object, ["evaluation_function"]),
         )
 
-    if getv(from_object, ["remote_custom_function"]) is not None:
-        setv(
+  if getv(from_object, ["remote_custom_function"]) is not None:
+    setv(
             to_object,
             ["evaluation_function"],
             getv(from_object, ["remote_custom_function"]),
         )
 
-    return to_object
+  if getv(from_object, ["code_execution_region"]) is not None:
+    setv(
+        to_object,
+        ["codeExecutionRegion"],
+        getv(from_object, ["code_execution_region"]),
+    )
+
+  return to_object
 
 
 def _DeleteEvaluationExperimentParameters_to_vertex(

agentplatform/_genai/types/common.py

@@ -1998,6 +1998,10 @@ class Metric(_common.BaseModel):
         default=None,
         description="""The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service.""",
     )
+    code_execution_region: Optional[str] = Field(
+        default=None,
+        description="""Optional. The region to use for code execution. If set, the Code Execution Sandbox will be invoked in the specified region regardless of the request's originating region. Supported regions: us-central1, us-east1, us-east4, us-west1, us-west4, southamerica-east1, europe-west2, europe-west3, asia-east1, asia-south1, asia-southeast1. If unset, the request's originating region is used.""",
+    )
     judge_model: Optional[str] = Field(
         default=None, description="""The judge model for the metric."""
     )
@@ -2238,6 +2242,9 @@ class MetricDict(TypedDict, total=False):
     remote_custom_function: Optional[str]
     """The evaluation function for the custom code execution metric. This custom code is run remotely in the evaluation service."""
 
+    code_execution_region: Optional[str]
+    """Optional. The region to use for code execution. If set, the Code Execution Sandbox will be invoked in the specified region regardless of the request's originating region. Supported regions: us-central1, us-east1, us-east4, us-west1, us-west4, southamerica-east1, europe-west2, europe-west3, asia-east1, asia-south1, asia-southeast1. If unset, the request's originating region is used."""
+
     judge_model: Optional[str]
     """The judge model for the metric."""
 
@@ -2350,6 +2357,10 @@ def evaluate(instance: dict[str, Any]) -> float:
   Instance is the evaluation instance, any fields populated in the instance
   are available to the function as instance[field_name].""",
     )
+    code_execution_region: Optional[str] = Field(
+        default=None,
+        description="""Optional. The region to use for code execution. If set, the Code Execution Sandbox will be invoked in the specified region regardless of the request's originating region. Supported regions: us-central1, us-east1, us-east4, us-west1, us-west4, southamerica-east1, europe-west2, europe-west3, asia-east1, asia-south1, asia-southeast1. If unset, the request's originating region is used.""",
+    )
 
 
 class CustomCodeExecutionSpecDict(TypedDict, total=False):
@@ -2369,6 +2380,9 @@ def evaluate(instance: dict[str, Any]) -> float:
   Instance is the evaluation instance, any fields populated in the instance
   are available to the function as instance[field_name]."""
 
+    code_execution_region: Optional[str]
+    """Optional. The region to use for code execution. If set, the Code Execution Sandbox will be invoked in the specified region regardless of the request's originating region. Supported regions: us-central1, us-east1, us-east4, us-west1, us-west4, southamerica-east1, europe-west2, europe-west3, asia-east1, asia-south1, asia-southeast1. If unset, the request's originating region is used."""
+
 
 CustomCodeExecutionSpecOrDict = Union[
     CustomCodeExecutionSpec, CustomCodeExecutionSpecDict

tests/unit/agentplatform/genai/replays/test_custom_code_execution_metric.py

@@ -43,40 +43,81 @@ def evaluate(instance):
     ],
 )
 def test_custom_code_execution(client, custom_metric):
-    """Tests that custom code execution metric produces a correctly structured EvaluationResult."""
+  """Tests that custom code execution metric produces a correctly structured EvaluationResult."""
 
-    prompts_df = pd.DataFrame(
+  prompts_df = pd.DataFrame(
         {
             "prompt": ["What is 2+2?", "What is 3+3?"],
             "response": ["4", "5"],
             "reference": ["4", "6"],
         }
     )
 
-    eval_dataset = types.EvaluationDataset(
+  eval_dataset = types.EvaluationDataset(
         eval_dataset_df=prompts_df,
         candidate_name="test_model",
     )
 
-    evaluation_result = client.evals.evaluate(
+  evaluation_result = client.evals.evaluate(
         dataset=eval_dataset,
         metrics=[custom_metric],
     )
 
-    assert isinstance(evaluation_result, types.EvaluationResult)
-
-    assert evaluation_result.summary_metrics is not None
-    assert evaluation_result.summary_metrics
-    for summary in evaluation_result.summary_metrics:
-        assert isinstance(summary, types.AggregatedMetricResult)
-        assert summary.metric_name == "my_custom_code_metric"
-
-    assert evaluation_result.eval_case_results is not None
-    assert evaluation_result.eval_case_results
-    for case_result in evaluation_result.eval_case_results:
-        assert isinstance(case_result, types.EvalCaseResult)
-        assert case_result.eval_case_index is not None
-        assert case_result.response_candidate_results is not None
+  assert isinstance(evaluation_result, types.EvaluationResult)
+
+  assert evaluation_result.summary_metrics is not None
+  assert evaluation_result.summary_metrics
+  for summary in evaluation_result.summary_metrics:
+    assert isinstance(summary, types.AggregatedMetricResult)
+    assert summary.metric_name == "my_custom_code_metric"
+
+  assert evaluation_result.eval_case_results is not None
+  assert evaluation_result.eval_case_results
+  for case_result in evaluation_result.eval_case_results:
+    assert isinstance(case_result, types.EvalCaseResult)
+    assert case_result.eval_case_index is not None
+    assert case_result.response_candidate_results is not None
+
+
+def test_custom_code_execution_with_region(client):
+  """Tests that code_execution_region is included in the custom code execution spec."""
+
+  prompts_df = pd.DataFrame({
+      "prompt": ["What is 2+2?", "What is 3+3?"],
+      "response": ["4", "5"],
+      "reference": ["4", "6"],
+  })
+
+  eval_dataset = types.EvaluationDataset(
+      eval_dataset_df=prompts_df,
+      candidate_name="test_model",
+  )
+
+  metric = types.Metric(
+      name="my_custom_code_metric",
+      remote_custom_function=CODE_SNIPPET,
+      code_execution_region="europe-west3",
+  )
+
+  evaluation_result = client.evals.evaluate(
+      dataset=eval_dataset,
+      metrics=[metric],
+  )
+
+  assert isinstance(evaluation_result, types.EvaluationResult)
+
+  assert evaluation_result.summary_metrics is not None
+  assert evaluation_result.summary_metrics
+  for summary in evaluation_result.summary_metrics:
+    assert isinstance(summary, types.AggregatedMetricResult)
+    assert summary.metric_name == "my_custom_code_metric"
+
+  assert evaluation_result.eval_case_results is not None
+  assert evaluation_result.eval_case_results
+  for case_result in evaluation_result.eval_case_results:
+    assert isinstance(case_result, types.EvalCaseResult)
+    assert case_result.eval_case_index is not None
+    assert case_result.response_candidate_results is not None
 
 
 @pytest.mark.parametrize(