diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py
index 7292d555c5..8e0559744e 100644
--- a/tests/unit/vertexai/genai/test_evals.py
+++ b/tests/unit/vertexai/genai/test_evals.py
@@ -158,14 +158,13 @@ def test_eval_run(self):
@pytest.mark.usefixtures("google_auth_mock")
@mock.patch.object(client.Client, "_get_api_client")
- @mock.patch.object(evals.Evals, "batch_eval")
- def test_eval_batch_eval(self, mock_evaluate, mock_get_api_client):
+ @mock.patch.object(evals.Evals, "batch_evaluate")
+ def test_eval_batch_evaluate(self, mock_evaluate, mock_get_api_client):
test_client = vertexai.Client(project=_TEST_PROJECT, location=_TEST_LOCATION)
- test_client.evals.batch_eval(
+ test_client.evals.batch_evaluate(
dataset=vertexai_genai_types.EvaluationDataset(),
metrics=[vertexai_genai_types.Metric(name="test")],
- output_config=vertexai_genai_types.OutputConfig(),
- autorater_config=vertexai_genai_types.AutoraterConfig(),
+ dest="gs://bucket/output",
config=vertexai_genai_types.EvaluateDatasetConfig(),
)
mock_evaluate.assert_called_once()
@@ -376,13 +375,14 @@ def test_inference_with_gcs_destination(
mock_generate_content_response
)
- gcs_dest_path = "gs://bucket/output.jsonl"
- config = vertexai_genai_types.EvalRunInferenceConfig(dest=gcs_dest_path)
+ gcs_dest_dir = "gs://bucket/output"
+ config = vertexai_genai_types.EvalRunInferenceConfig(dest=gcs_dest_dir)
inference_result = self.client.evals.run_inference(
model="gemini-pro", src=mock_df, config=config
)
+ expected_gcs_path = os.path.join(gcs_dest_dir, "inference_results.jsonl")
expected_df_to_save = pd.DataFrame(
{
"prompt": ["test prompt"],
@@ -393,7 +393,7 @@ def test_inference_with_gcs_destination(
pd.testing.assert_frame_equal(saved_df, expected_df_to_save)
mock_gcs_utils.return_value.upload_dataframe.assert_called_once_with(
df=mock.ANY,
- gcs_destination_blob_path=gcs_dest_path,
+ gcs_destination_blob_path=expected_gcs_path,
file_type="jsonl",
)
pd.testing.assert_frame_equal(
@@ -401,7 +401,7 @@ def test_inference_with_gcs_destination(
)
assert inference_result.candidate_name == "gemini-pro"
assert inference_result.gcs_source == vertexai_genai_types.GcsSource(
- uris=[gcs_dest_path]
+ uris=[expected_gcs_path]
)
@mock.patch.object(_evals_common, "Models")
@@ -434,16 +434,17 @@ def test_inference_with_local_destination(
mock_generate_content_response
)
- local_dest_path = "/tmp/test/output_dir/results.jsonl"
- config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_path)
+ local_dest_dir = "/tmp/test/output_dir"
+ config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_dir)
inference_result = self.client.evals.run_inference(
model="gemini-pro", src=mock_df, config=config
)
- mock_makedirs.assert_called_once_with("/tmp/test/output_dir", exist_ok=True)
+ mock_makedirs.assert_called_once_with(local_dest_dir, exist_ok=True)
+ expected_save_path = os.path.join(local_dest_dir, "inference_results.jsonl")
mock_df_to_json.assert_called_once_with(
- local_dest_path, orient="records", lines=True
+ expected_save_path, orient="records", lines=True
)
expected_df = pd.DataFrame(
{
@@ -457,7 +458,7 @@ def test_inference_with_local_destination(
@mock.patch.object(_evals_common, "Models")
@mock.patch.object(_evals_utils, "EvalDatasetLoader")
- def test_inference_from_request_column_save_locally(
+ def test_inference_from_request_column_save_to_local_dir(
self, mock_eval_dataset_loader, mock_models
):
mock_df = pd.DataFrame(
@@ -494,8 +495,8 @@ def test_inference_from_request_column_save_locally(
mock_generate_content_responses
)
- local_dest_path = "/tmp/output.jsonl"
- config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_path)
+ local_dest_dir = "/tmp/test_output_dir"
+ config = vertexai_genai_types.EvalRunInferenceConfig(dest=local_dest_dir)
inference_result = self.client.evals.run_inference(
model="gemini-pro", src=mock_df, config=config
@@ -530,13 +531,15 @@ def test_inference_from_request_column_save_locally(
expected_df.sort_values(by="request").reset_index(drop=True),
)
- with open(local_dest_path, "r") as f:
+ saved_file_path = os.path.join(local_dest_dir, "inference_results.jsonl")
+ with open(saved_file_path, "r") as f:
saved_records = [json.loads(line) for line in f]
expected_records = expected_df.to_dict(orient="records")
assert sorted(saved_records, key=lambda x: x["request"]) == sorted(
expected_records, key=lambda x: x["request"]
)
- os.remove(local_dest_path)
+ os.remove(saved_file_path)
+ os.rmdir(local_dest_dir)
assert inference_result.candidate_name == "gemini-pro"
assert inference_result.gcs_source is None
diff --git a/vertexai/_genai/_evals_common.py b/vertexai/_genai/_evals_common.py
index a99649bc4b..7a86e8003d 100644
--- a/vertexai/_genai/_evals_common.py
+++ b/vertexai/_genai/_evals_common.py
@@ -617,19 +617,13 @@ def _execute_inference(
if dest:
file_name = "inference_results.jsonl"
- full_dest_path = dest
is_gcs_path = dest.startswith(_evals_utils.GCS_PREFIX)
if is_gcs_path:
- if not dest.endswith("/"):
- pass
- else:
- full_dest_path = os.path.join(dest, file_name)
+ full_dest_path = os.path.join(dest, file_name)
else:
- if os.path.isdir(dest):
- full_dest_path = os.path.join(dest, file_name)
-
- os.makedirs(os.path.dirname(full_dest_path), exist_ok=True)
+ os.makedirs(dest, exist_ok=True)
+ full_dest_path = os.path.join(dest, file_name)
logger.info("Saving inference results to: %s", full_dest_path)
try:
diff --git a/vertexai/_genai/_evals_utils.py b/vertexai/_genai/_evals_utils.py
index 7f07f1f108..9cb7444eed 100644
--- a/vertexai/_genai/_evals_utils.py
+++ b/vertexai/_genai/_evals_utils.py
@@ -13,22 +13,26 @@
# limitations under the License.
#
"""Utility functions for evals."""
+
import io
import json
import logging
import os
import re
import time
-from typing import Any, Union, Optional
+from typing import Any, Optional, Union
from google.cloud import bigquery
from google.cloud import storage
from google.genai._api_client import BaseApiClient
+from google.genai._common import get_value_by_path as getv
+from google.genai._common import set_value_by_path as setv
import pandas as pd
import yaml
from . import types
+
logger = logging.getLogger(__name__)
@@ -525,3 +529,261 @@ def MULTI_TURN_SAFETY(self) -> LazyLoadedPrebuiltMetric:
PrebuiltMetric = PrebuiltMetricLoader()
+
+
+class BatchEvaluateRequestPreparer:
+ """Prepares data for requests."""
+
+ @staticmethod
+ def _EvaluationDataset_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, ["gcs_source"]) is not None:
+ setv(
+ to_object,
+ ["gcs_source"],
+ getv(from_object, ["gcs_source"]),
+ )
+
+ if getv(from_object, ["bigquery_source"]) is not None:
+ setv(
+ to_object,
+ ["bigquery_source"],
+ getv(from_object, ["bigquery_source"]),
+ )
+
+ return to_object
+
+ @staticmethod
+ def _Metric_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, ["prompt_template"]) is not None:
+ setv(
+ to_object,
+ ["pointwise_metric_spec", "prompt_template"],
+ getv(from_object, ["prompt_template"]),
+ )
+
+ if getv(from_object, ["judge_model"]) is not None:
+ setv(
+ parent_object,
+ ["autorater_config", "autorater_model"],
+ getv(from_object, ["judge_model"]),
+ )
+
+ if getv(from_object, ["judge_model_sampling_count"]) is not None:
+ setv(
+ parent_object,
+ ["autorater_config", "sampling_count"],
+ getv(from_object, ["judge_model_sampling_count"]),
+ )
+
+ if getv(from_object, ["judge_model_system_instruction"]) is not None:
+ setv(
+ to_object,
+ ["pointwise_metric_spec", "system_instruction"],
+ getv(from_object, ["judge_model_system_instruction"]),
+ )
+
+ if getv(from_object, ["return_raw_output"]) is not None:
+ setv(
+ to_object,
+ [
+ "pointwise_metric_spec",
+ "custom_output_format_config",
+ "return_raw_output",
+ ],
+ getv(from_object, ["return_raw_output"]),
+ )
+
+ return to_object
+
+ @staticmethod
+ def _OutputConfig_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, ["gcs_destination"]) is not None:
+ setv(
+ to_object,
+ ["gcsDestination"],
+ getv(from_object, ["gcs_destination"]),
+ )
+
+ return to_object
+
+ @staticmethod
+ def _EvaluationDataset_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, ["dataset", "gcs_source"]) is not None:
+ setv(
+ to_object,
+ ["gcs_source"],
+ getv(from_object, ["dataset", "gcs_source"]),
+ )
+
+ if getv(from_object, ["dataset", "bigquery_source"]) is not None:
+ setv(
+ to_object,
+ ["bigquery_source"],
+ getv(from_object, ["dataset", "bigquery_source"]),
+ )
+
+ return to_object
+
+ @staticmethod
+ def _AutoraterConfig_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, ["sampling_count"]) is not None:
+ setv(to_object, ["samplingCount"], getv(from_object, ["sampling_count"]))
+
+ if getv(from_object, ["flip_enabled"]) is not None:
+ setv(to_object, ["flipEnabled"], getv(from_object, ["flip_enabled"]))
+
+ if getv(from_object, ["autorater_model"]) is not None:
+ setv(
+ to_object,
+ ["autoraterModel"],
+ getv(from_object, ["autorater_model"]),
+ )
+
+ return to_object
+
+ @staticmethod
+ def EvaluateDatasetOperation_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, ["name"]) is not None:
+ setv(to_object, ["name"], getv(from_object, ["name"]))
+
+ if getv(from_object, ["metadata"]) is not None:
+ setv(to_object, ["metadata"], getv(from_object, ["metadata"]))
+
+ if getv(from_object, ["done"]) is not None:
+ setv(to_object, ["done"], getv(from_object, ["done"]))
+
+ if getv(from_object, ["error"]) is not None:
+ setv(to_object, ["error"], getv(from_object, ["error"]))
+
+ if getv(from_object, ["response"]) is not None:
+ setv(
+ to_object,
+ ["response"],
+ BatchEvaluateRequestPreparer._EvaluationDataset_from_vertex(
+ getv(from_object, ["response"]), to_object
+ ),
+ )
+
+ return to_object
+
+ @staticmethod
+ def EvaluateDatasetRequestParameters_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, ["dataset"]) is not None:
+ setv(
+ to_object,
+ ["dataset"],
+ BatchEvaluateRequestPreparer._EvaluationDataset_to_vertex(
+ getv(from_object, ["dataset"]), to_object
+ ),
+ )
+
+ if getv(from_object, ["metrics"]) is not None:
+ setv(
+ to_object,
+ ["metrics"],
+ [
+ BatchEvaluateRequestPreparer._Metric_to_vertex(item, to_object)
+ for item in getv(from_object, ["metrics"])
+ ],
+ )
+
+ if getv(from_object, ["output_config"]) is not None:
+ setv(
+ to_object,
+ ["outputConfig"],
+ BatchEvaluateRequestPreparer._OutputConfig_to_vertex(
+ getv(from_object, ["output_config"]), to_object
+ ),
+ )
+
+ if getv(from_object, ["autorater_config"]) is not None:
+ setv(
+ to_object,
+ ["autoraterConfig"],
+ BatchEvaluateRequestPreparer._AutoraterConfig_to_vertex(
+ getv(from_object, ["autorater_config"]), to_object
+ ),
+ )
+
+ if getv(from_object, ["config"]) is not None:
+ setv(to_object, ["config"], getv(from_object, ["config"]))
+
+ return to_object
+
+ @staticmethod
+ def prepare_metric_payload(
+ request_dict: dict[str, Any], resolved_metrics: list[types.MetricSubclass]
+ ) -> dict[str, Any]:
+ """Prepares the metric payload for the evaluation request.
+
+ Args:
+ request_dict: The dictionary containing the request details.
+ resolved_metrics: A list of resolved metric objects.
+
+ Returns:
+ The updated request dictionary with the prepared metric payload.
+ """
+ metrics_payload = []
+
+ for metric in resolved_metrics:
+ metric_payload_item = {}
+ metric_payload_item["aggregation_metrics"] = [
+ "AVERAGE",
+ "STANDARD_DEVIATION",
+ ]
+
+ metric_name = getv(metric, ["name"]).lower()
+
+ if metric_name == "exact_match":
+ metric_payload_item["exact_match_spec"] = {}
+ elif metric_name == "bleu":
+ metric_payload_item["bleu_spec"] = {}
+ elif metric_name.startswith("rouge"):
+ rouge_type = metric_name.replace("_", "")
+ metric_payload_item["rouge_spec"] = {"rouge_type": rouge_type}
+
+ elif hasattr(metric, "prompt_template") and metric.prompt_template:
+ pointwise_spec = {"metric_prompt_template": metric.prompt_template}
+ system_instruction = getv(metric, ["judge_model_system_instruction"])
+ if system_instruction:
+ pointwise_spec["system_instruction"] = system_instruction
+ metric_payload_item["pointwise_metric_spec"] = pointwise_spec
+ else:
+ raise ValueError(
+ "Unsupported metric type or invalid metric name:" f" {metric_name}"
+ )
+
+ metrics_payload.append(metric_payload_item)
+ request_dict["metrics"] = metrics_payload
+ return request_dict
diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py
index 293b64814a..e14742cb3d 100644
--- a/vertexai/_genai/evals.py
+++ b/vertexai/_genai/evals.py
@@ -28,6 +28,7 @@
import pandas as pd
from . import _evals_common
+from . import _evals_utils
from . import types
@@ -550,6 +551,12 @@ def _AutoraterConfig_to_vertex(
parent_object: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
to_object: dict[str, Any] = {}
+ if getv(from_object, ["sampling_count"]) is not None:
+ setv(to_object, ["samplingCount"], getv(from_object, ["sampling_count"]))
+
+ if getv(from_object, ["flip_enabled"]) is not None:
+ setv(to_object, ["flipEnabled"], getv(from_object, ["flip_enabled"]))
+
if getv(from_object, ["autorater_model"]) is not None:
setv(
to_object,
@@ -557,12 +564,6 @@ def _AutoraterConfig_to_vertex(
getv(from_object, ["autorater_model"]),
)
- if getv(from_object, ["flip_enabled"]) is not None:
- setv(to_object, ["flipEnabled"], getv(from_object, ["flip_enabled"]))
-
- if getv(from_object, ["sampling_count"]) is not None:
- setv(to_object, ["samplingCount"], getv(from_object, ["sampling_count"]))
-
return to_object
@@ -663,136 +664,6 @@ def _EvaluateInstancesRequestParameters_to_vertex(
return to_object
-def _EvaluationDataset_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, ["gcs_source"]) is not None:
- setv(
- to_object,
- ["dataset", "gcs_source"],
- getv(from_object, ["gcs_source"]),
- )
-
- if getv(from_object, ["bigquery_source"]) is not None:
- setv(
- to_object,
- ["dataset", "bigquery_source"],
- getv(from_object, ["bigquery_source"]),
- )
-
- return to_object
-
-
-def _Metric_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, ["prompt_template"]) is not None:
- setv(
- to_object,
- ["pointwise_metric_spec", "prompt_template"],
- getv(from_object, ["prompt_template"]),
- )
-
- if getv(from_object, ["judge_model"]) is not None:
- setv(
- parent_object,
- ["autorater_config", "autorater_model"],
- getv(from_object, ["judge_model"]),
- )
-
- if getv(from_object, ["judge_model_sampling_count"]) is not None:
- setv(
- parent_object,
- ["autorater_config", "sampling_count"],
- getv(from_object, ["judge_model_sampling_count"]),
- )
-
- if getv(from_object, ["judge_model_system_instruction"]) is not None:
- setv(
- to_object,
- ["pointwise_metric_spec", "system_instruction"],
- getv(from_object, ["judge_model_system_instruction"]),
- )
-
- if getv(from_object, ["return_raw_output"]) is not None:
- setv(
- to_object,
- [
- "pointwise_metric_spec",
- "custom_output_format_config",
- "return_raw_output",
- ],
- getv(from_object, ["return_raw_output"]),
- )
-
- return to_object
-
-
-def _OutputConfig_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, ["gcs_destination"]) is not None:
- setv(
- to_object,
- ["gcsDestination"],
- getv(from_object, ["gcs_destination"]),
- )
-
- return to_object
-
-
-def _EvaluateDatasetRequestParameters_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, ["dataset"]) is not None:
- setv(
- to_object,
- ["dataset"],
- _EvaluationDataset_to_vertex(getv(from_object, ["dataset"]), to_object),
- )
-
- if getv(from_object, ["metrics"]) is not None:
- setv(
- to_object,
- ["metrics"],
- [
- _Metric_to_vertex(item, to_object)
- for item in getv(from_object, ["metrics"])
- ],
- )
-
- if getv(from_object, ["output_config"]) is not None:
- setv(
- to_object,
- ["outputConfig"],
- _OutputConfig_to_vertex(getv(from_object, ["output_config"]), to_object),
- )
-
- if getv(from_object, ["autorater_config"]) is not None:
- setv(
- to_object,
- ["autoraterConfig"],
- _AutoraterConfig_to_vertex(
- getv(from_object, ["autorater_config"]), to_object
- ),
- )
-
- if getv(from_object, ["config"]) is not None:
- setv(to_object, ["config"], getv(from_object, ["config"]))
-
- return to_object
-
-
def _EvaluateInstancesResponse_from_vertex(
from_object: Union[dict[str, Any], object],
parent_object: Optional[dict[str, Any]] = None,
@@ -919,56 +790,6 @@ def _EvaluateInstancesResponse_from_vertex(
return to_object
-def _EvaluationDataset_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, ["dataset", "gcs_source"]) is not None:
- setv(
- to_object,
- ["gcs_source"],
- getv(from_object, ["dataset", "gcs_source"]),
- )
-
- if getv(from_object, ["dataset", "bigquery_source"]) is not None:
- setv(
- to_object,
- ["bigquery_source"],
- getv(from_object, ["dataset", "bigquery_source"]),
- )
-
- return to_object
-
-
-def _EvaluateDatasetOperation_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, ["name"]) is not None:
- setv(to_object, ["name"], getv(from_object, ["name"]))
-
- if getv(from_object, ["metadata"]) is not None:
- setv(to_object, ["metadata"], getv(from_object, ["metadata"]))
-
- if getv(from_object, ["done"]) is not None:
- setv(to_object, ["done"], getv(from_object, ["done"]))
-
- if getv(from_object, ["error"]) is not None:
- setv(to_object, ["error"], getv(from_object, ["error"]))
-
- if getv(from_object, ["response"]) is not None:
- setv(
- to_object,
- ["response"],
- _EvaluationDataset_from_vertex(getv(from_object, ["response"]), to_object),
- )
-
- return to_object
-
-
class Evals(_api_module.BaseModule):
def _evaluate_instances(
self,
@@ -1044,68 +865,8 @@ def _evaluate_instances(
return_value = types.EvaluateInstancesResponse._from_response(
response=response_dict, kwargs=parameter_model.model_dump()
)
- self._api_client._verify_response(return_value)
-
- return return_value
-
- def batch_eval(
- self,
- *,
- dataset: types.EvaluationDatasetOrDict,
- metrics: list[types.MetricOrDict],
- output_config: types.OutputConfigOrDict,
- autorater_config: Optional[types.AutoraterConfigOrDict] = None,
- config: Optional[types.EvaluateDatasetConfigOrDict] = None,
- ) -> types.EvaluateDatasetOperation:
- """Evaluates a dataset based on a set of given metrics."""
-
- parameter_model = types._EvaluateDatasetRequestParameters(
- dataset=dataset,
- metrics=metrics,
- output_config=output_config,
- autorater_config=autorater_config,
- config=config,
- )
-
- request_url_dict: Optional[dict[str, str]]
- if not self._api_client.vertexai:
- raise ValueError("This method is only supported in the Vertex AI client.")
- else:
- request_dict = _EvaluateDatasetRequestParameters_to_vertex(parameter_model)
- request_url_dict = request_dict.get("_url")
- if request_url_dict:
- path = ":evaluateDataset".format_map(request_url_dict)
- else:
- path = ":evaluateDataset"
- query_params = request_dict.get("_query")
- if query_params:
- path = f"{path}?{urlencode(query_params)}"
- # TODO: remove the hack that pops config.
- request_dict.pop("config", None)
-
- http_options: Optional[genai_types.HttpOptions] = None
- if (
- parameter_model.config is not None
- and parameter_model.config.http_options is not None
- ):
- http_options = parameter_model.config.http_options
-
- request_dict = _common.convert_to_dict(request_dict)
- request_dict = _common.encode_unserializable_types(request_dict)
-
- response = self._api_client.request("post", path, request_dict, http_options)
-
- response_dict = "" if not response.body else json.loads(response.body)
-
- if self._api_client.vertexai:
- response_dict = _EvaluateDatasetOperation_from_vertex(response_dict)
-
- return_value = types.EvaluateDatasetOperation._from_response(
- response=response_dict, kwargs=parameter_model.model_dump()
- )
self._api_client._verify_response(return_value)
-
return return_value
def run(self) -> types.EvaluateInstancesResponse:
@@ -1203,9 +964,11 @@ def evaluate(
config = types.EvaluateMethodConfig.model_validate(config)
if isinstance(dataset, list):
dataset = [
- types.EvaluationDataset.model_validate(ds_item)
- if isinstance(ds_item, dict)
- else ds_item
+ (
+ types.EvaluationDataset.model_validate(ds_item)
+ if isinstance(ds_item, dict)
+ else ds_item
+ )
for ds_item in dataset
]
else:
@@ -1220,6 +983,76 @@ def evaluate(
dest=config.dest,
)
+ def batch_evaluate(
+ self,
+ *,
+ dataset: types.EvaluationDatasetOrDict,
+ metrics: list[types.MetricOrDict],
+ dest: str,
+ config: Optional[types.EvaluateDatasetConfigOrDict] = None,
+ ) -> types.EvaluateDatasetOperation:
+ """Evaluates a dataset based on a set of given metrics."""
+
+ resolved_metrics = _evals_common._resolve_metrics(metrics, self._api_client)
+ output_config = types.OutputConfig(
+ gcs_destination=types.GcsDestination(output_uri_prefix=dest)
+ )
+ parameter_model = types.EvaluateDatasetRequestParameters(
+ dataset=dataset,
+ metrics=resolved_metrics,
+ output_config=output_config,
+ config=config,
+ )
+
+ request_url_dict: Optional[dict[str, str]]
+ if not self._api_client.vertexai:
+ raise ValueError("This method is only supported in the Vertex AI client.")
+ else:
+ request_dict = _evals_utils.BatchEvaluateRequestPreparer.EvaluateDatasetRequestParameters_to_vertex(
+ parameter_model
+ )
+ request_url_dict = request_dict.get("_url")
+ if request_url_dict:
+ path = ":evaluateDataset".format_map(request_url_dict)
+ else:
+ path = ":evaluateDataset"
+
+ request_dict = _evals_utils.BatchEvaluateRequestPreparer.prepare_metric_payload(
+ request_dict, resolved_metrics
+ )
+
+ query_params = request_dict.get("_query")
+ if query_params:
+ path = f"{path}?{urlencode(query_params)}"
+ # TODO: remove the hack that pops config.
+ request_dict.pop("config", None)
+
+ http_options: Optional[genai_types.HttpOptions] = None
+ if (
+ parameter_model.config is not None
+ and parameter_model.config.http_options is not None
+ ):
+ http_options = parameter_model.config.http_options
+
+ request_dict = _common.convert_to_dict(request_dict)
+ request_dict = _common.encode_unserializable_types(request_dict)
+
+ response = self._api_client.request("post", path, request_dict, http_options)
+
+ response_dict = "" if not response.body else json.loads(response.body)
+
+ if self._api_client.vertexai:
+ response_dict = _evals_utils.BatchEvaluateRequestPreparer.EvaluateDatasetOperation_from_vertex(
+ response_dict
+ )
+
+ return_value = types.EvaluateDatasetOperation._from_response(
+ response=response_dict, kwargs=parameter_model.model_dump()
+ )
+ self._api_client._verify_response(return_value)
+
+ return return_value
+
class AsyncEvals(_api_module.BaseModule):
async def _evaluate_instances(
@@ -1298,26 +1131,27 @@ async def _evaluate_instances(
return_value = types.EvaluateInstancesResponse._from_response(
response=response_dict, kwargs=parameter_model.model_dump()
)
- self._api_client._verify_response(return_value)
+ self._api_client._verify_response(return_value)
return return_value
- async def batch_eval(
+ async def batch_evaluate(
self,
*,
dataset: types.EvaluationDatasetOrDict,
metrics: list[types.MetricOrDict],
- output_config: types.OutputConfigOrDict,
- autorater_config: Optional[types.AutoraterConfigOrDict] = None,
+ dest: str,
config: Optional[types.EvaluateDatasetConfigOrDict] = None,
) -> types.EvaluateDatasetOperation:
"""Evaluates a dataset based on a set of given metrics."""
-
- parameter_model = types._EvaluateDatasetRequestParameters(
+ resolved_metrics = _evals_common._resolve_metrics(metrics, self._api_client)
+ output_config = types.OutputConfig(
+ gcs_destination=types.GcsDestination(output_uri_prefix=dest)
+ )
+ parameter_model = types.EvaluateDatasetRequestParameters(
dataset=dataset,
- metrics=metrics,
+ metrics=resolved_metrics,
output_config=output_config,
- autorater_config=autorater_config,
config=config,
)
@@ -1325,13 +1159,19 @@ async def batch_eval(
if not self._api_client.vertexai:
raise ValueError("This method is only supported in the Vertex AI client.")
else:
- request_dict = _EvaluateDatasetRequestParameters_to_vertex(parameter_model)
+ request_dict = _evals_utils.BatchEvaluateRequestPreparer.EvaluateDatasetRequestParameters_to_vertex(
+ parameter_model
+ )
request_url_dict = request_dict.get("_url")
if request_url_dict:
path = ":evaluateDataset".format_map(request_url_dict)
else:
path = ":evaluateDataset"
+ request_dict = _evals_utils.BatchEvaluateRequestPreparer.prepare_metric_payload(
+ request_dict, resolved_metrics
+ )
+
query_params = request_dict.get("_query")
if query_params:
path = f"{path}?{urlencode(query_params)}"
@@ -1355,7 +1195,9 @@ async def batch_eval(
response_dict = "" if not response.body else json.loads(response.body)
if self._api_client.vertexai:
- response_dict = _EvaluateDatasetOperation_from_vertex(response_dict)
+ response_dict = _evals_utils.BatchEvaluateRequestPreparer.EvaluateDatasetOperation_from_vertex(
+ response_dict
+ )
return_value = types.EvaluateDatasetOperation._from_response(
response=response_dict, kwargs=parameter_model.model_dump()
diff --git a/vertexai/_genai/types.py b/vertexai/_genai/types.py
index f61e2f973e..818d36d2d0 100644
--- a/vertexai/_genai/types.py
+++ b/vertexai/_genai/types.py
@@ -2043,1473 +2043,1470 @@ class EvaluateInstancesResponseDict(TypedDict, total=False):
]
-class Message(_common.BaseModel):
- """Represents a single message turn in a conversation."""
+class OptimizeConfig(_common.BaseModel):
+ """Config for Prompt Optimizer."""
- turn_id: Optional[str] = Field(
- default=None, description="""Unique identifier for the message turn."""
- )
- content: Optional[genai_types.Content] = Field(
- default=None,
- description="""Content of the message, including function call.""",
- )
- creation_timestamp: Optional[datetime.datetime] = Field(
- default=None,
- description="""Timestamp indicating when the message was created.""",
- )
- author: Optional[str] = Field(
- default=None,
- description="""Name of the entity that produced the message.""",
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
-class MessageDict(TypedDict, total=False):
- """Represents a single message turn in a conversation."""
+class OptimizeConfigDict(TypedDict, total=False):
+ """Config for Prompt Optimizer."""
- turn_id: Optional[str]
- """Unique identifier for the message turn."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
- content: Optional[genai_types.Content]
- """Content of the message, including function call."""
- creation_timestamp: Optional[datetime.datetime]
- """Timestamp indicating when the message was created."""
+OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict]
- author: Optional[str]
- """Name of the entity that produced the message."""
+class _OptimizeRequestParameters(_common.BaseModel):
+ """Parameters for the optimize_prompt method."""
-MessageOrDict = Union[Message, MessageDict]
+ config: Optional[OptimizeConfig] = Field(default=None, description="""""")
-class AgentData(_common.BaseModel):
- """Container for all agent-specific data."""
+class _OptimizeRequestParametersDict(TypedDict, total=False):
+ """Parameters for the optimize_prompt method."""
- tool_use_trajectory: Optional[list[Message]] = Field(
- default=None,
- description="""Tool use trajectory in chronological order.""",
- )
- intermediate_responses: Optional[list[Message]] = Field(
- default=None,
- description="""Intermediate responses generated by sub-agents to convey progress or status in a multi-agent system, distinct from the final response.""",
- )
+ config: Optional[OptimizeConfigDict]
+ """"""
-class AgentDataDict(TypedDict, total=False):
- """Container for all agent-specific data."""
+_OptimizeRequestParametersOrDict = Union[
+ _OptimizeRequestParameters, _OptimizeRequestParametersDict
+]
- tool_use_trajectory: Optional[list[MessageDict]]
- """Tool use trajectory in chronological order."""
- intermediate_responses: Optional[list[MessageDict]]
- """Intermediate responses generated by sub-agents to convey progress or status in a multi-agent system, distinct from the final response."""
+class OptimizeResponse(_common.BaseModel):
+ """Response for the optimize_prompt method."""
+ pass
-AgentDataOrDict = Union[AgentData, AgentDataDict]
+class OptimizeResponseDict(TypedDict, total=False):
+ """Response for the optimize_prompt method."""
-class ResponseCandidate(_common.BaseModel):
- """A model-generated content to the prompt."""
+ pass
- response: Optional[genai_types.Content] = Field(
+
+OptimizeResponseOrDict = Union[OptimizeResponse, OptimizeResponseDict]
+
+
+class GcsDestination(_common.BaseModel):
+ """The Google Cloud Storage location where the output is to be written to."""
+
+ output_uri_prefix: Optional[str] = Field(
default=None,
- description="""The final model-generated response to the `prompt`.""",
+ description="""Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist.""",
)
- agent_data: Optional[AgentData] = Field(
+
+
+class GcsDestinationDict(TypedDict, total=False):
+ """The Google Cloud Storage location where the output is to be written to."""
+
+ output_uri_prefix: Optional[str]
+ """Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist."""
+
+
+GcsDestinationOrDict = Union[GcsDestination, GcsDestinationDict]
+
+
+class PscInterfaceConfig(_common.BaseModel):
+ """Configuration for PSC-I."""
+
+ network_attachment: Optional[str] = Field(
default=None,
- description="""Agent-specific data including tool use trajectory and intermediate responses for more complex interactions.""",
+ description="""Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""",
)
-class ResponseCandidateDict(TypedDict, total=False):
- """A model-generated content to the prompt."""
-
- response: Optional[genai_types.Content]
- """The final model-generated response to the `prompt`."""
+class PscInterfaceConfigDict(TypedDict, total=False):
+ """Configuration for PSC-I."""
- agent_data: Optional[AgentDataDict]
- """Agent-specific data including tool use trajectory and intermediate responses for more complex interactions."""
+ network_attachment: Optional[str]
+ """Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I."""
-ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict]
+PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict]
-class EvalCase(_common.BaseModel):
- """A comprehensive representation of a GenAI interaction for evaluation."""
+class Scheduling(_common.BaseModel):
+ """All parameters related to queuing and scheduling of custom jobs."""
- prompt: Optional[genai_types.Content] = Field(
+ disable_retries: Optional[bool] = Field(
default=None,
- description="""The most recent user message (current input).""",
+ description="""Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""",
)
- responses: Optional[list[ResponseCandidate]] = Field(
+ max_wait_duration: Optional[str] = Field(
default=None,
- description="""Model-generated replies to the last user message. Multiple responses are allowed to support use cases such as comparing different model outputs.""",
+ description="""Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""",
)
- reference: Optional[ResponseCandidate] = Field(
+ restart_job_on_worker_restart: Optional[bool] = Field(
default=None,
- description="""User-provided, golden reference model reply to prompt in context of chat history.""",
- )
- system_instruction: Optional[genai_types.Content] = Field(
- default=None, description="""System instruction for the model."""
+ description="""Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""",
)
- conversation_history: Optional[list[Message]] = Field(
+ strategy: Optional[Strategy] = Field(
default=None,
- description="""List of all prior messages in the conversation (chat history).""",
+ description="""Optional. This determines which type of scheduling strategy to use.""",
)
- eval_case_id: Optional[str] = Field(
+ timeout: Optional[str] = Field(
default=None,
- description="""Unique identifier for the evaluation case.""",
+ description="""Optional. The maximum job running time. The default is 7 days.""",
)
- # Allow extra fields to support custom metric prompts and stay backward compatible.
- model_config = ConfigDict(frozen=True, extra="allow")
-
-class EvalCaseDict(TypedDict, total=False):
- """A comprehensive representation of a GenAI interaction for evaluation."""
- prompt: Optional[genai_types.Content]
- """The most recent user message (current input)."""
+class SchedulingDict(TypedDict, total=False):
+ """All parameters related to queuing and scheduling of custom jobs."""
- responses: Optional[list[ResponseCandidateDict]]
- """Model-generated replies to the last user message. Multiple responses are allowed to support use cases such as comparing different model outputs."""
+ disable_retries: Optional[bool]
+ """Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false."""
- reference: Optional[ResponseCandidateDict]
- """User-provided, golden reference model reply to prompt in context of chat history."""
+ max_wait_duration: Optional[str]
+ """Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours."""
- system_instruction: Optional[genai_types.Content]
- """System instruction for the model."""
+ restart_job_on_worker_restart: Optional[bool]
+ """Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job."""
- conversation_history: Optional[list[MessageDict]]
- """List of all prior messages in the conversation (chat history)."""
+ strategy: Optional[Strategy]
+ """Optional. This determines which type of scheduling strategy to use."""
- eval_case_id: Optional[str]
- """Unique identifier for the evaluation case."""
+ timeout: Optional[str]
+ """Optional. The maximum job running time. The default is 7 days."""
-EvalCaseOrDict = Union[EvalCase, EvalCaseDict]
+SchedulingOrDict = Union[Scheduling, SchedulingDict]
-class GcsSource(_common.BaseModel):
- """The Google Cloud Storage location for the input content."""
+class EnvVar(_common.BaseModel):
+ """Represents an environment variable present in a Container or Python Module."""
- uris: Optional[list[str]] = Field(
+ name: Optional[str] = Field(
default=None,
- description="""Required. Google Cloud Storage URI(-s) to the input file(s). May contain wildcards. For more information on wildcards, see https://cloud.google.com/storage/docs/wildcards.""",
+ description="""Required. Name of the environment variable. Must be a valid C identifier.""",
)
-
-
-class GcsSourceDict(TypedDict, total=False):
- """The Google Cloud Storage location for the input content."""
-
- uris: Optional[list[str]]
- """Required. Google Cloud Storage URI(-s) to the input file(s). May contain wildcards. For more information on wildcards, see https://cloud.google.com/storage/docs/wildcards."""
-
-
-GcsSourceOrDict = Union[GcsSource, GcsSourceDict]
-
-
-class BigQuerySource(_common.BaseModel):
- """The BigQuery location for the input content."""
-
- input_uri: Optional[str] = Field(
+ value: Optional[str] = Field(
default=None,
- description="""Required. BigQuery URI to a table, up to 2000 characters long. Accepted forms: * BigQuery path. For example: `bq://projectId.bqDatasetId.bqTableId`.""",
+ description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""",
)
-class BigQuerySourceDict(TypedDict, total=False):
- """The BigQuery location for the input content."""
+class EnvVarDict(TypedDict, total=False):
+ """Represents an environment variable present in a Container or Python Module."""
- input_uri: Optional[str]
- """Required. BigQuery URI to a table, up to 2000 characters long. Accepted forms: * BigQuery path. For example: `bq://projectId.bqDatasetId.bqTableId`."""
+ name: Optional[str]
+ """Required. Name of the environment variable. Must be a valid C identifier."""
+ value: Optional[str]
+ """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not."""
-BigQuerySourceOrDict = Union[BigQuerySource, BigQuerySourceDict]
+EnvVarOrDict = Union[EnvVar, EnvVarDict]
-class EvaluationDataset(_common.BaseModel):
- """The dataset used for evaluation."""
- eval_cases: Optional[list[EvalCase]] = Field(
- default=None, description="""The evaluation cases to be evaluated."""
- )
- eval_dataset_df: Optional["pd.DataFrame"] = Field(
+class ContainerSpec(_common.BaseModel):
+ """The spec of a Container."""
+
+ args: Optional[list[str]] = Field(
default=None,
- description="""The evaluation dataset in the form of a Pandas DataFrame.""",
+ description="""The arguments to be passed when starting the container.""",
)
- candidate_name: Optional[str] = Field(
+ command: Optional[list[str]] = Field(
default=None,
- description="""The name of the candidate model or agent for this evaluation dataset.""",
+ description="""The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""",
)
- gcs_source: Optional[GcsSource] = Field(
+ env: Optional[list[EnvVar]] = Field(
default=None,
- description="""The GCS source for the evaluation dataset.""",
+ description="""Environment variables to be passed to the container. Maximum limit is 100.""",
)
- bigquery_source: Optional[BigQuerySource] = Field(
+ image_uri: Optional[str] = Field(
default=None,
- description="""The BigQuery source for the evaluation dataset.""",
+ description="""Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""",
)
- def show(self) -> None:
- """Shows the evaluation dataset."""
- from . import _evals_visualization
- _evals_visualization.display_evaluation_dataset(self)
+class ContainerSpecDict(TypedDict, total=False):
+ """The spec of a Container."""
+ args: Optional[list[str]]
+ """The arguments to be passed when starting the container."""
-class EvaluationDatasetDict(TypedDict, total=False):
- """The dataset used for evaluation."""
+ command: Optional[list[str]]
+ """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided."""
- eval_cases: Optional[list[EvalCaseDict]]
- """The evaluation cases to be evaluated."""
+ env: Optional[list[EnvVarDict]]
+ """Environment variables to be passed to the container. Maximum limit is 100."""
- eval_dataset_df: Optional["pd.DataFrame"]
- """The evaluation dataset in the form of a Pandas DataFrame."""
+ image_uri: Optional[str]
+ """Required. The URI of a container image in the Container Registry that is to be run on each worker replica."""
- candidate_name: Optional[str]
- """The name of the candidate model or agent for this evaluation dataset."""
- gcs_source: Optional[GcsSourceDict]
- """The GCS source for the evaluation dataset."""
-
- bigquery_source: Optional[BigQuerySourceDict]
- """The BigQuery source for the evaluation dataset."""
-
-
-EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict]
+ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict]
-class Metric(_common.BaseModel):
- """The metric used for evaluation."""
+class DiskSpec(_common.BaseModel):
+ """Represents the spec of disk options."""
- name: Optional[str] = Field(default=None, description="""The name of the metric.""")
- custom_function: Optional[Callable] = Field(
- default=None,
- description="""The custom function that defines the end-to-end logic for metric computation.""",
- )
- prompt_template: Optional[str] = Field(
- default=None, description="""The prompt template for the metric."""
- )
- judge_model: Optional[str] = Field(
- default=None, description="""The judge model for the metric."""
- )
- judge_model_sampling_count: Optional[int] = Field(
- default=None, description="""The sampling count for the judge model."""
- )
- judge_model_system_instruction: Optional[str] = Field(
- default=None,
- description="""The system instruction for the judge model.""",
- )
- return_raw_output: Optional[bool] = Field(
- default=None,
- description="""Whether to return the raw output from the judge model.""",
- )
- parse_and_reduce_fn: Optional[Callable] = Field(
+ boot_disk_size_gb: Optional[int] = Field(
default=None,
- description="""The parse and reduce function for the judge model.""",
+ description="""Size in GB of the boot disk (default is 100GB).""",
)
- aggregate_summary_fn: Optional[Callable] = Field(
+ boot_disk_type: Optional[str] = Field(
default=None,
- description="""The aggregate summary function for the judge model.""",
+ description="""Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""",
)
- # Allow extra fields to support metric-specific config fields.
- model_config = ConfigDict(extra="allow")
-
- _is_predefined: bool = PrivateAttr(default=False)
- """A boolean indicating whether the metric is predefined."""
-
- _config_source: Optional[str] = PrivateAttr(default=None)
- """An optional string indicating the source of the metric configuration."""
- _version: Optional[str] = PrivateAttr(default=None)
- """An optional string indicating the version of the metric."""
+class DiskSpecDict(TypedDict, total=False):
+ """Represents the spec of disk options."""
- @model_validator(mode="after")
- @classmethod
- def validate_name(cls, model: "Metric") -> "Metric":
- if not model.name:
- raise ValueError("Metric name cannot be empty.")
- model.name = model.name.lower()
- return model
+ boot_disk_size_gb: Optional[int]
+ """Size in GB of the boot disk (default is 100GB)."""
- def to_yaml_file(self, file_path: str, version: Optional[str] = None) -> None:
- """Dumps the metric object to a YAML file.
+ boot_disk_type: Optional[str]
+ """Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced"."""
- Args:
- file_path: The path to the YAML file.
- version: Optional version string to include in the YAML output.
- Raises:
- ImportError: If the pyyaml library is not installed.
- """
- if yaml is None:
- raise ImportError(
- "YAML serialization requires the pyyaml library. Please install"
- " it using 'pip install google-cloud-aiplatform[evaluation]'."
- )
+DiskSpecOrDict = Union[DiskSpec, DiskSpecDict]
- fields_to_exclude_callables = set()
- for field_name, field_info in self.model_fields.items():
- annotation = field_info.annotation
- origin = typing.get_origin(annotation)
- is_field_callable_type = False
- if annotation is Callable or origin is Callable:
- is_field_callable_type = True
- elif origin is Union:
- args = typing.get_args(annotation)
- if any(
- arg is Callable or typing.get_origin(arg) is Callable
- for arg in args
- ):
- is_field_callable_type = True
+class ReservationAffinity(_common.BaseModel):
+ """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity."""
- if is_field_callable_type:
- fields_to_exclude_callables.add(field_name)
+ key: Optional[str] = Field(
+ default=None,
+ description="""Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""",
+ )
+ reservation_affinity_type: Optional[Type] = Field(
+ default=None,
+ description="""Required. Specifies the reservation affinity type.""",
+ )
+ values: Optional[list[str]] = Field(
+ default=None,
+ description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""",
+ )
- data_to_dump = self.model_dump(
- exclude_unset=True,
- exclude_none=True,
- mode="json",
- exclude=fields_to_exclude_callables
- if fields_to_exclude_callables
- else None,
- )
- if version:
- data_to_dump["version"] = version
+class ReservationAffinityDict(TypedDict, total=False):
+ """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity."""
- with open(file_path, "w", encoding="utf-8") as f:
- yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True)
+ key: Optional[str]
+ """Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value."""
+ reservation_affinity_type: Optional[Type]
+ """Required. Specifies the reservation affinity type."""
-class LLMMetric(Metric):
- """A metric that uses LLM-as-a-judge for evaluation."""
+ values: Optional[list[str]]
+ """Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block."""
- @field_validator("prompt_template", mode="before")
- @classmethod
- def validate_prompt_template(cls, value: Union[str, "MetricPromptBuilder"]) -> str:
- """Validates prompt template to be a non-empty string."""
- if value is None:
- raise ValueError("Prompt template cannot be empty.")
- if isinstance(value, MetricPromptBuilder):
- value = str(value)
- if not value.strip():
- raise ValueError("Prompt template cannot be an empty string.")
- return value
- @field_validator("judge_model_sampling_count")
- @classmethod
- def validate_judge_model_sampling_count(cls, value: Optional[int]) -> Optional[int]:
- """Validates judge_model_sampling_count to be between 1 and 32."""
- if value is not None and (value < 1 or value > 32):
- raise ValueError("judge_model_sampling_count must be between 1 and 32.")
- return value
+ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict]
-class MetricDict(TypedDict, total=False):
- """The metric used for evaluation."""
+class MachineSpec(_common.BaseModel):
+ """Specification of a single machine."""
- name: Optional[str]
- """The name of the metric."""
+ accelerator_count: Optional[int] = Field(
+ default=None,
+ description="""The number of accelerators to attach to the machine.""",
+ )
+ accelerator_type: Optional[AcceleratorType] = Field(
+ default=None,
+ description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""",
+ )
+ machine_type: Optional[str] = Field(
+ default=None,
+ description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""",
+ )
+ multihost_gpu_node_count: Optional[int] = Field(
+ default=None,
+ description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""",
+ )
+ reservation_affinity: Optional[ReservationAffinity] = Field(
+ default=None,
+ description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""",
+ )
+ tpu_topology: Optional[str] = Field(
+ default=None,
+ description="""Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""",
+ )
- custom_function: Optional[Callable]
- """The custom function that defines the end-to-end logic for metric computation."""
- prompt_template: Optional[str]
- """The prompt template for the metric."""
+class MachineSpecDict(TypedDict, total=False):
+ """Specification of a single machine."""
- judge_model: Optional[str]
- """The judge model for the metric."""
+ accelerator_count: Optional[int]
+ """The number of accelerators to attach to the machine."""
- judge_model_sampling_count: Optional[int]
- """The sampling count for the judge model."""
+ accelerator_type: Optional[AcceleratorType]
+ """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count."""
- judge_model_system_instruction: Optional[str]
- """The system instruction for the judge model."""
+ machine_type: Optional[str]
+ """Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required."""
- return_raw_output: Optional[bool]
- """Whether to return the raw output from the judge model."""
+ multihost_gpu_node_count: Optional[int]
+ """Optional. Immutable. The number of nodes per replica for multihost GPU deployments."""
- parse_and_reduce_fn: Optional[Callable]
- """The parse and reduce function for the judge model."""
+ reservation_affinity: Optional[ReservationAffinityDict]
+ """Optional. Immutable. Configuration controlling how this resource pool consumes reservation."""
- aggregate_summary_fn: Optional[Callable]
- """The aggregate summary function for the judge model."""
+ tpu_topology: Optional[str]
+ """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1")."""
-MetricOrDict = Union[Metric, MetricDict]
+MachineSpecOrDict = Union[MachineSpec, MachineSpecDict]
-class GcsDestination(_common.BaseModel):
- """The Google Cloud Storage location where the output is to be written to."""
+class NfsMount(_common.BaseModel):
+ """Represents a mount configuration for Network File System (NFS) to mount."""
- output_uri_prefix: Optional[str] = Field(
+ mount_point: Optional[str] = Field(
default=None,
- description="""Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist.""",
+ description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""",
)
-
-
-class GcsDestinationDict(TypedDict, total=False):
- """The Google Cloud Storage location where the output is to be written to."""
-
- output_uri_prefix: Optional[str]
- """Required. Google Cloud Storage URI to output directory. If the uri doesn't end with '/', a '/' will be automatically appended. The directory is created if it doesn't exist."""
-
-
-GcsDestinationOrDict = Union[GcsDestination, GcsDestinationDict]
-
-
-class OutputConfig(_common.BaseModel):
- """Config for evaluation output."""
-
- gcs_destination: Optional[GcsDestination] = Field(
+ path: Optional[str] = Field(
default=None,
- description="""Cloud storage destination for evaluation output.""",
+ description="""Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""",
+ )
+ server: Optional[str] = Field(
+ default=None, description="""Required. IP address of the NFS server."""
)
-class OutputConfigDict(TypedDict, total=False):
- """Config for evaluation output."""
+class NfsMountDict(TypedDict, total=False):
+ """Represents a mount configuration for Network File System (NFS) to mount."""
- gcs_destination: Optional[GcsDestinationDict]
- """Cloud storage destination for evaluation output."""
+ mount_point: Optional[str]
+ """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/"""
+ path: Optional[str]
+ """Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`"""
-OutputConfigOrDict = Union[OutputConfig, OutputConfigDict]
+ server: Optional[str]
+ """Required. IP address of the NFS server."""
-class EvaluateDatasetConfig(_common.BaseModel):
- """Config for evaluate instances."""
+NfsMountOrDict = Union[NfsMount, NfsMountDict]
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
- )
-
-
-class EvaluateDatasetConfigDict(TypedDict, total=False):
- """Config for evaluate instances."""
-
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
-
-
-EvaluateDatasetConfigOrDict = Union[EvaluateDatasetConfig, EvaluateDatasetConfigDict]
+class PythonPackageSpec(_common.BaseModel):
+ """The spec of a Python packaged code."""
-class _EvaluateDatasetRequestParameters(_common.BaseModel):
- """Parameters for batch dataset evaluation."""
-
- dataset: Optional[EvaluationDataset] = Field(default=None, description="""""")
- metrics: Optional[list[Metric]] = Field(default=None, description="""""")
- output_config: Optional[OutputConfig] = Field(default=None, description="""""")
- autorater_config: Optional[AutoraterConfig] = Field(
- default=None, description=""""""
+ args: Optional[list[str]] = Field(
+ default=None,
+ description="""Command line arguments to be passed to the Python task.""",
+ )
+ env: Optional[list[EnvVar]] = Field(
+ default=None,
+ description="""Environment variables to be passed to the python module. Maximum limit is 100.""",
+ )
+ executor_image_uri: Optional[str] = Field(
+ default=None,
+ description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""",
+ )
+ package_uris: Optional[list[str]] = Field(
+ default=None,
+ description="""Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""",
+ )
+ python_module: Optional[str] = Field(
+ default=None,
+ description="""Required. The Python module name to run after installing the packages.""",
)
- config: Optional[EvaluateDatasetConfig] = Field(default=None, description="""""")
-class _EvaluateDatasetRequestParametersDict(TypedDict, total=False):
- """Parameters for batch dataset evaluation."""
+class PythonPackageSpecDict(TypedDict, total=False):
+ """The spec of a Python packaged code."""
- dataset: Optional[EvaluationDatasetDict]
- """"""
+ args: Optional[list[str]]
+ """Command line arguments to be passed to the Python task."""
- metrics: Optional[list[MetricDict]]
- """"""
+ env: Optional[list[EnvVarDict]]
+ """Environment variables to be passed to the python module. Maximum limit is 100."""
- output_config: Optional[OutputConfigDict]
- """"""
+ executor_image_uri: Optional[str]
+ """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list."""
- autorater_config: Optional[AutoraterConfigDict]
- """"""
+ package_uris: Optional[list[str]]
+ """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100."""
- config: Optional[EvaluateDatasetConfigDict]
- """"""
+ python_module: Optional[str]
+ """Required. The Python module name to run after installing the packages."""
-_EvaluateDatasetRequestParametersOrDict = Union[
- _EvaluateDatasetRequestParameters, _EvaluateDatasetRequestParametersDict
-]
+PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict]
-class EvaluateDatasetOperation(_common.BaseModel):
+class WorkerPoolSpec(_common.BaseModel):
+ """Represents the spec of a worker pool in a job."""
- name: Optional[str] = Field(
- default=None,
- description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
+ container_spec: Optional[ContainerSpec] = Field(
+ default=None, description="""The custom container task."""
)
- metadata: Optional[dict[str, Any]] = Field(
+ disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""")
+ machine_spec: Optional[MachineSpec] = Field(
default=None,
- description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
+ description="""Optional. Immutable. The specification of a single machine.""",
)
- done: Optional[bool] = Field(
- default=None,
- description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
+ nfs_mounts: Optional[list[NfsMount]] = Field(
+ default=None, description="""Optional. List of NFS mount spec."""
)
- error: Optional[dict[str, Any]] = Field(
+ python_package_spec: Optional[PythonPackageSpec] = Field(
+ default=None, description="""The Python packaged task."""
+ )
+ replica_count: Optional[int] = Field(
default=None,
- description="""The error result of the operation in case of failure or cancellation.""",
+ description="""Optional. The number of worker replicas to use for this worker pool.""",
)
- response: Optional[EvaluationDataset] = Field(default=None, description="""""")
-class EvaluateDatasetOperationDict(TypedDict, total=False):
-
- name: Optional[str]
- """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
+class WorkerPoolSpecDict(TypedDict, total=False):
+ """Represents the spec of a worker pool in a job."""
- metadata: Optional[dict[str, Any]]
- """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
+ container_spec: Optional[ContainerSpecDict]
+ """The custom container task."""
- done: Optional[bool]
- """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
+ disk_spec: Optional[DiskSpecDict]
+ """Disk spec."""
- error: Optional[dict[str, Any]]
- """The error result of the operation in case of failure or cancellation."""
+ machine_spec: Optional[MachineSpecDict]
+ """Optional. Immutable. The specification of a single machine."""
- response: Optional[EvaluationDatasetDict]
- """"""
+ nfs_mounts: Optional[list[NfsMountDict]]
+ """Optional. List of NFS mount spec."""
+ python_package_spec: Optional[PythonPackageSpecDict]
+ """The Python packaged task."""
-EvaluateDatasetOperationOrDict = Union[
- EvaluateDatasetOperation, EvaluateDatasetOperationDict
-]
+ replica_count: Optional[int]
+ """Optional. The number of worker replicas to use for this worker pool."""
-class OptimizeConfig(_common.BaseModel):
- """Config for Prompt Optimizer."""
+WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict]
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
- )
+class CustomJobSpec(_common.BaseModel):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-class OptimizeConfigDict(TypedDict, total=False):
- """Config for Prompt Optimizer."""
+ base_output_directory: Optional[GcsDestination] = Field(
+ default=None,
+ description="""The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""",
+ )
+ enable_dashboard_access: Optional[bool] = Field(
+ default=None,
+ description="""Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""",
+ )
+ enable_web_access: Optional[bool] = Field(
+ default=None,
+ description="""Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""",
+ )
+ experiment: Optional[str] = Field(
+ default=None,
+ description="""Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""",
+ )
+ experiment_run: Optional[str] = Field(
+ default=None,
+ description="""Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""",
+ )
+ models: Optional[list[str]] = Field(
+ default=None,
+ description="""Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""",
+ )
+ network: Optional[str] = Field(
+ default=None,
+ description="""Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""",
+ )
+ persistent_resource_id: Optional[str] = Field(
+ default=None,
+ description="""Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""",
+ )
+ protected_artifact_location_id: Optional[str] = Field(
+ default=None,
+ description="""The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""",
+ )
+ psc_interface_config: Optional[PscInterfaceConfig] = Field(
+ default=None,
+ description="""Optional. Configuration for PSC-I for CustomJob.""",
+ )
+ reserved_ip_ranges: Optional[list[str]] = Field(
+ default=None,
+ description="""Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""",
+ )
+ scheduling: Optional[Scheduling] = Field(
+ default=None, description="""Scheduling options for a CustomJob."""
+ )
+ service_account: Optional[str] = Field(
+ default=None,
+ description="""Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""",
+ )
+ tensorboard: Optional[str] = Field(
+ default=None,
+ description="""Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""",
+ )
+ worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field(
+ default=None,
+ description="""Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""",
+ )
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
+class CustomJobSpecDict(TypedDict, total=False):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-OptimizeConfigOrDict = Union[OptimizeConfig, OptimizeConfigDict]
+ base_output_directory: Optional[GcsDestinationDict]
+ """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`"""
+ enable_dashboard_access: Optional[bool]
+ """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials)."""
-class _OptimizeRequestParameters(_common.BaseModel):
- """Parameters for the optimize_prompt method."""
+ enable_web_access: Optional[bool]
+ """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials)."""
- config: Optional[OptimizeConfig] = Field(default=None, description="""""")
+ experiment: Optional[str]
+ """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`"""
+ experiment_run: Optional[str]
+ """Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`"""
-class _OptimizeRequestParametersDict(TypedDict, total=False):
- """Parameters for the optimize_prompt method."""
+ models: Optional[list[str]]
+ """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version."""
- config: Optional[OptimizeConfigDict]
- """"""
+ network: Optional[str]
+ """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network."""
+ persistent_resource_id: Optional[str]
+ """Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected."""
-_OptimizeRequestParametersOrDict = Union[
- _OptimizeRequestParameters, _OptimizeRequestParametersDict
-]
+ protected_artifact_location_id: Optional[str]
+ """The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations"""
+ psc_interface_config: Optional[PscInterfaceConfigDict]
+ """Optional. Configuration for PSC-I for CustomJob."""
-class OptimizeResponse(_common.BaseModel):
- """Response for the optimize_prompt method."""
+ reserved_ip_ranges: Optional[list[str]]
+ """Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range']."""
- pass
+ scheduling: Optional[SchedulingDict]
+ """Scheduling options for a CustomJob."""
+ service_account: Optional[str]
+ """Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used."""
-class OptimizeResponseDict(TypedDict, total=False):
- """Response for the optimize_prompt method."""
+ tensorboard: Optional[str]
+ """Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`"""
- pass
+ worker_pool_specs: Optional[list[WorkerPoolSpecDict]]
+ """Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value."""
-OptimizeResponseOrDict = Union[OptimizeResponse, OptimizeResponseDict]
+CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict]
-class PscInterfaceConfig(_common.BaseModel):
- """Configuration for PSC-I."""
+class EncryptionSpec(_common.BaseModel):
+ """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
- network_attachment: Optional[str] = Field(
+ kms_key_name: Optional[str] = Field(
default=None,
- description="""Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I.""",
+ description="""Required. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. The key needs to be in the same region as where the compute resource is created.""",
)
-class PscInterfaceConfigDict(TypedDict, total=False):
- """Configuration for PSC-I."""
+class EncryptionSpecDict(TypedDict, total=False):
+ """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
- network_attachment: Optional[str]
- """Optional. The name of the Compute Engine [network attachment](https://cloud.google.com/vpc/docs/about-network-attachments) to attach to the resource within the region and user project. To specify this field, you must have already [created a network attachment] (https://cloud.google.com/vpc/docs/create-manage-network-attachments#create-network-attachments). This field is only used for resources using PSC-I."""
+ kms_key_name: Optional[str]
+ """Required. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. The key needs to be in the same region as where the compute resource is created."""
-PscInterfaceConfigOrDict = Union[PscInterfaceConfig, PscInterfaceConfigDict]
+EncryptionSpecOrDict = Union[EncryptionSpec, EncryptionSpecDict]
-class Scheduling(_common.BaseModel):
- """All parameters related to queuing and scheduling of custom jobs."""
+class GoogleRpcStatus(_common.BaseModel):
+ """The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs.
- disable_retries: Optional[bool] = Field(
- default=None,
- description="""Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false.""",
- )
- max_wait_duration: Optional[str] = Field(
- default=None,
- description="""Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours.""",
- )
- restart_job_on_worker_restart: Optional[bool] = Field(
+ It is used by [gRPC](https://github.com/grpc). Each `Status` message
+ contains three pieces of data: error code, error message, and error details.
+ You can find out more about this error model and how to work with it in the
+ [API Design Guide](https://cloud.google.com/apis/design/errors).
+ """
+
+ code: Optional[int] = Field(
default=None,
- description="""Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job.""",
+ description="""The status code, which should be an enum value of google.rpc.Code.""",
)
- strategy: Optional[Strategy] = Field(
+ details: Optional[list[dict[str, Any]]] = Field(
default=None,
- description="""Optional. This determines which type of scheduling strategy to use.""",
+ description="""A list of messages that carry the error details. There is a common set of message types for APIs to use.""",
)
- timeout: Optional[str] = Field(
+ message: Optional[str] = Field(
default=None,
- description="""Optional. The maximum job running time. The default is 7 days.""",
+ description="""A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client.""",
)
-class SchedulingDict(TypedDict, total=False):
- """All parameters related to queuing and scheduling of custom jobs."""
-
- disable_retries: Optional[bool]
- """Optional. Indicates if the job should retry for internal errors after the job starts running. If true, overrides `Scheduling.restart_job_on_worker_restart` to false."""
+class GoogleRpcStatusDict(TypedDict, total=False):
+ """The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs.
- max_wait_duration: Optional[str]
- """Optional. This is the maximum duration that a job will wait for the requested resources to be provisioned if the scheduling strategy is set to [Strategy.DWS_FLEX_START]. If set to 0, the job will wait indefinitely. The default is 24 hours."""
+ It is used by [gRPC](https://github.com/grpc). Each `Status` message
+ contains three pieces of data: error code, error message, and error details.
+ You can find out more about this error model and how to work with it in the
+ [API Design Guide](https://cloud.google.com/apis/design/errors).
+ """
- restart_job_on_worker_restart: Optional[bool]
- """Optional. Restarts the entire CustomJob if a worker gets restarted. This feature can be used by distributed training jobs that are not resilient to workers leaving and joining a job."""
+ code: Optional[int]
+ """The status code, which should be an enum value of google.rpc.Code."""
- strategy: Optional[Strategy]
- """Optional. This determines which type of scheduling strategy to use."""
+ details: Optional[list[dict[str, Any]]]
+ """A list of messages that carry the error details. There is a common set of message types for APIs to use."""
- timeout: Optional[str]
- """Optional. The maximum job running time. The default is 7 days."""
+ message: Optional[str]
+ """A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client."""
-SchedulingOrDict = Union[Scheduling, SchedulingDict]
+GoogleRpcStatusOrDict = Union[GoogleRpcStatus, GoogleRpcStatusDict]
-class EnvVar(_common.BaseModel):
- """Represents an environment variable present in a Container or Python Module."""
+class CustomJob(_common.BaseModel):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
+ display_name: Optional[str] = Field(
+ default=None,
+ description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""",
+ )
+ job_spec: Optional[CustomJobSpec] = Field(
+ default=None, description="""Required. Job spec."""
+ )
+ create_time: Optional[datetime.datetime] = Field(
+ default=None,
+ description="""Output only. Time when the CustomJob was created.""",
+ )
+ encryption_spec: Optional[EncryptionSpec] = Field(
+ default=None,
+ description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""",
+ )
+ end_time: Optional[datetime.datetime] = Field(
+ default=None,
+ description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""",
+ )
+ error: Optional[GoogleRpcStatus] = Field(
+ default=None,
+ description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""",
+ )
+ labels: Optional[dict[str, str]] = Field(
+ default=None,
+ description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""",
+ )
name: Optional[str] = Field(
default=None,
- description="""Required. Name of the environment variable. Must be a valid C identifier.""",
+ description="""Output only. Resource name of a CustomJob.""",
)
- value: Optional[str] = Field(
+ satisfies_pzi: Optional[bool] = Field(
+ default=None, description="""Output only. Reserved for future use."""
+ )
+ satisfies_pzs: Optional[bool] = Field(
+ default=None, description="""Output only. Reserved for future use."""
+ )
+ start_time: Optional[datetime.datetime] = Field(
default=None,
- description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""",
+ description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""",
+ )
+ state: Optional[JobState] = Field(
+ default=None,
+ description="""Output only. The detailed state of the job.""",
+ )
+ update_time: Optional[datetime.datetime] = Field(
+ default=None,
+ description="""Output only. Time when the CustomJob was most recently updated.""",
+ )
+ web_access_uris: Optional[dict[str, str]] = Field(
+ default=None,
+ description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""",
)
-class EnvVarDict(TypedDict, total=False):
- """Represents an environment variable present in a Container or Python Module."""
+class CustomJobDict(TypedDict, total=False):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
- name: Optional[str]
- """Required. Name of the environment variable. Must be a valid C identifier."""
+ display_name: Optional[str]
+ """Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters."""
- value: Optional[str]
- """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not."""
+ job_spec: Optional[CustomJobSpecDict]
+ """Required. Job spec."""
+ create_time: Optional[datetime.datetime]
+ """Output only. Time when the CustomJob was created."""
-EnvVarOrDict = Union[EnvVar, EnvVarDict]
+ encryption_spec: Optional[EncryptionSpecDict]
+ """Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key."""
+ end_time: Optional[datetime.datetime]
+ """Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`."""
-class ContainerSpec(_common.BaseModel):
- """The spec of a Container."""
+ error: Optional[GoogleRpcStatusDict]
+ """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`."""
- args: Optional[list[str]] = Field(
- default=None,
- description="""The arguments to be passed when starting the container.""",
- )
- command: Optional[list[str]] = Field(
- default=None,
- description="""The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided.""",
- )
- env: Optional[list[EnvVar]] = Field(
- default=None,
- description="""Environment variables to be passed to the container. Maximum limit is 100.""",
- )
- image_uri: Optional[str] = Field(
- default=None,
- description="""Required. The URI of a container image in the Container Registry that is to be run on each worker replica.""",
- )
+ labels: Optional[dict[str, str]]
+ """The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels."""
+ name: Optional[str]
+ """Output only. Resource name of a CustomJob."""
-class ContainerSpecDict(TypedDict, total=False):
- """The spec of a Container."""
+ satisfies_pzi: Optional[bool]
+ """Output only. Reserved for future use."""
- args: Optional[list[str]]
- """The arguments to be passed when starting the container."""
+ satisfies_pzs: Optional[bool]
+ """Output only. Reserved for future use."""
- command: Optional[list[str]]
- """The command to be invoked when the container is started. It overrides the entrypoint instruction in Dockerfile when provided."""
+ start_time: Optional[datetime.datetime]
+ """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state."""
- env: Optional[list[EnvVarDict]]
- """Environment variables to be passed to the container. Maximum limit is 100."""
+ state: Optional[JobState]
+ """Output only. The detailed state of the job."""
- image_uri: Optional[str]
- """Required. The URI of a container image in the Container Registry that is to be run on each worker replica."""
+ update_time: Optional[datetime.datetime]
+ """Output only. Time when the CustomJob was most recently updated."""
+ web_access_uris: Optional[dict[str, str]]
+ """Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell."""
-ContainerSpecOrDict = Union[ContainerSpec, ContainerSpecDict]
+CustomJobOrDict = Union[CustomJob, CustomJobDict]
-class DiskSpec(_common.BaseModel):
- """Represents the spec of disk options."""
- boot_disk_size_gb: Optional[int] = Field(
- default=None,
- description="""Size in GB of the boot disk (default is 100GB).""",
- )
- boot_disk_type: Optional[str] = Field(
- default=None,
- description="""Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced".""",
+class BaseConfig(_common.BaseModel):
+
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
-class DiskSpecDict(TypedDict, total=False):
- """Represents the spec of disk options."""
+class BaseConfigDict(TypedDict, total=False):
- boot_disk_size_gb: Optional[int]
- """Size in GB of the boot disk (default is 100GB)."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
- boot_disk_type: Optional[str]
- """Type of the boot disk. For non-A3U machines, the default value is "pd-ssd", for A3U machines, the default value is "hyperdisk-balanced". Valid values: "pd-ssd" (Persistent Disk Solid State Drive), "pd-standard" (Persistent Disk Hard Disk Drive) or "hyperdisk-balanced"."""
+BaseConfigOrDict = Union[BaseConfig, BaseConfigDict]
-DiskSpecOrDict = Union[DiskSpec, DiskSpecDict]
+class _CustomJobParameters(_common.BaseModel):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-class ReservationAffinity(_common.BaseModel):
- """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity."""
+ custom_job: Optional[CustomJob] = Field(default=None, description="""""")
+ config: Optional[BaseConfig] = Field(default=None, description="""""")
- key: Optional[str] = Field(
- default=None,
- description="""Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value.""",
- )
- reservation_affinity_type: Optional[Type] = Field(
- default=None,
- description="""Required. Specifies the reservation affinity type.""",
- )
- values: Optional[list[str]] = Field(
- default=None,
- description="""Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block.""",
- )
+class _CustomJobParametersDict(TypedDict, total=False):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-class ReservationAffinityDict(TypedDict, total=False):
- """A ReservationAffinity can be used to configure a Vertex AI resource (e.g., a DeployedModel) to draw its Compute Engine resources from a Shared Reservation, or exclusively from on-demand capacity."""
+ custom_job: Optional[CustomJobDict]
+ """"""
- key: Optional[str]
- """Optional. Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use `compute.googleapis.com/reservation-name` as the key and specify the name of your reservation as its value."""
+ config: Optional[BaseConfigDict]
+ """"""
- reservation_affinity_type: Optional[Type]
- """Required. Specifies the reservation affinity type."""
- values: Optional[list[str]]
- """Optional. Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation or reservation block."""
+_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict]
-ReservationAffinityOrDict = Union[ReservationAffinity, ReservationAffinityDict]
+class _GetCustomJobParameters(_common.BaseModel):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
+ name: Optional[str] = Field(default=None, description="""""")
+ config: Optional[BaseConfig] = Field(default=None, description="""""")
-class MachineSpec(_common.BaseModel):
- """Specification of a single machine."""
- accelerator_count: Optional[int] = Field(
- default=None,
- description="""The number of accelerators to attach to the machine.""",
- )
- accelerator_type: Optional[AcceleratorType] = Field(
- default=None,
- description="""Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count.""",
- )
- machine_type: Optional[str] = Field(
+class _GetCustomJobParametersDict(TypedDict, total=False):
+ """Represents a job that runs custom workloads such as a Docker container or a Python package."""
+
+ name: Optional[str]
+ """"""
+
+ config: Optional[BaseConfigDict]
+ """"""
+
+
+_GetCustomJobParametersOrDict = Union[
+ _GetCustomJobParameters, _GetCustomJobParametersDict
+]
+
+
+class SecretRef(_common.BaseModel):
+ """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
+
+ secret: Optional[str] = Field(
default=None,
- description="""Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required.""",
+ description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""",
)
- multihost_gpu_node_count: Optional[int] = Field(
+ version: Optional[str] = Field(
default=None,
- description="""Optional. Immutable. The number of nodes per replica for multihost GPU deployments.""",
+ description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""",
)
- reservation_affinity: Optional[ReservationAffinity] = Field(
+
+
+class SecretRefDict(TypedDict, total=False):
+ """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
+
+ secret: Optional[str]
+ """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}."""
+
+ version: Optional[str]
+ """The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias."""
+
+
+SecretRefOrDict = Union[SecretRef, SecretRefDict]
+
+
+class SecretEnvVar(_common.BaseModel):
+ """Represents an environment variable where the value is a secret in Cloud Secret Manager."""
+
+ name: Optional[str] = Field(
default=None,
- description="""Optional. Immutable. Configuration controlling how this resource pool consumes reservation.""",
+ description="""Required. Name of the secret environment variable.""",
)
- tpu_topology: Optional[str] = Field(
+ secret_ref: Optional[SecretRef] = Field(
default=None,
- description="""Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1").""",
+ description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""",
)
-class MachineSpecDict(TypedDict, total=False):
- """Specification of a single machine."""
+class SecretEnvVarDict(TypedDict, total=False):
+ """Represents an environment variable where the value is a secret in Cloud Secret Manager."""
- accelerator_count: Optional[int]
- """The number of accelerators to attach to the machine."""
+ name: Optional[str]
+ """Required. Name of the secret environment variable."""
- accelerator_type: Optional[AcceleratorType]
- """Immutable. The type of accelerator(s) that may be attached to the machine as per accelerator_count."""
+ secret_ref: Optional[SecretRefDict]
+ """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
- machine_type: Optional[str]
- """Immutable. The type of the machine. See the [list of machine types supported for prediction](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#machine-types) See the [list of machine types supported for custom training](https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types). For DeployedModel this field is optional, and the default value is `n1-standard-2`. For BatchPredictionJob or as part of WorkerPoolSpec this field is required."""
- multihost_gpu_node_count: Optional[int]
- """Optional. Immutable. The number of nodes per replica for multihost GPU deployments."""
+SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict]
- reservation_affinity: Optional[ReservationAffinityDict]
- """Optional. Immutable. Configuration controlling how this resource pool consumes reservation."""
- tpu_topology: Optional[str]
- """Immutable. The topology of the TPUs. Corresponds to the TPU topologies available from GKE. (Example: tpu_topology: "2x2x1")."""
+class ReasoningEngineSpecDeploymentSpec(_common.BaseModel):
+ """The specification of a Reasoning Engine deployment."""
+
+ env: Optional[list[EnvVar]] = Field(
+ default=None,
+ description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""",
+ )
+ secret_env: Optional[list[SecretEnvVar]] = Field(
+ default=None,
+ description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""",
+ )
-MachineSpecOrDict = Union[MachineSpec, MachineSpecDict]
+class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
+ """The specification of a Reasoning Engine deployment."""
+ env: Optional[list[EnvVarDict]]
+ """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API."""
-class NfsMount(_common.BaseModel):
- """Represents a mount configuration for Network File System (NFS) to mount."""
+ secret_env: Optional[list[SecretEnvVarDict]]
+ """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent."""
- mount_point: Optional[str] = Field(
+
+ReasoningEngineSpecDeploymentSpecOrDict = Union[
+ ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict
+]
+
+
+class ReasoningEngineSpecPackageSpec(_common.BaseModel):
+ """User provided package spec like pickled object and package requirements."""
+
+ dependency_files_gcs_uri: Optional[str] = Field(
default=None,
- description="""Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/""",
+ description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""",
)
- path: Optional[str] = Field(
+ pickle_object_gcs_uri: Optional[str] = Field(
default=None,
- description="""Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`""",
+ description="""Optional. The Cloud Storage URI of the pickled python object.""",
)
- server: Optional[str] = Field(
- default=None, description="""Required. IP address of the NFS server."""
+ python_version: Optional[str] = Field(
+ default=None,
+ description="""Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10.""",
+ )
+ requirements_gcs_uri: Optional[str] = Field(
+ default=None,
+ description="""Optional. The Cloud Storage URI of the `requirements.txt` file""",
)
-class NfsMountDict(TypedDict, total=False):
- """Represents a mount configuration for Network File System (NFS) to mount."""
+class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False):
+ """User provided package spec like pickled object and package requirements."""
- mount_point: Optional[str]
- """Required. Destination mount path. The NFS will be mounted for the user under /mnt/nfs/"""
+ dependency_files_gcs_uri: Optional[str]
+ """Optional. The Cloud Storage URI of the dependency files in tar.gz format."""
- path: Optional[str]
- """Required. Source path exported from NFS server. Has to start with '/', and combined with the ip address, it indicates the source mount path in the form of `server:path`"""
+ pickle_object_gcs_uri: Optional[str]
+ """Optional. The Cloud Storage URI of the pickled python object."""
- server: Optional[str]
- """Required. IP address of the NFS server."""
+ python_version: Optional[str]
+ """Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10."""
+ requirements_gcs_uri: Optional[str]
+ """Optional. The Cloud Storage URI of the `requirements.txt` file"""
-NfsMountOrDict = Union[NfsMount, NfsMountDict]
+ReasoningEngineSpecPackageSpecOrDict = Union[
+ ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict
+]
-class PythonPackageSpec(_common.BaseModel):
- """The spec of a Python packaged code."""
- args: Optional[list[str]] = Field(
- default=None,
- description="""Command line arguments to be passed to the Python task.""",
- )
- env: Optional[list[EnvVar]] = Field(
+class ReasoningEngineSpec(_common.BaseModel):
+ """The specification of a Reasoning Engine."""
+
+ agent_framework: Optional[str] = Field(
default=None,
- description="""Environment variables to be passed to the python module. Maximum limit is 100.""",
+ description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""",
)
- executor_image_uri: Optional[str] = Field(
+ class_methods: Optional[list[dict[str, Any]]] = Field(
default=None,
- description="""Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list.""",
+ description="""Optional. Declarations for object class methods in OpenAPI specification format.""",
)
- package_uris: Optional[list[str]] = Field(
+ deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field(
default=None,
- description="""Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100.""",
+ description="""Optional. The specification of a Reasoning Engine deployment.""",
)
- python_module: Optional[str] = Field(
+ package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field(
default=None,
- description="""Required. The Python module name to run after installing the packages.""",
+ description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes.""",
)
-class PythonPackageSpecDict(TypedDict, total=False):
- """The spec of a Python packaged code."""
-
- args: Optional[list[str]]
- """Command line arguments to be passed to the Python task."""
+class ReasoningEngineSpecDict(TypedDict, total=False):
+ """The specification of a Reasoning Engine."""
- env: Optional[list[EnvVarDict]]
- """Environment variables to be passed to the python module. Maximum limit is 100."""
+ agent_framework: Optional[str]
+ """Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"."""
- executor_image_uri: Optional[str]
- """Required. The URI of a container image in Artifact Registry that will run the provided Python package. Vertex AI provides a wide range of executor images with pre-installed packages to meet users' various use cases. See the list of [pre-built containers for training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers). You must use an image from this list."""
+ class_methods: Optional[list[dict[str, Any]]]
+ """Optional. Declarations for object class methods in OpenAPI specification format."""
- package_uris: Optional[list[str]]
- """Required. The Google Cloud Storage location of the Python package files which are the training program and its dependent packages. The maximum number of package URIs is 100."""
+ deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict]
+ """Optional. The specification of a Reasoning Engine deployment."""
- python_module: Optional[str]
- """Required. The Python module name to run after installing the packages."""
+ package_spec: Optional[ReasoningEngineSpecPackageSpecDict]
+ """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes."""
-PythonPackageSpecOrDict = Union[PythonPackageSpec, PythonPackageSpecDict]
+ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict]
-class WorkerPoolSpec(_common.BaseModel):
- """Represents the spec of a worker pool in a job."""
+class CreateAgentEngineConfig(_common.BaseModel):
+ """Config for create agent engine."""
- container_spec: Optional[ContainerSpec] = Field(
- default=None, description="""The custom container task."""
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
- disk_spec: Optional[DiskSpec] = Field(default=None, description="""Disk spec.""")
- machine_spec: Optional[MachineSpec] = Field(
+ display_name: Optional[str] = Field(
default=None,
- description="""Optional. Immutable. The specification of a single machine.""",
- )
- nfs_mounts: Optional[list[NfsMount]] = Field(
- default=None, description="""Optional. List of NFS mount spec."""
+ description="""The user-defined name of the Agent Engine.
+
+ The display name can be up to 128 characters long and can comprise any
+ UTF-8 characters.
+ """,
)
- python_package_spec: Optional[PythonPackageSpec] = Field(
- default=None, description="""The Python packaged task."""
+ description: Optional[str] = Field(
+ default=None, description="""The description of the Agent Engine."""
)
- replica_count: Optional[int] = Field(
+ spec: Optional[ReasoningEngineSpec] = Field(
default=None,
- description="""Optional. The number of worker replicas to use for this worker pool.""",
+ description="""Optional. Configurations of the ReasoningEngine.""",
)
-class WorkerPoolSpecDict(TypedDict, total=False):
- """Represents the spec of a worker pool in a job."""
+class CreateAgentEngineConfigDict(TypedDict, total=False):
+ """Config for create agent engine."""
- container_spec: Optional[ContainerSpecDict]
- """The custom container task."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
- disk_spec: Optional[DiskSpecDict]
- """Disk spec."""
+ display_name: Optional[str]
+ """The user-defined name of the Agent Engine.
- machine_spec: Optional[MachineSpecDict]
- """Optional. Immutable. The specification of a single machine."""
+ The display name can be up to 128 characters long and can comprise any
+ UTF-8 characters.
+ """
- nfs_mounts: Optional[list[NfsMountDict]]
- """Optional. List of NFS mount spec."""
+ description: Optional[str]
+ """The description of the Agent Engine."""
- python_package_spec: Optional[PythonPackageSpecDict]
- """The Python packaged task."""
+ spec: Optional[ReasoningEngineSpecDict]
+ """Optional. Configurations of the ReasoningEngine."""
- replica_count: Optional[int]
- """Optional. The number of worker replicas to use for this worker pool."""
+CreateAgentEngineConfigOrDict = Union[
+ CreateAgentEngineConfig, CreateAgentEngineConfigDict
+]
-WorkerPoolSpecOrDict = Union[WorkerPoolSpec, WorkerPoolSpecDict]
+class _CreateAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for creating agent engines."""
-class CustomJobSpec(_common.BaseModel):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
+ config: Optional[CreateAgentEngineConfig] = Field(default=None, description="""""")
- base_output_directory: Optional[GcsDestination] = Field(
- default=None,
- description="""The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`""",
- )
- enable_dashboard_access: Optional[bool] = Field(
- default=None,
- description="""Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""",
- )
- enable_web_access: Optional[bool] = Field(
+
+class _CreateAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for creating agent engines."""
+
+ config: Optional[CreateAgentEngineConfigDict]
+ """"""
+
+
+_CreateAgentEngineRequestParametersOrDict = Union[
+ _CreateAgentEngineRequestParameters, _CreateAgentEngineRequestParametersDict
+]
+
+
+class ReasoningEngine(_common.BaseModel):
+ """An agent engine."""
+
+ create_time: Optional[datetime.datetime] = Field(
default=None,
- description="""Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials).""",
+ description="""Output only. Timestamp when this ReasoningEngine was created.""",
)
- experiment: Optional[str] = Field(
+ description: Optional[str] = Field(
default=None,
- description="""Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`""",
+ description="""Optional. The description of the ReasoningEngine.""",
)
- experiment_run: Optional[str] = Field(
+ display_name: Optional[str] = Field(
default=None,
- description="""Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`""",
+ description="""Required. The display name of the ReasoningEngine.""",
)
- models: Optional[list[str]] = Field(
+ etag: Optional[str] = Field(
default=None,
- description="""Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version.""",
+ description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""",
)
- network: Optional[str] = Field(
+ name: Optional[str] = Field(
default=None,
- description="""Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network.""",
+ description="""Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""",
)
- persistent_resource_id: Optional[str] = Field(
+ spec: Optional[ReasoningEngineSpec] = Field(
default=None,
- description="""Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.""",
+ description="""Optional. Configurations of the ReasoningEngine""",
)
- protected_artifact_location_id: Optional[str] = Field(
+ update_time: Optional[datetime.datetime] = Field(
default=None,
- description="""The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations""",
+ description="""Output only. Timestamp when this ReasoningEngine was most recently updated.""",
)
- psc_interface_config: Optional[PscInterfaceConfig] = Field(
+
+
+class ReasoningEngineDict(TypedDict, total=False):
+ """An agent engine."""
+
+ create_time: Optional[datetime.datetime]
+ """Output only. Timestamp when this ReasoningEngine was created."""
+
+ description: Optional[str]
+ """Optional. The description of the ReasoningEngine."""
+
+ display_name: Optional[str]
+ """Required. The display name of the ReasoningEngine."""
+
+ etag: Optional[str]
+ """Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens."""
+
+ name: Optional[str]
+ """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`"""
+
+ spec: Optional[ReasoningEngineSpecDict]
+ """Optional. Configurations of the ReasoningEngine"""
+
+ update_time: Optional[datetime.datetime]
+ """Output only. Timestamp when this ReasoningEngine was most recently updated."""
+
+
+ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict]
+
+
+class AgentEngineOperation(_common.BaseModel):
+ """Operation that has an agent engine as a response."""
+
+ name: Optional[str] = Field(
default=None,
- description="""Optional. Configuration for PSC-I for CustomJob.""",
+ description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
)
- reserved_ip_ranges: Optional[list[str]] = Field(
+ metadata: Optional[dict[str, Any]] = Field(
default=None,
- description="""Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range'].""",
- )
- scheduling: Optional[Scheduling] = Field(
- default=None, description="""Scheduling options for a CustomJob."""
+ description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
)
- service_account: Optional[str] = Field(
+ done: Optional[bool] = Field(
default=None,
- description="""Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used.""",
+ description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
)
- tensorboard: Optional[str] = Field(
+ error: Optional[dict[str, Any]] = Field(
default=None,
- description="""Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`""",
+ description="""The error result of the operation in case of failure or cancellation.""",
)
- worker_pool_specs: Optional[list[WorkerPoolSpec]] = Field(
- default=None,
- description="""Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value.""",
+ response: Optional[ReasoningEngine] = Field(
+ default=None, description="""The created Agent Engine."""
)
-class CustomJobSpecDict(TypedDict, total=False):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-
- base_output_directory: Optional[GcsDestinationDict]
- """The Cloud Storage location to store the output of this CustomJob or HyperparameterTuningJob. For HyperparameterTuningJob, the baseOutputDirectory of each child CustomJob backing a Trial is set to a subdirectory of name id under its parent HyperparameterTuningJob's baseOutputDirectory. The following Vertex AI environment variables will be passed to containers or python modules when this field is set: For CustomJob: * AIP_MODEL_DIR = `/model/` * AIP_CHECKPOINT_DIR = `/checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `/logs/` For CustomJob backing a Trial of HyperparameterTuningJob: * AIP_MODEL_DIR = `//model/` * AIP_CHECKPOINT_DIR = `//checkpoints/` * AIP_TENSORBOARD_LOG_DIR = `//logs/`"""
-
- enable_dashboard_access: Optional[bool]
- """Optional. Whether you want Vertex AI to enable access to the customized dashboard in training chief container. If set to `true`, you can access the dashboard at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials)."""
+class AgentEngineOperationDict(TypedDict, total=False):
+ """Operation that has an agent engine as a response."""
- enable_web_access: Optional[bool]
- """Optional. Whether you want Vertex AI to enable [interactive shell access](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) to training containers. If set to `true`, you can access interactive shells at the URIs given by CustomJob.web_access_uris or Trial.web_access_uris (within HyperparameterTuningJob.trials)."""
+ name: Optional[str]
+ """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
- experiment: Optional[str]
- """Optional. The Experiment associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}`"""
+ metadata: Optional[dict[str, Any]]
+ """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
- experiment_run: Optional[str]
- """Optional. The Experiment Run associated with this job. Format: `projects/{project}/locations/{location}/metadataStores/{metadataStores}/contexts/{experiment-name}-{experiment-run-name}`"""
+ done: Optional[bool]
+ """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
- models: Optional[list[str]]
- """Optional. The name of the Model resources for which to generate a mapping to artifact URIs. Applicable only to some of the Google-provided custom jobs. Format: `projects/{project}/locations/{location}/models/{model}` In order to retrieve a specific version of the model, also provide the version ID or version alias. Example: `projects/{project}/locations/{location}/models/{model}@2` or `projects/{project}/locations/{location}/models/{model}@golden` If no version ID or alias is specified, the "default" version will be returned. The "default" version alias is created for the first version of the model, and can be moved to other versions later on. There will be exactly one default version."""
+ error: Optional[dict[str, Any]]
+ """The error result of the operation in case of failure or cancellation."""
- network: Optional[str]
- """Optional. The full name of the Compute Engine [network](/compute/docs/networks-and-firewalls#networks) to which the Job should be peered. For example, `projects/12345/global/networks/myVPC`. [Format](/compute/docs/reference/rest/v1/networks/insert) is of the form `projects/{project}/global/networks/{network}`. Where {project} is a project number, as in `12345`, and {network} is a network name. To specify this field, you must have already [configured VPC Network Peering for Vertex AI](https://cloud.google.com/vertex-ai/docs/general/vpc-peering). If this field is left unspecified, the job is not peered with any network."""
+ response: Optional[ReasoningEngineDict]
+ """The created Agent Engine."""
- persistent_resource_id: Optional[str]
- """Optional. The ID of the PersistentResource in the same Project and Location which to run If this is specified, the job will be run on existing machines held by the PersistentResource instead of on-demand short-live machines. The network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected."""
- protected_artifact_location_id: Optional[str]
- """The ID of the location to store protected artifacts. e.g. us-central1. Populate only when the location is different than CustomJob location. List of supported locations: https://cloud.google.com/vertex-ai/docs/general/locations"""
+AgentEngineOperationOrDict = Union[AgentEngineOperation, AgentEngineOperationDict]
- psc_interface_config: Optional[PscInterfaceConfigDict]
- """Optional. Configuration for PSC-I for CustomJob."""
- reserved_ip_ranges: Optional[list[str]]
- """Optional. A list of names for the reserved ip ranges under the VPC network that can be used for this job. If set, we will deploy the job within the provided ip ranges. Otherwise, the job will be deployed to any ip ranges under the provided VPC network. Example: ['vertex-ai-ip-range']."""
+class DeleteAgentEngineConfig(_common.BaseModel):
+ """Config for deleting agent engine."""
- scheduling: Optional[SchedulingDict]
- """Scheduling options for a CustomJob."""
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
+ )
- service_account: Optional[str]
- """Specifies the service account for workload run-as account. Users submitting jobs must have act-as permission on this run-as account. If unspecified, the [Vertex AI Custom Code Service Agent](https://cloud.google.com/vertex-ai/docs/general/access-control#service-agents) for the CustomJob's project is used."""
- tensorboard: Optional[str]
- """Optional. The name of a Vertex AI Tensorboard resource to which this CustomJob will upload Tensorboard logs. Format: `projects/{project}/locations/{location}/tensorboards/{tensorboard}`"""
+class DeleteAgentEngineConfigDict(TypedDict, total=False):
+ """Config for deleting agent engine."""
- worker_pool_specs: Optional[list[WorkerPoolSpecDict]]
- """Required. The spec of the worker pools including machine type and Docker image. All worker pools except the first one are optional and can be skipped by providing an empty value."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
-CustomJobSpecOrDict = Union[CustomJobSpec, CustomJobSpecDict]
+DeleteAgentEngineConfigOrDict = Union[
+ DeleteAgentEngineConfig, DeleteAgentEngineConfigDict
+]
-class EncryptionSpec(_common.BaseModel):
- """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
+class _DeleteAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for deleting agent engines."""
- kms_key_name: Optional[str] = Field(
- default=None,
- description="""Required. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. The key needs to be in the same region as where the compute resource is created.""",
+ name: Optional[str] = Field(
+ default=None, description="""Name of the agent engine."""
+ )
+ force: Optional[bool] = Field(
+ default=False,
+ description="""If set to true, any child resources will also be deleted.""",
)
+ config: Optional[DeleteAgentEngineConfig] = Field(default=None, description="""""")
-class EncryptionSpecDict(TypedDict, total=False):
- """Represents a customer-managed encryption key spec that can be applied to a top-level resource."""
+class _DeleteAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for deleting agent engines."""
- kms_key_name: Optional[str]
- """Required. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. The key needs to be in the same region as where the compute resource is created."""
+ name: Optional[str]
+ """Name of the agent engine."""
+
+ force: Optional[bool]
+ """If set to true, any child resources will also be deleted."""
+ config: Optional[DeleteAgentEngineConfigDict]
+ """"""
-EncryptionSpecOrDict = Union[EncryptionSpec, EncryptionSpecDict]
+_DeleteAgentEngineRequestParametersOrDict = Union[
+ _DeleteAgentEngineRequestParameters, _DeleteAgentEngineRequestParametersDict
+]
-class GoogleRpcStatus(_common.BaseModel):
- """The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs.
- It is used by [gRPC](https://github.com/grpc). Each `Status` message
- contains three pieces of data: error code, error message, and error details.
- You can find out more about this error model and how to work with it in the
- [API Design Guide](https://cloud.google.com/apis/design/errors).
- """
+class DeleteAgentEngineOperation(_common.BaseModel):
+ """Operation for deleting agent engines."""
- code: Optional[int] = Field(
+ name: Optional[str] = Field(
default=None,
- description="""The status code, which should be an enum value of google.rpc.Code.""",
+ description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
)
- details: Optional[list[dict[str, Any]]] = Field(
+ metadata: Optional[dict[str, Any]] = Field(
default=None,
- description="""A list of messages that carry the error details. There is a common set of message types for APIs to use.""",
+ description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
)
- message: Optional[str] = Field(
+ done: Optional[bool] = Field(
default=None,
- description="""A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client.""",
+ description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
+ )
+ error: Optional[dict[str, Any]] = Field(
+ default=None,
+ description="""The error result of the operation in case of failure or cancellation.""",
)
-class GoogleRpcStatusDict(TypedDict, total=False):
- """The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs.
+class DeleteAgentEngineOperationDict(TypedDict, total=False):
+ """Operation for deleting agent engines."""
- It is used by [gRPC](https://github.com/grpc). Each `Status` message
- contains three pieces of data: error code, error message, and error details.
- You can find out more about this error model and how to work with it in the
- [API Design Guide](https://cloud.google.com/apis/design/errors).
- """
+ name: Optional[str]
+ """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
- code: Optional[int]
- """The status code, which should be an enum value of google.rpc.Code."""
+ metadata: Optional[dict[str, Any]]
+ """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
- details: Optional[list[dict[str, Any]]]
- """A list of messages that carry the error details. There is a common set of message types for APIs to use."""
+ done: Optional[bool]
+ """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
- message: Optional[str]
- """A developer-facing error message, which should be in English. Any user-facing error message should be localized and sent in the google.rpc.Status.details field, or localized by the client."""
+ error: Optional[dict[str, Any]]
+ """The error result of the operation in case of failure or cancellation."""
-GoogleRpcStatusOrDict = Union[GoogleRpcStatus, GoogleRpcStatusDict]
+DeleteAgentEngineOperationOrDict = Union[
+ DeleteAgentEngineOperation, DeleteAgentEngineOperationDict
+]
-class CustomJob(_common.BaseModel):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
+class GetAgentEngineConfig(_common.BaseModel):
+ """Config for create agent engine."""
- display_name: Optional[str] = Field(
- default=None,
- description="""Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters.""",
- )
- job_spec: Optional[CustomJobSpec] = Field(
- default=None, description="""Required. Job spec."""
- )
- create_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Time when the CustomJob was created.""",
- )
- encryption_spec: Optional[EncryptionSpec] = Field(
- default=None,
- description="""Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.""",
- )
- end_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`.""",
- )
- error: Optional[GoogleRpcStatus] = Field(
- default=None,
- description="""Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`.""",
- )
- labels: Optional[dict[str, str]] = Field(
- default=None,
- description="""The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""",
- )
- name: Optional[str] = Field(
- default=None,
- description="""Output only. Resource name of a CustomJob.""",
- )
- satisfies_pzi: Optional[bool] = Field(
- default=None, description="""Output only. Reserved for future use."""
- )
- satisfies_pzs: Optional[bool] = Field(
- default=None, description="""Output only. Reserved for future use."""
- )
- start_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state.""",
- )
- state: Optional[JobState] = Field(
- default=None,
- description="""Output only. The detailed state of the job.""",
- )
- update_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Time when the CustomJob was most recently updated.""",
- )
- web_access_uris: Optional[dict[str, str]] = Field(
- default=None,
- description="""Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell.""",
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
-class CustomJobDict(TypedDict, total=False):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-
- display_name: Optional[str]
- """Required. The display name of the CustomJob. The name can be up to 128 characters long and can consist of any UTF-8 characters."""
-
- job_spec: Optional[CustomJobSpecDict]
- """Required. Job spec."""
-
- create_time: Optional[datetime.datetime]
- """Output only. Time when the CustomJob was created."""
-
- encryption_spec: Optional[EncryptionSpecDict]
- """Customer-managed encryption key options for a CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key."""
+class GetAgentEngineConfigDict(TypedDict, total=False):
+ """Config for create agent engine."""
- end_time: Optional[datetime.datetime]
- """Output only. Time when the CustomJob entered any of the following states: `JOB_STATE_SUCCEEDED`, `JOB_STATE_FAILED`, `JOB_STATE_CANCELLED`."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
- error: Optional[GoogleRpcStatusDict]
- """Output only. Only populated when job's state is `JOB_STATE_FAILED` or `JOB_STATE_CANCELLED`."""
- labels: Optional[dict[str, str]]
- """The labels with user-defined metadata to organize CustomJobs. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels."""
+GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict]
- name: Optional[str]
- """Output only. Resource name of a CustomJob."""
- satisfies_pzi: Optional[bool]
- """Output only. Reserved for future use."""
+class _GetAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for getting agent engines."""
- satisfies_pzs: Optional[bool]
- """Output only. Reserved for future use."""
+ name: Optional[str] = Field(
+ default=None, description="""Name of the agent engine."""
+ )
+ config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""")
- start_time: Optional[datetime.datetime]
- """Output only. Time when the CustomJob for the first time entered the `JOB_STATE_RUNNING` state."""
- state: Optional[JobState]
- """Output only. The detailed state of the job."""
+class _GetAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for getting agent engines."""
- update_time: Optional[datetime.datetime]
- """Output only. Time when the CustomJob was most recently updated."""
+ name: Optional[str]
+ """Name of the agent engine."""
- web_access_uris: Optional[dict[str, str]]
- """Output only. URIs for accessing [interactive shells](https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell) (one URI for each training node). Only available if job_spec.enable_web_access is `true`. The keys are names of each node in the training job; for example, `workerpool0-0` for the primary node, `workerpool1-0` for the first node in the second worker pool, and `workerpool1-1` for the second node in the second worker pool. The values are the URIs for each node's interactive shell."""
+ config: Optional[GetAgentEngineConfigDict]
+ """"""
-CustomJobOrDict = Union[CustomJob, CustomJobDict]
+_GetAgentEngineRequestParametersOrDict = Union[
+ _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict
+]
-class BaseConfig(_common.BaseModel):
+class ListAgentEngineConfig(_common.BaseModel):
+ """Config for listing agent engines."""
http_options: Optional[HttpOptions] = Field(
default=None, description="""Used to override HTTP request options."""
)
+ page_size: Optional[int] = Field(default=None, description="""""")
+ page_token: Optional[str] = Field(default=None, description="""""")
+ filter: Optional[str] = Field(
+ default=None,
+ description="""An expression for filtering the results of the request.
+ For field names both snake_case and camelCase are supported.""",
+ )
-class BaseConfigDict(TypedDict, total=False):
+class ListAgentEngineConfigDict(TypedDict, total=False):
+ """Config for listing agent engines."""
http_options: Optional[HttpOptionsDict]
"""Used to override HTTP request options."""
-
-BaseConfigOrDict = Union[BaseConfig, BaseConfigDict]
-
-
-class _CustomJobParameters(_common.BaseModel):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-
- custom_job: Optional[CustomJob] = Field(default=None, description="""""")
- config: Optional[BaseConfig] = Field(default=None, description="""""")
-
-
-class _CustomJobParametersDict(TypedDict, total=False):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
-
- custom_job: Optional[CustomJobDict]
+ page_size: Optional[int]
""""""
- config: Optional[BaseConfigDict]
+ page_token: Optional[str]
""""""
+ filter: Optional[str]
+ """An expression for filtering the results of the request.
+ For field names both snake_case and camelCase are supported."""
-_CustomJobParametersOrDict = Union[_CustomJobParameters, _CustomJobParametersDict]
+ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict]
-class _GetCustomJobParameters(_common.BaseModel):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
- name: Optional[str] = Field(default=None, description="""""")
- config: Optional[BaseConfig] = Field(default=None, description="""""")
+class _ListAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for listing agent engines."""
+ config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""")
-class _GetCustomJobParametersDict(TypedDict, total=False):
- """Represents a job that runs custom workloads such as a Docker container or a Python package."""
- name: Optional[str]
- """"""
+class _ListAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for listing agent engines."""
- config: Optional[BaseConfigDict]
+ config: Optional[ListAgentEngineConfigDict]
""""""
-_GetCustomJobParametersOrDict = Union[
- _GetCustomJobParameters, _GetCustomJobParametersDict
+_ListAgentEngineRequestParametersOrDict = Union[
+ _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict
]
-class SecretRef(_common.BaseModel):
- """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
+class ListReasoningEnginesResponse(_common.BaseModel):
+ """Response for listing agent engines."""
- secret: Optional[str] = Field(
- default=None,
- description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""",
- )
- version: Optional[str] = Field(
+ next_page_token: Optional[str] = Field(default=None, description="""""")
+ reasoning_engines: Optional[list[ReasoningEngine]] = Field(
default=None,
- description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""",
+ description="""List of agent engines.
+ """,
)
-class SecretRefDict(TypedDict, total=False):
- """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
+class ListReasoningEnginesResponseDict(TypedDict, total=False):
+ """Response for listing agent engines."""
- secret: Optional[str]
- """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}."""
+ next_page_token: Optional[str]
+ """"""
- version: Optional[str]
- """The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias."""
+ reasoning_engines: Optional[list[ReasoningEngineDict]]
+ """List of agent engines.
+ """
-SecretRefOrDict = Union[SecretRef, SecretRefDict]
+ListReasoningEnginesResponseOrDict = Union[
+ ListReasoningEnginesResponse, ListReasoningEnginesResponseDict
+]
-class SecretEnvVar(_common.BaseModel):
- """Represents an environment variable where the value is a secret in Cloud Secret Manager."""
+class GetAgentEngineOperationConfig(_common.BaseModel):
- name: Optional[str] = Field(
- default=None,
- description="""Required. Name of the secret environment variable.""",
- )
- secret_ref: Optional[SecretRef] = Field(
- default=None,
- description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""",
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
-class SecretEnvVarDict(TypedDict, total=False):
- """Represents an environment variable where the value is a secret in Cloud Secret Manager."""
-
- name: Optional[str]
- """Required. Name of the secret environment variable."""
+class GetAgentEngineOperationConfigDict(TypedDict, total=False):
- secret_ref: Optional[SecretRefDict]
- """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
-SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict]
+GetAgentEngineOperationConfigOrDict = Union[
+ GetAgentEngineOperationConfig, GetAgentEngineOperationConfigDict
+]
-class ReasoningEngineSpecDeploymentSpec(_common.BaseModel):
- """The specification of a Reasoning Engine deployment."""
+class _GetAgentEngineOperationParameters(_common.BaseModel):
+ """Parameters for the GET method."""
- env: Optional[list[EnvVar]] = Field(
+ operation_name: Optional[str] = Field(
default=None,
- description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""",
+ description="""The server-assigned name for the operation.""",
)
- secret_env: Optional[list[SecretEnvVar]] = Field(
+ config: Optional[GetAgentEngineOperationConfig] = Field(
default=None,
- description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""",
+ description="""Used to override the default configuration.""",
)
-class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False):
- """The specification of a Reasoning Engine deployment."""
+class _GetAgentEngineOperationParametersDict(TypedDict, total=False):
+ """Parameters for the GET method."""
- env: Optional[list[EnvVarDict]]
- """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API."""
+ operation_name: Optional[str]
+ """The server-assigned name for the operation."""
- secret_env: Optional[list[SecretEnvVarDict]]
- """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent."""
+ config: Optional[GetAgentEngineOperationConfigDict]
+ """Used to override the default configuration."""
-ReasoningEngineSpecDeploymentSpecOrDict = Union[
- ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict
+_GetAgentEngineOperationParametersOrDict = Union[
+ _GetAgentEngineOperationParameters, _GetAgentEngineOperationParametersDict
]
-class ReasoningEngineSpecPackageSpec(_common.BaseModel):
- """User provided package spec like pickled object and package requirements."""
+class QueryAgentEngineConfig(_common.BaseModel):
+ """Config for querying agent engines."""
- dependency_files_gcs_uri: Optional[str] = Field(
- default=None,
- description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""",
- )
- pickle_object_gcs_uri: Optional[str] = Field(
- default=None,
- description="""Optional. The Cloud Storage URI of the pickled python object.""",
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
)
- python_version: Optional[str] = Field(
- default=None,
- description="""Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10.""",
+ class_method: Optional[str] = Field(
+ default=None, description="""The class method to call."""
)
- requirements_gcs_uri: Optional[str] = Field(
- default=None,
- description="""Optional. The Cloud Storage URI of the `requirements.txt` file""",
+ input: Optional[dict[str, Any]] = Field(
+ default=None, description="""The input to the class method."""
)
+ include_all_fields: Optional[bool] = Field(default=False, description="""""")
-class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False):
- """User provided package spec like pickled object and package requirements."""
+class QueryAgentEngineConfigDict(TypedDict, total=False):
+ """Config for querying agent engines."""
- dependency_files_gcs_uri: Optional[str]
- """Optional. The Cloud Storage URI of the dependency files in tar.gz format."""
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
- pickle_object_gcs_uri: Optional[str]
- """Optional. The Cloud Storage URI of the pickled python object."""
+ class_method: Optional[str]
+ """The class method to call."""
- python_version: Optional[str]
- """Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10."""
+ input: Optional[dict[str, Any]]
+ """The input to the class method."""
- requirements_gcs_uri: Optional[str]
- """Optional. The Cloud Storage URI of the `requirements.txt` file"""
+ include_all_fields: Optional[bool]
+ """"""
-ReasoningEngineSpecPackageSpecOrDict = Union[
- ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict
-]
+QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict]
-class ReasoningEngineSpec(_common.BaseModel):
- """The specification of a Reasoning Engine."""
+class _QueryAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for querying agent engines."""
- agent_framework: Optional[str] = Field(
- default=None,
- description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""",
- )
- class_methods: Optional[list[dict[str, Any]]] = Field(
- default=None,
- description="""Optional. Declarations for object class methods in OpenAPI specification format.""",
- )
- deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field(
- default=None,
- description="""Optional. The specification of a Reasoning Engine deployment.""",
- )
- package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field(
- default=None,
- description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes.""",
+ name: Optional[str] = Field(
+ default=None, description="""Name of the agent engine."""
)
+ config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""")
-class ReasoningEngineSpecDict(TypedDict, total=False):
- """The specification of a Reasoning Engine."""
+class _QueryAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for querying agent engines."""
- agent_framework: Optional[str]
- """Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom"."""
+ name: Optional[str]
+ """Name of the agent engine."""
- class_methods: Optional[list[dict[str, Any]]]
- """Optional. Declarations for object class methods in OpenAPI specification format."""
+ config: Optional[QueryAgentEngineConfigDict]
+ """"""
- deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict]
- """Optional. The specification of a Reasoning Engine deployment."""
- package_spec: Optional[ReasoningEngineSpecPackageSpecDict]
- """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes."""
+_QueryAgentEngineRequestParametersOrDict = Union[
+ _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict
+]
-ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict]
+class QueryReasoningEngineResponse(_common.BaseModel):
+ """The response for querying an agent engine."""
+
+ output: Optional[Any] = Field(
+ default=None,
+ description="""Response provided by users in JSON object format.""",
+ )
-class CreateAgentEngineConfig(_common.BaseModel):
- """Config for create agent engine."""
+class QueryReasoningEngineResponseDict(TypedDict, total=False):
+ """The response for querying an agent engine."""
+
+ output: Optional[Any]
+ """Response provided by users in JSON object format."""
+
+
+QueryReasoningEngineResponseOrDict = Union[
+ QueryReasoningEngineResponse, QueryReasoningEngineResponseDict
+]
+
+
+class UpdateAgentEngineConfig(_common.BaseModel):
+ """Config for updating agent engine."""
http_options: Optional[HttpOptions] = Field(
default=None, description="""Used to override HTTP request options."""
@@ -3529,10 +3526,15 @@ class CreateAgentEngineConfig(_common.BaseModel):
default=None,
description="""Optional. Configurations of the ReasoningEngine.""",
)
+ update_mask: Optional[str] = Field(
+ default=None,
+ description="""The update mask to apply. For the `FieldMask` definition, see
+ https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""",
+ )
-class CreateAgentEngineConfigDict(TypedDict, total=False):
- """Config for create agent engine."""
+class UpdateAgentEngineConfigDict(TypedDict, total=False):
+ """Config for updating agent engine."""
http_options: Optional[HttpOptionsDict]
"""Used to override HTTP request options."""
@@ -3550,1014 +3552,901 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
spec: Optional[ReasoningEngineSpecDict]
"""Optional. Configurations of the ReasoningEngine."""
-
-CreateAgentEngineConfigOrDict = Union[
- CreateAgentEngineConfig, CreateAgentEngineConfigDict
-]
-
-
-class _CreateAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for creating agent engines."""
-
- config: Optional[CreateAgentEngineConfig] = Field(default=None, description="""""")
-
-
-class _CreateAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for creating agent engines."""
-
- config: Optional[CreateAgentEngineConfigDict]
- """"""
+ update_mask: Optional[str]
+ """The update mask to apply. For the `FieldMask` definition, see
+ https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
-_CreateAgentEngineRequestParametersOrDict = Union[
- _CreateAgentEngineRequestParameters, _CreateAgentEngineRequestParametersDict
+UpdateAgentEngineConfigOrDict = Union[
+ UpdateAgentEngineConfig, UpdateAgentEngineConfigDict
]
-class ReasoningEngine(_common.BaseModel):
- """An agent engine."""
+class _UpdateAgentEngineRequestParameters(_common.BaseModel):
+ """Parameters for updating agent engines."""
- create_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Timestamp when this ReasoningEngine was created.""",
- )
- description: Optional[str] = Field(
- default=None,
- description="""Optional. The description of the ReasoningEngine.""",
- )
- display_name: Optional[str] = Field(
- default=None,
- description="""Required. The display name of the ReasoningEngine.""",
- )
- etag: Optional[str] = Field(
- default=None,
- description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""",
- )
name: Optional[str] = Field(
- default=None,
- description="""Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""",
- )
- spec: Optional[ReasoningEngineSpec] = Field(
- default=None,
- description="""Optional. Configurations of the ReasoningEngine""",
- )
- update_time: Optional[datetime.datetime] = Field(
- default=None,
- description="""Output only. Timestamp when this ReasoningEngine was most recently updated.""",
+ default=None, description="""Name of the agent engine."""
)
+ config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""")
-class ReasoningEngineDict(TypedDict, total=False):
- """An agent engine."""
-
- create_time: Optional[datetime.datetime]
- """Output only. Timestamp when this ReasoningEngine was created."""
-
- description: Optional[str]
- """Optional. The description of the ReasoningEngine."""
-
- display_name: Optional[str]
- """Required. The display name of the ReasoningEngine."""
-
- etag: Optional[str]
- """Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens."""
+class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False):
+ """Parameters for updating agent engines."""
name: Optional[str]
- """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`"""
-
- spec: Optional[ReasoningEngineSpecDict]
- """Optional. Configurations of the ReasoningEngine"""
+ """Name of the agent engine."""
- update_time: Optional[datetime.datetime]
- """Output only. Timestamp when this ReasoningEngine was most recently updated."""
+ config: Optional[UpdateAgentEngineConfigDict]
+ """"""
-ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict]
+_UpdateAgentEngineRequestParametersOrDict = Union[
+ _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict
+]
-class AgentEngineOperation(_common.BaseModel):
- """Operation that has an agent engine as a response."""
+class PromptOptimizerVAPOConfig(_common.BaseModel):
+ """VAPO Prompt Optimizer Config."""
- name: Optional[str] = Field(
- default=None,
- description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
- )
- metadata: Optional[dict[str, Any]] = Field(
- default=None,
- description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
- )
- done: Optional[bool] = Field(
- default=None,
- description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
- )
- error: Optional[dict[str, Any]] = Field(
- default=None,
- description="""The error result of the operation in case of failure or cancellation.""",
- )
- response: Optional[ReasoningEngine] = Field(
- default=None, description="""The created Agent Engine."""
+ config_path: Optional[str] = Field(
+ default=None, description="""The gcs path to the config file."""
)
+ service_account: Optional[str] = Field(default=None, description="""""")
+ wait_for_completion: Optional[bool] = Field(default=True, description="""""")
-class AgentEngineOperationDict(TypedDict, total=False):
- """Operation that has an agent engine as a response."""
+class PromptOptimizerVAPOConfigDict(TypedDict, total=False):
+ """VAPO Prompt Optimizer Config."""
- name: Optional[str]
- """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
+ config_path: Optional[str]
+ """The gcs path to the config file."""
- metadata: Optional[dict[str, Any]]
- """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
+ service_account: Optional[str]
+ """"""
- done: Optional[bool]
- """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
+ wait_for_completion: Optional[bool]
+ """"""
- error: Optional[dict[str, Any]]
- """The error result of the operation in case of failure or cancellation."""
- response: Optional[ReasoningEngineDict]
- """The created Agent Engine."""
+PromptOptimizerVAPOConfigOrDict = Union[
+ PromptOptimizerVAPOConfig, PromptOptimizerVAPOConfigDict
+]
-AgentEngineOperationOrDict = Union[AgentEngineOperation, AgentEngineOperationDict]
+class PromptTemplate(_common.BaseModel):
+ """A prompt template for creating prompts with variables."""
+ text: Optional[str] = Field(
+ default=None, description="""The prompt template text."""
+ )
+ _VARIABLE_NAME_REGEX: ClassVar[str] = r"\{([_a-zA-Z][_a-zA-Z0-9]*)\}"
-class DeleteAgentEngineConfig(_common.BaseModel):
- """Config for deleting agent engine."""
-
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
- )
-
-
-class DeleteAgentEngineConfigDict(TypedDict, total=False):
- """Config for deleting agent engine."""
+ @field_validator("text")
+ @classmethod
+ def text_must_not_be_empty(cls, value: str) -> str:
+ if not value.strip():
+ raise ValueError(
+ "Prompt template text cannot be empty or consist only of" " whitespace."
+ )
+ return value
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
+ @computed_field
+ @property
+ def variables(self) -> set[str]:
+ return set(re.findall(self._VARIABLE_NAME_REGEX, self.text))
+ def _split_template_by_variables(self) -> list[Tuple[str, str]]:
+ parts = []
+ last_end = 0
+ for match in re.finditer(self._VARIABLE_NAME_REGEX, self.text):
+ start, end = match.span()
+ var_name = match.group(1)
+ if start > last_end:
+ parts.append(("text", self.text[last_end:start]))
+ parts.append(("var", var_name))
+ last_end = end
+ if last_end < len(self.text):
+ parts.append(("text", self.text[last_end:]))
+ return parts
-DeleteAgentEngineConfigOrDict = Union[
- DeleteAgentEngineConfig, DeleteAgentEngineConfigDict
-]
+ def _merge_adjacent_text_parts(
+ self, parts: list[genai_types.Part]
+ ) -> list[genai_types.Part]:
+ if not parts:
+ return []
+ merged = []
+ current_text_buffer = []
-class _DeleteAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for deleting agent engines."""
+ for part in parts:
+ is_purely_text = part.text is not None and all(
+ getattr(part, field) is None
+ for field in part.model_fields
+ if field != "text"
+ )
- name: Optional[str] = Field(
- default=None, description="""Name of the agent engine."""
- )
- force: Optional[bool] = Field(
- default=False,
- description="""If set to true, any child resources will also be deleted.""",
- )
- config: Optional[DeleteAgentEngineConfig] = Field(default=None, description="""""")
+ if is_purely_text:
+ current_text_buffer.append(part.text)
+ else:
+ if current_text_buffer:
+ merged.append(genai_types.Part(text="".join(current_text_buffer)))
+ current_text_buffer = []
+ merged.append(part)
+ if current_text_buffer:
+ merged.append(genai_types.Part(text="".join(current_text_buffer)))
-class _DeleteAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for deleting agent engines."""
+ return merged
- name: Optional[str]
- """Name of the agent engine."""
+ def _is_multimodal_json_string(
+ self,
+ value: Any,
+ ) -> bool:
+ """Checks if the input value is a multimodal JSON string."""
+ if not isinstance(value, str):
+ return False
+ try:
+ data = json.loads(value)
+ # Check for the specific structure: {"contents": [{"parts": [...]}]}
+ # or {"parts": [...]} if assemble returns a single Content JSON
+ if isinstance(data, dict):
+ if "contents" in data and isinstance(data["contents"], list):
+ if not data["contents"]:
+ return False
+ first_content = data["contents"][0]
+ if isinstance(first_content, dict) and "parts" in first_content:
+ try:
+ genai_types.Content.model_validate(first_content)
+ return True
+ except ValueError:
+ return False
+ # Adding a check if 'data' itself is a Content-like object with parts
+ elif "parts" in data and isinstance(data["parts"], list):
+ try:
+ genai_types.Content.model_validate(data)
+ return True
+ except ValueError:
+ return False
+ return False
+ except json.JSONDecodeError:
+ return False
- force: Optional[bool]
- """If set to true, any child resources will also be deleted."""
+ def _parse_multimodal_json_string_into_parts(
+ self,
+ value: str,
+ ) -> list[genai_types.Part]:
+ """Parses a multimodal JSON string and returns its list of Parts."""
+ try:
+ content = genai_types.Content.model_validate_json(value)
+ return content.parts
+ except Exception:
+ return [genai_types.Part(text=value)]
- config: Optional[DeleteAgentEngineConfigDict]
- """"""
+ def assemble(self, **kwargs: Any) -> str:
+ """Assembles the prompt template with the given keyword arguments.
+ Supports both text and multimodal content. The `assemble` method
+ substitutes variables from the prompt template text with provided
+ values.
-_DeleteAgentEngineRequestParametersOrDict = Union[
- _DeleteAgentEngineRequestParameters, _DeleteAgentEngineRequestParametersDict
-]
+ Key Behaviors of `assemble()`:
+ 1. Variable Substitution: Replaces all defined variables with their
+ corresponding keyword argument values. Raises ValueError if a
+ template
+ variable is missing a value or if an extraneous kwarg is provided.
+ 2. Multimodal Handling:
+ - Detects if any variable's value is a JSON string representing
+ multimodal content (specifically, `{"contents": [{"parts": [...]}]}`
+ or `{"role": "user", "parts": [...]}`).
+ - If multimodal content is detected for a variable, its `Part`
+ objects
+ are extracted and inserted into the assembled sequence.
+ - Text segments from the template and simple text variable values
+ become `Part(text=...)`.
+ 3. Output Format:
+ - If ALL substituted variables were simple text AND the assembled
+ result (after merging adjacent text parts) consists of a single,
+ purely textual `Part`, `assemble()` returns a raw Python string.
+ - Otherwise (if any variable was multimodal, or if the assembly
+ results in multiple parts or non-textual parts), `assemble()`
+ returns
+ a JSON string representing a single `google.genai.types.Content`
+ object with `role="user"` and the assembled parts.
+ 4. Text Part Merging: Consecutively assembled text parts are
+ automatically merged into a single text `Part` to create a more
+ concise list of parts.
+ This dual output format (raw string or JSON string of `Content`) allows
+ the downstream inference functions to seamlessly handle both simple text
+ prompts and more complex multimodal prompts generated from the same
+ templating mechanism.
+ """
+ current_variables = self.variables
+ for var_name_in_kwarg in kwargs:
+ if var_name_in_kwarg not in current_variables:
+ raise ValueError(
+ f"Invalid variable name '{var_name_in_kwarg}' provided to"
+ " assemble. Valid variables in template are:"
+ f" {current_variables}"
+ )
+ # Check if all template variables are provided in kwargs
+ for tpl_var in current_variables:
+ if tpl_var not in kwargs:
+ raise ValueError(f"Missing value for template variable '{tpl_var}'.")
-class DeleteAgentEngineOperation(_common.BaseModel):
- """Operation for deleting agent engines."""
+ template_segments = self._split_template_by_variables()
- name: Optional[str] = Field(
- default=None,
- description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
- )
- metadata: Optional[dict[str, Any]] = Field(
- default=None,
- description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
- )
- done: Optional[bool] = Field(
- default=None,
- description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
- )
- error: Optional[dict[str, Any]] = Field(
- default=None,
- description="""The error result of the operation in case of failure or cancellation.""",
- )
+ raw_assembled_parts: list[genai_types.Part] = []
+ contains_multimodal_variable_type = False
+ for segment_type, segment_value in template_segments:
+ if segment_type == "text":
+ if segment_value:
+ raw_assembled_parts.append(genai_types.Part(text=segment_value))
+ elif segment_type == "var":
+ var_value = kwargs.get(segment_value)
-class DeleteAgentEngineOperationDict(TypedDict, total=False):
- """Operation for deleting agent engines."""
+ str_var_value = str(var_value)
- name: Optional[str]
- """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
+ if self._is_multimodal_json_string(str_var_value):
+ multimodal_parts = self._parse_multimodal_json_string_into_parts(
+ str_var_value
+ )
+ if multimodal_parts:
+ contains_multimodal_variable_type = True
+ raw_assembled_parts.extend(multimodal_parts)
+ else:
+ raw_assembled_parts.append(genai_types.Part(text=str_var_value))
+ else:
+ raw_assembled_parts.append(genai_types.Part(text=str_var_value))
- metadata: Optional[dict[str, Any]]
- """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
+ final_assembled_parts = self._merge_adjacent_text_parts(raw_assembled_parts)
- done: Optional[bool]
- """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
+ # Condition for returning raw text string:
+ # 1. No multimodal variable was *originally* a multimodal JSON string.
+ # 2. After merging, there's exactly one part.
+ # 3. That single part is purely textual.
+ if (
+ not contains_multimodal_variable_type
+ and len(final_assembled_parts) == 1
+ and final_assembled_parts[0].text is not None
+ and all(
+ getattr(final_assembled_parts[0], field) is None
+ for field in final_assembled_parts[0].model_fields
+ if field not in ["text", "role"]
+ )
+ ):
+ return final_assembled_parts[0].text
- error: Optional[dict[str, Any]]
- """The error result of the operation in case of failure or cancellation."""
+ # Otherwise, construct a Content object (as JSON string).
+ final_content_obj = genai_types.Content(parts=final_assembled_parts)
+ return final_content_obj.model_dump_json(exclude_none=True)
+ def __str__(self) -> str:
+ return self.text
-DeleteAgentEngineOperationOrDict = Union[
- DeleteAgentEngineOperation, DeleteAgentEngineOperationDict
-]
+ def __repr__(self) -> str:
+ return f"PromptTemplate(text='{self.text}')"
-class GetAgentEngineConfig(_common.BaseModel):
- """Config for create agent engine."""
+class MetricPromptBuilder(PromptTemplate):
+ """Builder class for structured LLM-based metric prompt template."""
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
+ criteria: Optional[dict[str, str]] = Field(
+ None,
+ description="""A dictionary of criteria used to evaluate the model responses.
+ The keys are criterion names, and the values are the corresponding
+ criterion definitions.
+ """,
)
-
-class GetAgentEngineConfigDict(TypedDict, total=False):
- """Config for create agent engine."""
-
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
-
-
-GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict]
-
-
-class _GetAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for getting agent engines."""
-
- name: Optional[str] = Field(
- default=None, description="""Name of the agent engine."""
+ rating_scores: Optional[dict[str, str]] = Field(
+ None,
+ description="""A dictionary mapping of rating score names to their definitions.""",
)
- config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""")
-
-
-class _GetAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for getting agent engines."""
-
- name: Optional[str]
- """Name of the agent engine."""
-
- config: Optional[GetAgentEngineConfigDict]
- """"""
-
-_GetAgentEngineRequestParametersOrDict = Union[
- _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict
-]
-
-
-class ListAgentEngineConfig(_common.BaseModel):
- """Config for listing agent engines."""
+ @staticmethod
+ def _get_default_instruction() -> str:
+ """Returns the default instruction for evaluation."""
+ return (
+ "You are an expert evaluator. Your task is to evaluate the quality"
+ " of the responses generated by AI models. We will provide you with"
+ " the user prompt and an AI-generated responses.\nYou should first"
+ " read the user input carefully for analyzing the task, and then"
+ " evaluate the quality of the responses based on the Criteria"
+ " provided in the Evaluation section below.\nYou will assign the"
+ " response a rating following the Rating Scores and Evaluation"
+ " Steps. Give step by step explanations for your rating, and only"
+ " choose ratings from the Rating Scores."
+ )
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
- )
- page_size: Optional[int] = Field(default=None, description="""""")
- page_token: Optional[str] = Field(default=None, description="""""")
- filter: Optional[str] = Field(
- default=None,
- description="""An expression for filtering the results of the request.
- For field names both snake_case and camelCase are supported.""",
+ instruction: Optional[str] = Field(
+ default_factory=lambda: MetricPromptBuilder._get_default_instruction(),
+ description="""The general instruction to guide the model in performing the evaluation.
+ If not provided, a default instruction for evaluation will be used.
+ """,
)
+ metric_definition: Optional[str] = Field(
+ None,
+ description="""An optional high-level description of the metric to be evaluated.
+ If not provided, this field will not be included in the prompt template.
+ """,
+ )
-class ListAgentEngineConfigDict(TypedDict, total=False):
- """Config for listing agent engines."""
-
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
-
- page_size: Optional[int]
- """"""
-
- page_token: Optional[str]
- """"""
-
- filter: Optional[str]
- """An expression for filtering the results of the request.
- For field names both snake_case and camelCase are supported."""
-
-
-ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict]
-
-
-class _ListAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for listing agent engines."""
-
- config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""")
-
-
-class _ListAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for listing agent engines."""
-
- config: Optional[ListAgentEngineConfigDict]
- """"""
-
-
-_ListAgentEngineRequestParametersOrDict = Union[
- _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict
-]
-
+ @staticmethod
+ def _get_default_evaluation_steps() -> dict[str, str]:
+ """Returns the default evaluation steps for metric evaluation."""
+ return {
+ "Step 1": (
+ "Assess the response in aspects of all criteria provided."
+ " Provide assessment according to each criterion."
+ ),
+ "Step 2": (
+ "Score based on the Rating Scores. Give a brief rationale to"
+ " explain your evaluation considering each individual"
+ " criterion."
+ ),
+ }
-class ListReasoningEnginesResponse(_common.BaseModel):
- """Response for listing agent engines."""
+ evaluation_steps: Optional[dict[str, str]] = Field(
+ default_factory=lambda: MetricPromptBuilder._get_default_evaluation_steps(),
+ description="""An optional dictionary of evaluation steps.
+ The keys are the names of the evaluation steps, and the values are
+ descriptions of the corresponding evaluation steps. If not provided,
+ default metric evaluation steps will be used.
+ """,
+ )
- next_page_token: Optional[str] = Field(default=None, description="""""")
- reasoning_engines: Optional[list[ReasoningEngine]] = Field(
- default=None,
- description="""List of agent engines.
+ few_shot_examples: Optional[list[str]] = Field(
+ None,
+ description="""An optional list of few-shot examples to guide the model's evaluation.
+ These examples demonstrate how to apply the criteria, rating scores,
+ and evaluation steps to assess model responses. Providing few-shot examples
+ can improve the accuracy of the evaluation. If not provided, this field
+ will not be included in the prompt template.
""",
)
+ @staticmethod
+ def _serialize_dict_in_order(elements: Optional[dict[str, str]]) -> str:
+ """Serializes dictionary to ordered string value without brackets."""
+ if elements is None:
+ return ""
+ return "\n".join(f"{key}: {value}" for key, value in sorted(elements.items()))
-class ListReasoningEnginesResponseDict(TypedDict, total=False):
- """Response for listing agent engines."""
-
- next_page_token: Optional[str]
- """"""
-
- reasoning_engines: Optional[list[ReasoningEngineDict]]
- """List of agent engines.
- """
-
+ @model_validator(mode="before")
+ @classmethod
+ def _prepare_fields_and_construct_text(cls, data: Any) -> Any:
+ """Pydantic model validator (before mode) to prepare and construct prompt text.
-ListReasoningEnginesResponseOrDict = Union[
- ListReasoningEnginesResponse, ListReasoningEnginesResponseDict
-]
+ This validator performs the following actions:
+ 1. Apply default logic for fields (instruction, evaluation_steps).
+ 2. Construct the 'text' string from all components.
+ 3. Ensure 'text' is in the data dictionary for PromptTemplate
+ initialization.
+ Args:
+ data: Input data for the model, either a dictionary or an existing
+ model instance.
-class GetAgentEngineOperationConfig(_common.BaseModel):
+ Returns:
+ Processed data dictionary with the 'text' field constructed.
+ """
+ if not isinstance(data, dict):
+ return data
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
- )
+ if "text" in data:
+ raise ValueError(
+ "The 'text' field is automatically constructed and should not"
+ " be provided manually."
+ )
+ if data.get("criteria") is None or data.get("rating_scores") is None:
+ raise ValueError(
+ "Both 'criteria' and 'rating_scores' are required to construct"
+ " theLLM-based metric prompt template text."
+ )
-class GetAgentEngineOperationConfigDict(TypedDict, total=False):
+ instruction = data.get("instruction", cls._get_default_instruction())
+ metric_definition = data.get("metric_definition")
+ evaluation_steps = data.get(
+ "evaluation_steps", cls._get_default_evaluation_steps()
+ )
+ criteria = data.get("criteria")
+ rating_scores = data.get("rating_scores")
+ few_shot_examples = data.get("few_shot_examples")
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
+ template_parts = [
+ "# Instruction",
+ instruction,
+ "\n",
+ "# Evaluation",
+ ]
+ sections = {
+ "Metric Definition": metric_definition,
+ "Criteria": cls._serialize_dict_in_order(criteria),
+ "Rating Scores": cls._serialize_dict_in_order(rating_scores),
+ "Evaluation Steps": cls._serialize_dict_in_order(evaluation_steps),
+ "Evaluation Examples": (
+ "\n".join(few_shot_examples) if few_shot_examples else None
+ ),
+ }
-GetAgentEngineOperationConfigOrDict = Union[
- GetAgentEngineOperationConfig, GetAgentEngineOperationConfigDict
-]
+ for title, content in sections.items():
+ if content:
+ template_parts.extend([f"## {title}", f"{content}\n"])
+ template_parts.extend(
+ [
+ "\n",
+ "# User Inputs and AI-generated Response",
+ "## User Prompt",
+ "{prompt}",
+ "\n",
+ "## AI-generated Response",
+ "{response}",
+ ]
+ )
-class _GetAgentEngineOperationParameters(_common.BaseModel):
- """Parameters for the GET method."""
+ constructed_text = "\n".join(template_parts)
- operation_name: Optional[str] = Field(
- default=None,
- description="""The server-assigned name for the operation.""",
- )
- config: Optional[GetAgentEngineOperationConfig] = Field(
- default=None,
- description="""Used to override the default configuration.""",
- )
+ data["text"] = constructed_text
+ return data
+ def __str__(self) -> str:
+ """Returns the fully constructed prompt template text."""
+ return self.text
-class _GetAgentEngineOperationParametersDict(TypedDict, total=False):
- """Parameters for the GET method."""
- operation_name: Optional[str]
- """The server-assigned name for the operation."""
+class PromptTemplateDict(TypedDict, total=False):
+ """A prompt template for creating prompts with variables."""
- config: Optional[GetAgentEngineOperationConfigDict]
- """Used to override the default configuration."""
+ text: Optional[str]
+ """The prompt template text."""
-_GetAgentEngineOperationParametersOrDict = Union[
- _GetAgentEngineOperationParameters, _GetAgentEngineOperationParametersDict
-]
+PromptTemplateOrDict = Union[PromptTemplate, PromptTemplateDict]
-class QueryAgentEngineConfig(_common.BaseModel):
- """Config for querying agent engines."""
+class EvalRunInferenceConfig(_common.BaseModel):
+ """Optional parameters for inference."""
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
+ dest: Optional[str] = Field(
+ default=None,
+ description="""The destination path for the inference results.""",
)
- class_method: Optional[str] = Field(
- default=None, description="""The class method to call."""
+ prompt_template: Optional[Union[str, PromptTemplate]] = Field(
+ default=None,
+ description="""The prompt template to use for inference.""",
)
- input: Optional[dict[str, Any]] = Field(
- default=None, description="""The input to the class method."""
+ generate_content_config: Optional[genai_types.GenerateContentConfig] = Field(
+ default=None,
+ description="""The config for the generate content call.""",
)
- include_all_fields: Optional[bool] = Field(default=False, description="""""")
-
-
-class QueryAgentEngineConfigDict(TypedDict, total=False):
- """Config for querying agent engines."""
-
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
-
- class_method: Optional[str]
- """The class method to call."""
-
- input: Optional[dict[str, Any]]
- """The input to the class method."""
-
- include_all_fields: Optional[bool]
- """"""
-
-QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict]
-
-
-class _QueryAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for querying agent engines."""
-
- name: Optional[str] = Field(
- default=None, description="""Name of the agent engine."""
- )
- config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""")
+class EvalRunInferenceConfigDict(TypedDict, total=False):
+ """Optional parameters for inference."""
-class _QueryAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for querying agent engines."""
+ dest: Optional[str]
+ """The destination path for the inference results."""
- name: Optional[str]
- """Name of the agent engine."""
+ prompt_template: Optional[Union[str, PromptTemplateDict]]
+ """The prompt template to use for inference."""
- config: Optional[QueryAgentEngineConfigDict]
- """"""
+ generate_content_config: Optional[genai_types.GenerateContentConfig]
+ """The config for the generate content call."""
-_QueryAgentEngineRequestParametersOrDict = Union[
- _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict
-]
+EvalRunInferenceConfigOrDict = Union[EvalRunInferenceConfig, EvalRunInferenceConfigDict]
-class QueryReasoningEngineResponse(_common.BaseModel):
- """The response for querying an agent engine."""
+class Metric(_common.BaseModel):
+ """The metric used for evaluation."""
- output: Optional[Any] = Field(
+ name: Optional[str] = Field(default=None, description="""The name of the metric.""")
+ custom_function: Optional[Callable] = Field(
default=None,
- description="""Response provided by users in JSON object format.""",
+ description="""The custom function that defines the end-to-end logic for metric computation.""",
)
-
-
-class QueryReasoningEngineResponseDict(TypedDict, total=False):
- """The response for querying an agent engine."""
-
- output: Optional[Any]
- """Response provided by users in JSON object format."""
-
-
-QueryReasoningEngineResponseOrDict = Union[
- QueryReasoningEngineResponse, QueryReasoningEngineResponseDict
-]
-
-
-class UpdateAgentEngineConfig(_common.BaseModel):
- """Config for updating agent engine."""
-
- http_options: Optional[HttpOptions] = Field(
- default=None, description="""Used to override HTTP request options."""
+ prompt_template: Optional[str] = Field(
+ default=None, description="""The prompt template for the metric."""
)
- display_name: Optional[str] = Field(
+ judge_model: Optional[str] = Field(
+ default=None, description="""The judge model for the metric."""
+ )
+ judge_model_sampling_count: Optional[int] = Field(
+ default=None, description="""The sampling count for the judge model."""
+ )
+ judge_model_system_instruction: Optional[str] = Field(
default=None,
- description="""The user-defined name of the Agent Engine.
-
- The display name can be up to 128 characters long and can comprise any
- UTF-8 characters.
- """,
+ description="""The system instruction for the judge model.""",
)
- description: Optional[str] = Field(
- default=None, description="""The description of the Agent Engine."""
+ return_raw_output: Optional[bool] = Field(
+ default=None,
+ description="""Whether to return the raw output from the judge model.""",
)
- spec: Optional[ReasoningEngineSpec] = Field(
+ parse_and_reduce_fn: Optional[Callable] = Field(
default=None,
- description="""Optional. Configurations of the ReasoningEngine.""",
+ description="""The parse and reduce function for the judge model.""",
)
- update_mask: Optional[str] = Field(
+ aggregate_summary_fn: Optional[Callable] = Field(
default=None,
- description="""The update mask to apply. For the `FieldMask` definition, see
- https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask.""",
+ description="""The aggregate summary function for the judge model.""",
)
+ # Allow extra fields to support metric-specific config fields.
+ model_config = ConfigDict(extra="allow")
+
+ _is_predefined: bool = PrivateAttr(default=False)
+ """A boolean indicating whether the metric is predefined."""
-class UpdateAgentEngineConfigDict(TypedDict, total=False):
- """Config for updating agent engine."""
+ _config_source: Optional[str] = PrivateAttr(default=None)
+ """An optional string indicating the source of the metric configuration."""
+
+ _version: Optional[str] = PrivateAttr(default=None)
+ """An optional string indicating the version of the metric."""
- http_options: Optional[HttpOptionsDict]
- """Used to override HTTP request options."""
+ @model_validator(mode="after")
+ @classmethod
+ def validate_name(cls, model: "Metric") -> "Metric":
+ if not model.name:
+ raise ValueError("Metric name cannot be empty.")
+ model.name = model.name.lower()
+ return model
- display_name: Optional[str]
- """The user-defined name of the Agent Engine.
+ def to_yaml_file(self, file_path: str, version: Optional[str] = None) -> None:
+ """Dumps the metric object to a YAML file.
- The display name can be up to 128 characters long and can comprise any
- UTF-8 characters.
- """
+ Args:
+ file_path: The path to the YAML file.
+ version: Optional version string to include in the YAML output.
- description: Optional[str]
- """The description of the Agent Engine."""
+ Raises:
+ ImportError: If the pyyaml library is not installed.
+ """
+ if yaml is None:
+ raise ImportError(
+ "YAML serialization requires the pyyaml library. Please install"
+ " it using 'pip install google-cloud-aiplatform[evaluation]'."
+ )
- spec: Optional[ReasoningEngineSpecDict]
- """Optional. Configurations of the ReasoningEngine."""
+ fields_to_exclude_callables = set()
+ for field_name, field_info in self.model_fields.items():
+ annotation = field_info.annotation
+ origin = typing.get_origin(annotation)
- update_mask: Optional[str]
- """The update mask to apply. For the `FieldMask` definition, see
- https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
+ is_field_callable_type = False
+ if annotation is Callable or origin is Callable:
+ is_field_callable_type = True
+ elif origin is Union:
+ args = typing.get_args(annotation)
+ if any(
+ arg is Callable or typing.get_origin(arg) is Callable
+ for arg in args
+ ):
+ is_field_callable_type = True
+ if is_field_callable_type:
+ fields_to_exclude_callables.add(field_name)
-UpdateAgentEngineConfigOrDict = Union[
- UpdateAgentEngineConfig, UpdateAgentEngineConfigDict
-]
+ data_to_dump = self.model_dump(
+ exclude_unset=True,
+ exclude_none=True,
+ mode="json",
+ exclude=(
+ fields_to_exclude_callables if fields_to_exclude_callables else None
+ ),
+ )
+ if version:
+ data_to_dump["version"] = version
-class _UpdateAgentEngineRequestParameters(_common.BaseModel):
- """Parameters for updating agent engines."""
+ with open(file_path, "w", encoding="utf-8") as f:
+ yaml.dump(data_to_dump, f, sort_keys=False, allow_unicode=True)
- name: Optional[str] = Field(
- default=None, description="""Name of the agent engine."""
- )
- config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""")
+class LLMMetric(Metric):
+ """A metric that uses LLM-as-a-judge for evaluation."""
-class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False):
- """Parameters for updating agent engines."""
+ @field_validator("prompt_template", mode="before")
+ @classmethod
+ def validate_prompt_template(cls, value: Union[str, "MetricPromptBuilder"]) -> str:
+ """Validates prompt template to be a non-empty string."""
+ if value is None:
+ raise ValueError("Prompt template cannot be empty.")
+ if isinstance(value, MetricPromptBuilder):
+ value = str(value)
+ if not value.strip():
+ raise ValueError("Prompt template cannot be an empty string.")
+ return value
- name: Optional[str]
- """Name of the agent engine."""
+ @field_validator("judge_model_sampling_count")
+ @classmethod
+ def validate_judge_model_sampling_count(cls, value: Optional[int]) -> Optional[int]:
+ """Validates judge_model_sampling_count to be between 1 and 32."""
+ if value is not None and (value < 1 or value > 32):
+ raise ValueError("judge_model_sampling_count must be between 1 and 32.")
+ return value
- config: Optional[UpdateAgentEngineConfigDict]
- """"""
+class MetricDict(TypedDict, total=False):
+ """The metric used for evaluation."""
-_UpdateAgentEngineRequestParametersOrDict = Union[
- _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict
-]
+ name: Optional[str]
+ """The name of the metric."""
+ custom_function: Optional[Callable]
+ """The custom function that defines the end-to-end logic for metric computation."""
-class PromptOptimizerVAPOConfig(_common.BaseModel):
- """VAPO Prompt Optimizer Config."""
+ prompt_template: Optional[str]
+ """The prompt template for the metric."""
- config_path: Optional[str] = Field(
- default=None, description="""The gcs path to the config file."""
- )
- service_account: Optional[str] = Field(default=None, description="""""")
- wait_for_completion: Optional[bool] = Field(default=True, description="""""")
+ judge_model: Optional[str]
+ """The judge model for the metric."""
+ judge_model_sampling_count: Optional[int]
+ """The sampling count for the judge model."""
-class PromptOptimizerVAPOConfigDict(TypedDict, total=False):
- """VAPO Prompt Optimizer Config."""
+ judge_model_system_instruction: Optional[str]
+ """The system instruction for the judge model."""
- config_path: Optional[str]
- """The gcs path to the config file."""
+ return_raw_output: Optional[bool]
+ """Whether to return the raw output from the judge model."""
- service_account: Optional[str]
- """"""
+ parse_and_reduce_fn: Optional[Callable]
+ """The parse and reduce function for the judge model."""
- wait_for_completion: Optional[bool]
- """"""
+ aggregate_summary_fn: Optional[Callable]
+ """The aggregate summary function for the judge model."""
-PromptOptimizerVAPOConfigOrDict = Union[
- PromptOptimizerVAPOConfig, PromptOptimizerVAPOConfigDict
-]
+MetricOrDict = Union[Metric, MetricDict]
-class PromptTemplate(_common.BaseModel):
- """A prompt template for creating prompts with variables."""
+class Message(_common.BaseModel):
+ """Represents a single message turn in a conversation."""
- text: Optional[str] = Field(
- default=None, description="""The prompt template text."""
+ turn_id: Optional[str] = Field(
+ default=None, description="""Unique identifier for the message turn."""
+ )
+ content: Optional[genai_types.Content] = Field(
+ default=None,
+ description="""Content of the message, including function call.""",
+ )
+ creation_timestamp: Optional[datetime.datetime] = Field(
+ default=None,
+ description="""Timestamp indicating when the message was created.""",
+ )
+ author: Optional[str] = Field(
+ default=None,
+ description="""Name of the entity that produced the message.""",
)
- _VARIABLE_NAME_REGEX: ClassVar[str] = r"\{([_a-zA-Z][_a-zA-Z0-9]*)\}"
- @field_validator("text")
- @classmethod
- def text_must_not_be_empty(cls, value: str) -> str:
- if not value.strip():
- raise ValueError(
- "Prompt template text cannot be empty or consist only of" " whitespace."
- )
- return value
- @computed_field
- @property
- def variables(self) -> set[str]:
- return set(re.findall(self._VARIABLE_NAME_REGEX, self.text))
+class MessageDict(TypedDict, total=False):
+ """Represents a single message turn in a conversation."""
- def _split_template_by_variables(self) -> list[Tuple[str, str]]:
- parts = []
- last_end = 0
- for match in re.finditer(self._VARIABLE_NAME_REGEX, self.text):
- start, end = match.span()
- var_name = match.group(1)
- if start > last_end:
- parts.append(("text", self.text[last_end:start]))
- parts.append(("var", var_name))
- last_end = end
- if last_end < len(self.text):
- parts.append(("text", self.text[last_end:]))
- return parts
+ turn_id: Optional[str]
+ """Unique identifier for the message turn."""
- def _merge_adjacent_text_parts(
- self, parts: list[genai_types.Part]
- ) -> list[genai_types.Part]:
- if not parts:
- return []
+ content: Optional[genai_types.Content]
+ """Content of the message, including function call."""
- merged = []
- current_text_buffer = []
+ creation_timestamp: Optional[datetime.datetime]
+ """Timestamp indicating when the message was created."""
- for part in parts:
- is_purely_text = part.text is not None and all(
- getattr(part, field) is None
- for field in part.model_fields
- if field != "text"
- )
+ author: Optional[str]
+ """Name of the entity that produced the message."""
- if is_purely_text:
- current_text_buffer.append(part.text)
- else:
- if current_text_buffer:
- merged.append(genai_types.Part(text="".join(current_text_buffer)))
- current_text_buffer = []
- merged.append(part)
- if current_text_buffer:
- merged.append(genai_types.Part(text="".join(current_text_buffer)))
+MessageOrDict = Union[Message, MessageDict]
- return merged
- def _is_multimodal_json_string(
- self,
- value: Any,
- ) -> bool:
- """Checks if the input value is a multimodal JSON string."""
- if not isinstance(value, str):
- return False
- try:
- data = json.loads(value)
- # Check for the specific structure: {"contents": [{"parts": [...]}]}
- # or {"parts": [...]} if assemble returns a single Content JSON
- if isinstance(data, dict):
- if "contents" in data and isinstance(data["contents"], list):
- if not data["contents"]:
- return False
- first_content = data["contents"][0]
- if isinstance(first_content, dict) and "parts" in first_content:
- try:
- genai_types.Content.model_validate(first_content)
- return True
- except ValueError:
- return False
- # Adding a check if 'data' itself is a Content-like object with parts
- elif "parts" in data and isinstance(data["parts"], list):
- try:
- genai_types.Content.model_validate(data)
- return True
- except ValueError:
- return False
- return False
- except json.JSONDecodeError:
- return False
+class AgentData(_common.BaseModel):
+ """Container for all agent-specific data."""
+
+ tool_use_trajectory: Optional[list[Message]] = Field(
+ default=None,
+ description="""Tool use trajectory in chronological order.""",
+ )
+ intermediate_responses: Optional[list[Message]] = Field(
+ default=None,
+ description="""Intermediate responses generated by sub-agents to convey progress or status in a multi-agent system, distinct from the final response.""",
+ )
- def _parse_multimodal_json_string_into_parts(
- self,
- value: str,
- ) -> list[genai_types.Part]:
- """Parses a multimodal JSON string and returns its list of Parts."""
- try:
- content = genai_types.Content.model_validate_json(value)
- return content.parts
- except Exception:
- return [genai_types.Part(text=value)]
- def assemble(self, **kwargs: Any) -> str:
- """Assembles the prompt template with the given keyword arguments.
+class AgentDataDict(TypedDict, total=False):
+ """Container for all agent-specific data."""
- Supports both text and multimodal content. The `assemble` method
- substitutes variables from the prompt template text with provided
- values.
+ tool_use_trajectory: Optional[list[MessageDict]]
+ """Tool use trajectory in chronological order."""
- Key Behaviors of `assemble()`:
- 1. Variable Substitution: Replaces all defined variables with their
- corresponding keyword argument values. Raises ValueError if a
- template
- variable is missing a value or if an extraneous kwarg is provided.
- 2. Multimodal Handling:
- - Detects if any variable's value is a JSON string representing
- multimodal content (specifically, `{"contents": [{"parts": [...]}]}`
- or `{"role": "user", "parts": [...]}`).
- - If multimodal content is detected for a variable, its `Part`
- objects
- are extracted and inserted into the assembled sequence.
- - Text segments from the template and simple text variable values
- become `Part(text=...)`.
- 3. Output Format:
- - If ALL substituted variables were simple text AND the assembled
- result (after merging adjacent text parts) consists of a single,
- purely textual `Part`, `assemble()` returns a raw Python string.
- - Otherwise (if any variable was multimodal, or if the assembly
- results in multiple parts or non-textual parts), `assemble()`
- returns
- a JSON string representing a single `google.genai.types.Content`
- object with `role="user"` and the assembled parts.
- 4. Text Part Merging: Consecutively assembled text parts are
- automatically merged into a single text `Part` to create a more
- concise list of parts.
+ intermediate_responses: Optional[list[MessageDict]]
+ """Intermediate responses generated by sub-agents to convey progress or status in a multi-agent system, distinct from the final response."""
- This dual output format (raw string or JSON string of `Content`) allows
- the downstream inference functions to seamlessly handle both simple text
- prompts and more complex multimodal prompts generated from the same
- templating mechanism.
- """
- current_variables = self.variables
- for var_name_in_kwarg in kwargs:
- if var_name_in_kwarg not in current_variables:
- raise ValueError(
- f"Invalid variable name '{var_name_in_kwarg}' provided to"
- " assemble. Valid variables in template are:"
- f" {current_variables}"
- )
- # Check if all template variables are provided in kwargs
- for tpl_var in current_variables:
- if tpl_var not in kwargs:
- raise ValueError(f"Missing value for template variable '{tpl_var}'.")
- template_segments = self._split_template_by_variables()
+AgentDataOrDict = Union[AgentData, AgentDataDict]
- raw_assembled_parts: list[genai_types.Part] = []
- contains_multimodal_variable_type = False
- for segment_type, segment_value in template_segments:
- if segment_type == "text":
- if segment_value:
- raw_assembled_parts.append(genai_types.Part(text=segment_value))
- elif segment_type == "var":
- var_value = kwargs.get(segment_value)
+class ResponseCandidate(_common.BaseModel):
+ """A model-generated content to the prompt."""
- str_var_value = str(var_value)
+ response: Optional[genai_types.Content] = Field(
+ default=None,
+ description="""The final model-generated response to the `prompt`.""",
+ )
+ agent_data: Optional[AgentData] = Field(
+ default=None,
+ description="""Agent-specific data including tool use trajectory and intermediate responses for more complex interactions.""",
+ )
- if self._is_multimodal_json_string(str_var_value):
- multimodal_parts = self._parse_multimodal_json_string_into_parts(
- str_var_value
- )
- if multimodal_parts:
- contains_multimodal_variable_type = True
- raw_assembled_parts.extend(multimodal_parts)
- else:
- raw_assembled_parts.append(genai_types.Part(text=str_var_value))
- else:
- raw_assembled_parts.append(genai_types.Part(text=str_var_value))
- final_assembled_parts = self._merge_adjacent_text_parts(raw_assembled_parts)
+class ResponseCandidateDict(TypedDict, total=False):
+ """A model-generated content to the prompt."""
- # Condition for returning raw text string:
- # 1. No multimodal variable was *originally* a multimodal JSON string.
- # 2. After merging, there's exactly one part.
- # 3. That single part is purely textual.
- if (
- not contains_multimodal_variable_type
- and len(final_assembled_parts) == 1
- and final_assembled_parts[0].text is not None
- and all(
- getattr(final_assembled_parts[0], field) is None
- for field in final_assembled_parts[0].model_fields
- if field not in ["text", "role"]
- )
- ):
- return final_assembled_parts[0].text
+ response: Optional[genai_types.Content]
+ """The final model-generated response to the `prompt`."""
- # Otherwise, construct a Content object (as JSON string).
- final_content_obj = genai_types.Content(parts=final_assembled_parts)
- return final_content_obj.model_dump_json(exclude_none=True)
+ agent_data: Optional[AgentDataDict]
+ """Agent-specific data including tool use trajectory and intermediate responses for more complex interactions."""
- def __str__(self) -> str:
- return self.text
- def __repr__(self) -> str:
- return f"PromptTemplate(text='{self.text}')"
+ResponseCandidateOrDict = Union[ResponseCandidate, ResponseCandidateDict]
-class MetricPromptBuilder(PromptTemplate):
- """Builder class for structured LLM-based metric prompt template."""
+class EvalCase(_common.BaseModel):
+ """A comprehensive representation of a GenAI interaction for evaluation."""
- criteria: Optional[dict[str, str]] = Field(
- None,
- description="""A dictionary of criteria used to evaluate the model responses.
- The keys are criterion names, and the values are the corresponding
- criterion definitions.
- """,
+ prompt: Optional[genai_types.Content] = Field(
+ default=None,
+ description="""The most recent user message (current input).""",
)
-
- rating_scores: Optional[dict[str, str]] = Field(
- None,
- description="""A dictionary mapping of rating score names to their definitions.""",
+ responses: Optional[list[ResponseCandidate]] = Field(
+ default=None,
+ description="""Model-generated replies to the last user message. Multiple responses are allowed to support use cases such as comparing different model outputs.""",
+ )
+ reference: Optional[ResponseCandidate] = Field(
+ default=None,
+ description="""User-provided, golden reference model reply to prompt in context of chat history.""",
+ )
+ system_instruction: Optional[genai_types.Content] = Field(
+ default=None, description="""System instruction for the model."""
+ )
+ conversation_history: Optional[list[Message]] = Field(
+ default=None,
+ description="""List of all prior messages in the conversation (chat history).""",
+ )
+ eval_case_id: Optional[str] = Field(
+ default=None,
+ description="""Unique identifier for the evaluation case.""",
)
+ # Allow extra fields to support custom metric prompts and stay backward compatible.
+ model_config = ConfigDict(frozen=True, extra="allow")
- @staticmethod
- def _get_default_instruction() -> str:
- """Returns the default instruction for evaluation."""
- return (
- "You are an expert evaluator. Your task is to evaluate the quality"
- " of the responses generated by AI models. We will provide you with"
- " the user prompt and an AI-generated responses.\nYou should first"
- " read the user input carefully for analyzing the task, and then"
- " evaluate the quality of the responses based on the Criteria"
- " provided in the Evaluation section below.\nYou will assign the"
- " response a rating following the Rating Scores and Evaluation"
- " Steps. Give step by step explanations for your rating, and only"
- " choose ratings from the Rating Scores."
- )
- instruction: Optional[str] = Field(
- default_factory=lambda: MetricPromptBuilder._get_default_instruction(),
- description="""The general instruction to guide the model in performing the evaluation.
- If not provided, a default instruction for evaluation will be used.
- """,
- )
+class EvalCaseDict(TypedDict, total=False):
+ """A comprehensive representation of a GenAI interaction for evaluation."""
- metric_definition: Optional[str] = Field(
- None,
- description="""An optional high-level description of the metric to be evaluated.
- If not provided, this field will not be included in the prompt template.
- """,
- )
+ prompt: Optional[genai_types.Content]
+ """The most recent user message (current input)."""
- @staticmethod
- def _get_default_evaluation_steps() -> dict[str, str]:
- """Returns the default evaluation steps for metric evaluation."""
- return {
- "Step 1": (
- "Assess the response in aspects of all criteria provided."
- " Provide assessment according to each criterion."
- ),
- "Step 2": (
- "Score based on the Rating Scores. Give a brief rationale to"
- " explain your evaluation considering each individual"
- " criterion."
- ),
- }
+ responses: Optional[list[ResponseCandidateDict]]
+ """Model-generated replies to the last user message. Multiple responses are allowed to support use cases such as comparing different model outputs."""
- evaluation_steps: Optional[dict[str, str]] = Field(
- default_factory=lambda: MetricPromptBuilder._get_default_evaluation_steps(),
- description="""An optional dictionary of evaluation steps.
- The keys are the names of the evaluation steps, and the values are
- descriptions of the corresponding evaluation steps. If not provided,
- default metric evaluation steps will be used.
- """,
- )
+ reference: Optional[ResponseCandidateDict]
+ """User-provided, golden reference model reply to prompt in context of chat history."""
- few_shot_examples: Optional[list[str]] = Field(
- None,
- description="""An optional list of few-shot examples to guide the model's evaluation.
- These examples demonstrate how to apply the criteria, rating scores,
- and evaluation steps to assess model responses. Providing few-shot examples
- can improve the accuracy of the evaluation. If not provided, this field
- will not be included in the prompt template.
- """,
- )
+ system_instruction: Optional[genai_types.Content]
+ """System instruction for the model."""
- @staticmethod
- def _serialize_dict_in_order(elements: Optional[dict[str, str]]) -> str:
- """Serializes dictionary to ordered string value without brackets."""
- if elements is None:
- return ""
- return "\n".join(f"{key}: {value}" for key, value in sorted(elements.items()))
+ conversation_history: Optional[list[MessageDict]]
+ """List of all prior messages in the conversation (chat history)."""
- @model_validator(mode="before")
- @classmethod
- def _prepare_fields_and_construct_text(cls, data: Any) -> Any:
- """Pydantic model validator (before mode) to prepare and construct prompt text.
+ eval_case_id: Optional[str]
+ """Unique identifier for the evaluation case."""
- This validator performs the following actions:
- 1. Apply default logic for fields (instruction, evaluation_steps).
- 2. Construct the 'text' string from all components.
- 3. Ensure 'text' is in the data dictionary for PromptTemplate
- initialization.
- Args:
- data: Input data for the model, either a dictionary or an existing
- model instance.
+EvalCaseOrDict = Union[EvalCase, EvalCaseDict]
+
- Returns:
- Processed data dictionary with the 'text' field constructed.
- """
- if not isinstance(data, dict):
- return data
+class GcsSource(_common.BaseModel):
+ """Cloud storage source holds the dataset.
- if "text" in data:
- raise ValueError(
- "The 'text' field is automatically constructed and should not"
- " be provided manually."
- )
+ Currently only one Cloud Storage file path is supported.
+ """
- if data.get("criteria") is None or data.get("rating_scores") is None:
- raise ValueError(
- "Both 'criteria' and 'rating_scores' are required to construct"
- " theLLM-based metric prompt template text."
- )
+ uris: Optional[list[str]] = Field(
+ default=None,
+ description="""Required. Google Cloud Storage URI(-s) to the input file(s). May contain wildcards. For more information on wildcards, see https://cloud.google.com/storage/docs/wildcards.""",
+ )
- instruction = data.get("instruction", cls._get_default_instruction())
- metric_definition = data.get("metric_definition")
- evaluation_steps = data.get(
- "evaluation_steps", cls._get_default_evaluation_steps()
- )
- criteria = data.get("criteria")
- rating_scores = data.get("rating_scores")
- few_shot_examples = data.get("few_shot_examples")
- template_parts = [
- "# Instruction",
- instruction,
- "\n",
- "# Evaluation",
- ]
+class GcsSourceDict(TypedDict, total=False):
+ """Cloud storage source holds the dataset.
- sections = {
- "Metric Definition": metric_definition,
- "Criteria": cls._serialize_dict_in_order(criteria),
- "Rating Scores": cls._serialize_dict_in_order(rating_scores),
- "Evaluation Steps": cls._serialize_dict_in_order(evaluation_steps),
- "Evaluation Examples": (
- "\n".join(few_shot_examples) if few_shot_examples else None
- ),
- }
+ Currently only one Cloud Storage file path is supported.
+ """
- for title, content in sections.items():
- if content:
- template_parts.extend([f"## {title}", f"{content}\n"])
+ uris: Optional[list[str]]
+ """Required. Google Cloud Storage URI(-s) to the input file(s). May contain wildcards. For more information on wildcards, see https://cloud.google.com/storage/docs/wildcards."""
- template_parts.extend(
- [
- "\n",
- "# User Inputs and AI-generated Response",
- "## User Prompt",
- "{prompt}",
- "\n",
- "## AI-generated Response",
- "{response}",
- ]
- )
- constructed_text = "\n".join(template_parts)
+GcsSourceOrDict = Union[GcsSource, GcsSourceDict]
- data["text"] = constructed_text
- return data
- def __str__(self) -> str:
- """Returns the fully constructed prompt template text."""
- return self.text
+class BigQuerySource(_common.BaseModel):
+ """The BigQuery location for the input content."""
+ input_uri: Optional[str] = Field(
+ default=None,
+ description="""Required. BigQuery URI to a table, up to 2000 characters long. Accepted forms: * BigQuery path. For example: `bq://projectId.bqDatasetId.bqTableId`.""",
+ )
-class PromptTemplateDict(TypedDict, total=False):
- """A prompt template for creating prompts with variables."""
- text: Optional[str]
- """The prompt template text."""
+class BigQuerySourceDict(TypedDict, total=False):
+ """The BigQuery location for the input content."""
+ input_uri: Optional[str]
+ """Required. BigQuery URI to a table, up to 2000 characters long. Accepted forms: * BigQuery path. For example: `bq://projectId.bqDatasetId.bqTableId`."""
-PromptTemplateOrDict = Union[PromptTemplate, PromptTemplateDict]
+BigQuerySourceOrDict = Union[BigQuerySource, BigQuerySourceDict]
-class EvalRunInferenceConfig(_common.BaseModel):
- """Optional parameters for inference."""
- dest: Optional[str] = Field(
+class EvaluationDataset(_common.BaseModel):
+ """The dataset used for evaluation."""
+
+ eval_cases: Optional[list[EvalCase]] = Field(
+ default=None, description="""The evaluation cases to be evaluated."""
+ )
+ eval_dataset_df: Optional["pd.DataFrame"] = Field(
default=None,
- description="""The destination path for the inference results.""",
+ description="""The evaluation dataset in the form of a Pandas DataFrame.""",
)
- prompt_template: Optional[Union[str, PromptTemplate]] = Field(
+ candidate_name: Optional[str] = Field(
default=None,
- description="""The prompt template to use for inference.""",
+ description="""The name of the candidate model or agent for this evaluation dataset.""",
)
- generate_content_config: Optional[genai_types.GenerateContentConfig] = Field(
+ gcs_source: Optional[GcsSource] = Field(
default=None,
- description="""The config for the generate content call.""",
+ description="""The GCS source for the evaluation dataset.""",
+ )
+ bigquery_source: Optional[BigQuerySource] = Field(
+ default=None,
+ description="""The BigQuery source for the evaluation dataset.""",
)
+ def show(self) -> None:
+ """Shows the evaluation dataset."""
+ from . import _evals_visualization
-class EvalRunInferenceConfigDict(TypedDict, total=False):
- """Optional parameters for inference."""
+ _evals_visualization.display_evaluation_dataset(self)
- dest: Optional[str]
- """The destination path for the inference results."""
- prompt_template: Optional[Union[str, PromptTemplateDict]]
- """The prompt template to use for inference."""
+class EvaluationDatasetDict(TypedDict, total=False):
+ """The dataset used for evaluation."""
- generate_content_config: Optional[genai_types.GenerateContentConfig]
- """The config for the generate content call."""
+ eval_cases: Optional[list[EvalCaseDict]]
+ """The evaluation cases to be evaluated."""
+ eval_dataset_df: Optional["pd.DataFrame"]
+ """The evaluation dataset in the form of a Pandas DataFrame."""
-EvalRunInferenceConfigOrDict = Union[EvalRunInferenceConfig, EvalRunInferenceConfigDict]
+ candidate_name: Optional[str]
+ """The name of the candidate model or agent for this evaluation dataset."""
+
+ gcs_source: Optional[GcsSourceDict]
+ """The GCS source for the evaluation dataset."""
+
+ bigquery_source: Optional[BigQuerySourceDict]
+ """The BigQuery source for the evaluation dataset."""
+
+
+EvaluationDatasetOrDict = Union[EvaluationDataset, EvaluationDatasetDict]
class WinRateStats(_common.BaseModel):
@@ -4870,6 +4759,123 @@ class EvaluateMethodConfigDict(TypedDict, total=False):
EvaluateMethodConfigOrDict = Union[EvaluateMethodConfig, EvaluateMethodConfigDict]
+class OutputConfig(_common.BaseModel):
+ """Config for evaluation output."""
+
+ gcs_destination: Optional[GcsDestination] = Field(
+ default=None,
+ description="""Cloud storage destination for evaluation output.""",
+ )
+
+
+class OutputConfigDict(TypedDict, total=False):
+ """Config for evaluation output."""
+
+ gcs_destination: Optional[GcsDestinationDict]
+ """Cloud storage destination for evaluation output."""
+
+
+OutputConfigOrDict = Union[OutputConfig, OutputConfigDict]
+
+
+class EvaluateDatasetConfig(_common.BaseModel):
+ """Config for evaluate instances."""
+
+ http_options: Optional[HttpOptions] = Field(
+ default=None, description="""Used to override HTTP request options."""
+ )
+
+
+class EvaluateDatasetConfigDict(TypedDict, total=False):
+ """Config for evaluate instances."""
+
+ http_options: Optional[HttpOptionsDict]
+ """Used to override HTTP request options."""
+
+
+EvaluateDatasetConfigOrDict = Union[EvaluateDatasetConfig, EvaluateDatasetConfigDict]
+
+
+class EvaluateDatasetRequestParameters(_common.BaseModel):
+ """Parameters for batch dataset evaluation."""
+
+ dataset: Optional[EvaluationDataset] = Field(default=None, description="""""")
+ metrics: Optional[list[Metric]] = Field(default=None, description="""""")
+ output_config: Optional[OutputConfig] = Field(default=None, description="""""")
+ autorater_config: Optional[AutoraterConfig] = Field(
+ default=None, description=""""""
+ )
+ config: Optional[EvaluateDatasetConfig] = Field(default=None, description="""""")
+
+
+class EvaluateDatasetRequestParametersDict(TypedDict, total=False):
+ """Parameters for batch dataset evaluation."""
+
+ dataset: Optional[EvaluationDatasetDict]
+ """"""
+
+ metrics: Optional[list[MetricDict]]
+ """"""
+
+ output_config: Optional[OutputConfigDict]
+ """"""
+
+ autorater_config: Optional[AutoraterConfigDict]
+ """"""
+
+ config: Optional[EvaluateDatasetConfigDict]
+ """"""
+
+
+EvaluateDatasetRequestParametersOrDict = Union[
+ EvaluateDatasetRequestParameters, EvaluateDatasetRequestParametersDict
+]
+
+
+class EvaluateDatasetOperation(_common.BaseModel):
+
+ name: Optional[str] = Field(
+ default=None,
+ description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""",
+ )
+ metadata: Optional[dict[str, Any]] = Field(
+ default=None,
+ description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""",
+ )
+ done: Optional[bool] = Field(
+ default=None,
+ description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""",
+ )
+ error: Optional[dict[str, Any]] = Field(
+ default=None,
+ description="""The error result of the operation in case of failure or cancellation.""",
+ )
+ response: Optional[EvaluationDataset] = Field(default=None, description="""""")
+
+
+class EvaluateDatasetOperationDict(TypedDict, total=False):
+
+ name: Optional[str]
+ """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`."""
+
+ metadata: Optional[dict[str, Any]]
+ """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any."""
+
+ done: Optional[bool]
+ """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available."""
+
+ error: Optional[dict[str, Any]]
+ """The error result of the operation in case of failure or cancellation."""
+
+ response: Optional[EvaluationDatasetDict]
+ """"""
+
+
+EvaluateDatasetOperationOrDict = Union[
+ EvaluateDatasetOperation, EvaluateDatasetOperationDict
+]
+
+
class AgentEngine(_common.BaseModel):
"""An agent engine instance."""