Skip to content

googleapis/python-aiplatform

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3,591 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Google Gemini Enterprise Agent Platform SDK for Python

GA pypi versions unit-tests system-tests sample-tests

> Note: The Gemini Enterprise Agent Platform was formerly known as > Vertex AI.

> Note: The agent_engines module will be rebranded as runtimes in > an upcoming major release (not before 7/31/2026). See [Agent Runtime](https://docs.cloud.google.com/gemini-enterprise-agent-platform/scale#agent-runtime) for details.

Gemini Enterprise Agent Platform: Gemini Enterprise Agent Platform is Google Cloud's comprehensive platform for developers to build, scale, govern and optimize agents.

Installation

pip install google-cloud-aiplatform

With uv:

uv pip install google-cloud-aiplatform

Generative AI in the Gemini Enterprise Agent Platform SDK

To use Gen AI features from the Gemini Enterprise Agent Platform SDK, you can instantiate a Vertex SDK client with the following:

import vertexai
from vertexai import types

# Instantiate GenAI client from Vertex SDK
# Replace with your project ID and location
client = vertexai.Client(project='my-project', location='us-central1')

See the examples below for guidance on how to use specific features supported by the Vertex SDK client.

Gen AI Evaluation

To run evaluation, first generate model responses from a set of prompts.

import pandas as pd

prompts_df = pd.DataFrame({
    "prompt": [
        "What is the capital of France?",
        "Write a haiku about a cat.",
        "Write a Python function to calculate the factorial of a number.",
        "Translate 'How are you?' to French.",
    ],

    "reference": [
        "Paris",
        "Sunbeam on the floor,\nA furry puddle sleeping,\nTwitching tail tells tales.",
        "def factorial(n):\n    if n < 0:\n        return 'Factorial does not exist for negative numbers'\n    elif n == 0:\n        return 1\n    else:\n        fact = 1\n        i = 1\n        while i <= n:\n            fact *= i\n            i += 1\n        return fact",
        "Comment ça va ?",
    ]
})

inference_results = client.evals.run_inference(
    model="gemini-2.5-flash-preview-05-20",
    src=prompts_df
)

Then run evaluation by providing the inference results and specifying the metric types.

eval_result = client.evals.evaluate(
    dataset=inference_results,
    metrics=[
        types.Metric(name='exact_match'),
        types.Metric(name='rouge_l_sum'),
        types.RubricMetric.TEXT_QUALITY,
    ]
)

Agent Engine with Agent Development Kit (ADK)

First, define a function that looks up the exchange rate:

def get_exchange_rate(
    currency_from: str = "USD",
    currency_to: str = "EUR",
    currency_date: str = "latest",
):
    """Retrieves the exchange rate between two currencies on a specified date.

    Uses the Frankfurter API (https://api.frankfurter.app/) to obtain
    exchange rate data.

    Returns:
        dict: A dictionary containing the exchange rate information.
            Example: {"amount": 1.0, "base": "USD", "date": "2023-11-24",
                "rates": {"EUR": 0.95534}}
    """
    import requests
    response = requests.get(
        f"https://api.frankfurter.app/{currency_date}",
        params={"from": currency_from, "to": currency_to},
    )
    return response.json()

Next, define an ADK Agent:

from google.adk.agents import Agent
from vertexai.agent_engines import AdkApp

app = AdkApp(agent=Agent(
    model="gemini-3.5-flash",        # Required.
    name='currency_exchange_agent',  # Required.
    tools=[get_exchange_rate],       # Optional.
))

Test the agent locally using US dollars and Swedish Krona:

async for event in app.async_stream_query(
    user_id="user-id",
    message="What is the exchange rate from US dollars to SEK today?",
):
    print(event)

To deploy the agent to Agent Engine:

# Replace with client.runtimes.create in a future major release.
remote_app = client.agent_engines.create(
    agent=app,
    config={
        "requirements": ["google-cloud-aiplatform[agent_engines,adk]"],
    },
)

You can also run queries against the deployed agent:

async for event in remote_app.async_stream_query(
    user_id="user-id",
    message="What is the exchange rate from US dollars to SEK today?",
):
    print(event)

Prompt optimization

To do a zero-shot prompt optimization, use the optimize method.

prompt = "Generate system instructions for a question-answering assistant"
response = client.prompts.optimize(prompt=prompt)
print(response.raw_text_response)
if response.parsed_response:
  print(response.parsed_response.suggested_prompt)

To call the data-driven prompt optimization, call the launch_optimization_job method. In this case however, we need to provide a VAPO (Gemini Enterprise Agent Platform Prompt Optimizer) config. This config needs to have either service account or project number and the config path. Please refer to this [tutorial](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/prompts/data-driven-optimizer) for more details on config parameter.

from vertexai import types

project_number = PROJECT_NUMBER # replace with your project number
service_account = f"{project_number}-compute@developer.gserviceaccount.com"

vapo_config = vertexai.types.PromptOptimizerVAPOConfig(
    config_path="gs://your-bucket/config.json",
    service_account_project_number=project_number,
    wait_for_completion=False
)

# Set up logging to see the progress of the optimization job
logging.basicConfig(encoding='utf-8', level=logging.INFO, force=True)

result = client.prompts.launch_optimization_job(method=types.PromptOptimizerMethod.VAPO, config=vapo_config)

If you want to use the project number instead of the service account, you can instead use the following config:

vapo_config = vertexai.types.PromptOptimizerVAPOConfig(
    config_path="gs://your-bucket/config.json",
    service_account_project_number=project_number,
    wait_for_completion=False
)

We can also call the launch_optimization_job method asynchronously.

await client.aio.prompts.launch_optimization_job(method=types.PromptOptimizerMethod.VAPO, config=vapo_config)

Prompt Management

First define your prompt as a dictionary or types.Prompt object. Then call create_version.

prompt = {
    "prompt_data": {
        "contents": [{"parts": [{"text": "Hello, {name}! How are you?"}]}],
        "system_instruction": {"parts": [{"text": "Please answer in a short sentence."}]},
        "variables": [
            {"name": {"text": "Alice"}},
        ],
        "model": "gemini-2.5-flash",
    },
}

prompt_resource = client.prompts.create_version(
    prompt=prompt,
)

Note that you can also use the types.Prompt object to define your prompt. Some of the types used to do this are from the Gen AI SDK.

import types
from google.genai import types as genai_types

prompt = types.Prompt(
    prompt_data=types.PromptData(
      contents=[genai_types.Content(parts=[genai_types.Part(text="Hello, {name}! How are you?")])],
      system_instruction=genai_types.Content(parts=[genai_types.Part(text="Please answer in a short sentence.")]),
      variables=[
        {"name": genai_types.Part(text="Alice")},
      ],
      model="gemini-2.5-flash",
    ),
)

Retrieve a prompt by calling get() with the prompt_id.

retrieved_prompt = client.prompts.get(
    prompt_id=prompt_resource.prompt_id,
)

After creating or retrieving a prompt, you can call generate_content() with that prompt using the Gen AI SDK.

The following uses a utility function available on Prompt objects to transform a Prompt object into a list of Content objects for use with generate_content. To run this you need to have the Gen AI SDK installed, which you can do via pip install google-genai.

from google import genai
from google.genai import types as genai_types

# Create a Client in the Gen AI SDK
genai_client = genai.Client(vertexai=True, project="your-project", location="your-location")

# Call generate_content() with the prompt
response = genai_client.models.generate_content(
    model=retrieved_prompt.prompt_data.model,
    contents=retrieved_prompt.assemble_contents(),
)

Skill Registry

Create and manage skills in Skill Registry. You can optionally specify a custom string identifier using the skill_id configuration parameter.

# Create a skill
skill = client.skills.create(
    display_name="weather_skill",
    description="Retrieves the weather for a given location",
    config={
        "local_path": "./weather_skill_dir",
        "skill_id": "my-custom-weather-skill",
    },
)

Get an existing skill by its resource name.

fetched_skill = client.skills.get(name=skill.name)

Update an existing skill's metadata or underlying implementation.

# Update skill metadata
updated_skill = client.skills.update(
    name=skill.name,
    config={
        "display_name": "Updated Weather Skill",
        "description": "Provides localized current weather conditions and multi-day forecasts.",
    },
)

List all registered skills.

# List skills with custom page size
pager = client.skills.list(config={"page_size": 10})
for item in pager:
    print(item.name, item.display_name)

Search for skills semantically matched to a query.

# Retrieve skills matched to a semantic query
matched_skills = client.skills.retrieve(query="weather forecast")

List and view revisions for a skill using the ListSkillRevisions and GetSkillRevision API methods.

# List skill revisions
revisions_response = client.skills.revisions.list(name=skill.name)
for rev in revisions_response.skill_revisions:
    print(rev.name, rev.create_time)

# Get a specific skill revision by its resource name
if revisions_response.skill_revisions:
    target_revision_name = revisions_response.skill_revisions[0].name
    revision = client.skills.revisions.get(name=target_revision_name)

Delete a skill when it is no longer required.

# Delete a skill
client.skills.delete(name=skill.name)

Note

The following Generative AI modules in the Gemini Enterprise Agent Platform SDK are deprecated as of June 24, 2025 and will be removed on June 24, 2026: vertexai.generative_models, vertexai.language_models, vertexai.vision_models, vertexai.tuning, vertexai.caching. Please use the [Google Gen AI SDK](https://pypi.org/project/google-genai/) to access these features. See [the migration guide](https://cloud.google.com/vertex-ai/generative-ai/docs/deprecations/genai-vertexai-sdk) for details. You can continue using all other Gemini Enterprise Agent Platform SDK modules, as they are the recommended way to use the API.

Quick Start

In order to use this library, you first need to go through the following steps:

  1. Select or create a Cloud Platform project.
  2. Enable billing for your project.
  3. Enable the Gemini Enterprise Agent Platform API.
  4. Setup Authentication.

Supported Python Versions

Python >= 3.10

Deprecated Python Versions

Python <= 3.9.

The last version of this library compatible with Python 3.8 is google-cloud-aiplatform==1.90.0.

The last version of this library compatible with Python 3.7 is google-cloud-aiplatform==1.31.1.

The last version of this library compatible with Python 3.6 is google-cloud-aiplatform==1.12.1.

Overview

This section provides a brief overview of the Gemini Enterprise Agent Platform SDK for Python. You can also reference the notebooks in vertex-ai-samples for examples.

All publicly available SDK features can be found in the google/cloud/aiplatform directory. Under the hood, Vertex SDK builds on top of GAPIC, which stands for Google API CodeGen. The GAPIC library code sits in google/cloud/aiplatform_v1 and google/cloud/aiplatform_v1beta1, and it is auto-generated from Google's service proto files.

For most developers' programmatic needs, they can follow these steps to figure out which libraries to import:

  1. Look through google/cloud/aiplatform first -- Vertex SDK's APIs will almost always be easier to use and more concise comparing with GAPIC
  2. If the feature that you are looking for cannot be found there, look through aiplatform_v1 to see if it's available in GAPIC
  3. If it is still in beta phase, it will be available in aiplatform_v1beta1

If none of the above scenarios could help you find the right tools for your task, please feel free to open a github issue and send us a feature request.

Importing

Gemini Enterprise Agent Platform SDK resource based functionality can be used by importing the following namespace:

from google.cloud import aiplatform

Initialization

Initialize the SDK to store common configurations that you use with the SDK.

aiplatform.init(
    # your Google Cloud Project ID or number
    # environment default used is not set
    project='my-project',

    # the Gemini Enterprise Agent Platform region you will use
    # defaults to us-central1
    location='us-central1',

    # Google Cloud Storage bucket in same region as location
    # used to stage artifacts
    staging_bucket='gs://my_staging_bucket',

    # custom google.auth.credentials.Credentials
    # environment default credentials used if not set
    credentials=my_credentials,

    # customer managed encryption key resource name
    # will be applied to all Gemini Enterprise Agent Platform resources if set
    encryption_spec_key_name=my_encryption_key_name,

    # the name of the experiment to use to track
    # logged metrics and parameters
    experiment='my-experiment',

    # description of the experiment above
    experiment_description='my experiment description'
)

Datasets

Gemini Enterprise Agent Platform provides managed tabular, text, image, and video datasets. In the SDK, datasets can be used downstream to train models.

To create a tabular dataset:

my_dataset = aiplatform.TabularDataset.create(
    display_name="my-dataset", gcs_source=['gs://path/to/my/dataset.csv'])

You can also create and import a dataset in separate steps:

from google.cloud import aiplatform

my_dataset = aiplatform.TextDataset.create(
    display_name="my-dataset")

my_dataset.import_data(
    gcs_source=['gs://path/to/my/dataset.csv'],
    import_schema_uri=aiplatform.schema.dataset.ioformat.text.multi_label_classification
)

To get a previously created Dataset:

dataset = aiplatform.ImageDataset('projects/my-project/location/us-central1/datasets/{DATASET_ID}')

Gemini Enterprise Agent Platform supports a variety of dataset schemas. References to these schemas are available under the aiplatform.schema.dataset namespace. For more information on the supported dataset schemas please refer to the Preparing data docs.

Training

The Gemini Enterprise Agent Platform SDK for Python allows you train Custom and AutoML Models.

You can train custom models using a custom Python script, custom Python package, or container.

Preparing Your Custom Code

Gemini Enterprise Agent Platform custom training enables you to train on Gemini Enterprise Agent Platform datasets and produce Gemini Enterprise Agent Platform models. To do so your script must adhere to the following contract:

It must read datasets from the environment variables populated by the training service:

os.environ['AIP_DATA_FORMAT']  # provides format of data
os.environ['AIP_TRAINING_DATA_URI']  # uri to training split
os.environ['AIP_VALIDATION_DATA_URI']  # uri to validation split
os.environ['AIP_TEST_DATA_URI']  # uri to test split

Please visit Using a managed dataset in a custom training application for a detailed overview.

It must write the model artifact to the environment variable populated by the training service:

os.environ['AIP_MODEL_DIR']

Running Training

job = aiplatform.CustomTrainingJob(
    display_name="my-training-job",
    script_path="training_script.py",
    container_uri="us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-2:latest",
    requirements=["gcsfs==0.7.1"],
    model_serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest",
)

model = job.run(my_dataset,
                replica_count=1,
                machine_type="n1-standard-4",
                accelerator_type='NVIDIA_TESLA_K80',
                accelerator_count=1)

In the code block above my_dataset is managed dataset created in the Dataset section above. The model variable is a managed Gemini Enterprise Agent Platform model that can be deployed or exported.

AutoMLs

The Gemini Enterprise Agent Platform SDK for Python supports AutoML tabular, image, text, video, and forecasting.

To train an AutoML tabular model:

dataset = aiplatform.TabularDataset('projects/my-project/location/us-central1/datasets/{DATASET_ID}')

job = aiplatform.AutoMLTabularTrainingJob(
  display_name="train-automl",
  optimization_prediction_type="regression",
  optimization_objective="minimize-rmse",
)

model = job.run(
    dataset=dataset,
    target_column="target_column_name",
    training_fraction_split=0.6,
    validation_fraction_split=0.2,
    test_fraction_split=0.2,
    budget_milli_node_hours=1000,
    model_display_name="my-automl-model",
    disable_early_stopping=False,
)

Models

To get a model:

model = aiplatform.Model('/projects/my-project/locations/us-central1/models/{MODEL_ID}')

To upload a model:

model = aiplatform.Model.upload(
    display_name='my-model',
    artifact_uri="gs://python/to/my/model/dir",
    serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-2:latest",
)

To deploy a model:

endpoint = model.deploy(machine_type="n1-standard-4",
                        min_replica_count=1,
                        max_replica_count=5
                        machine_type='n1-standard-4',
                        accelerator_type='NVIDIA_TESLA_K80',
                        accelerator_count=1)

Please visit Importing models to Gemini Enterprise Agent Platform for a detailed overview:

Model Evaluation

The Gemini Enterprise Agent Platform SDK for Python currently supports getting model evaluation metrics for all AutoML models.

To list all model evaluations for a model:

model = aiplatform.Model('projects/my-project/locations/us-central1/models/{MODEL_ID}')

evaluations = model.list_model_evaluations()

To get the model evaluation resource for a given model:

model = aiplatform.Model('projects/my-project/locations/us-central1/models/{MODEL_ID}')

# returns the first evaluation with no arguments, you can also pass the evaluation ID
evaluation = model.get_model_evaluation()

eval_metrics = evaluation.metrics

You can also create a reference to your model evaluation directly by passing in the resource name of the model evaluation:

evaluation = aiplatform.ModelEvaluation(
  evaluation_name='projects/my-project/locations/us-central1/models/{MODEL_ID}/evaluations/{EVALUATION_ID}')

Alternatively, you can create a reference to your evaluation by passing in the model and evaluation IDs:

evaluation = aiplatform.ModelEvaluation(
  evaluation_name={EVALUATION_ID},
  model_id={MODEL_ID})

Batch Prediction

To create a batch prediction job:

model = aiplatform.Model('/projects/my-project/locations/us-central1/models/{MODEL_ID}')

batch_prediction_job = model.batch_predict(
  job_display_name='my-batch-prediction-job',
  instances_format='csv',
  machine_type='n1-standard-4',
  gcs_source=['gs://path/to/my/file.csv'],
  gcs_destination_prefix='gs://path/to/my/batch_prediction/results/',
  service_account='my-sa@my-project.iam.gserviceaccount.com'
)

You can also create a batch prediction job asynchronously by including the sync=False argument:

batch_prediction_job = model.batch_predict(..., sync=False)

# wait for resource to be created
batch_prediction_job.wait_for_resource_creation()

# get the state
batch_prediction_job.state

# block until job is complete
batch_prediction_job.wait()

Endpoints

To create an endpoint:

endpoint = aiplatform.Endpoint.create(display_name='my-endpoint')

To deploy a model to a created endpoint:

model = aiplatform.Model('/projects/my-project/locations/us-central1/models/{MODEL_ID}')

endpoint.deploy(model,
                min_replica_count=1,
                max_replica_count=5,
                machine_type='n1-standard-4',
                accelerator_type='NVIDIA_TESLA_K80',
                accelerator_count=1)

To get predictions from endpoints:

endpoint.predict(instances=[[6.7, 3.1, 4.7, 1.5], [4.6, 3.1, 1.5, 0.2]])

To undeploy models from an endpoint:

endpoint.undeploy_all()

To delete an endpoint:

endpoint.delete()

Pipelines

To create a Gemini Enterprise Agent Platform Pipeline run and monitor until completion:

# Instantiate PipelineJob object
pl = PipelineJob(
    display_name="My first pipeline",

    # Whether or not to enable caching
    # True = always cache pipeline step result
    # False = never cache pipeline step result
    # None = defer to cache option for each pipeline component in the pipeline definition
    enable_caching=False,

    # Local or GCS path to a compiled pipeline definition
    template_path="pipeline.json",

    # Dictionary containing input parameters for your pipeline
    parameter_values=parameter_values,

    # GCS path to act as the pipeline root
    pipeline_root=pipeline_root,
)

# Execute pipeline in Gemini Enterprise Agent Platform and monitor until completion
pl.run(
  # Email address of service account to use for the pipeline run
  # You must have iam.serviceAccounts.actAs permission on the service account to use it
  service_account=service_account,

  # Whether this function call should be synchronous (wait for pipeline run to finish before terminating)
  # or asynchronous (return immediately)
  sync=True
)

To create a Gemini Enterprise Agent Platform Pipeline without monitoring until completion, use submit instead of run:

# Instantiate PipelineJob object
pl = PipelineJob(
    display_name="My first pipeline",

    # Whether or not to enable caching
    # True = always cache pipeline step result
    # False = never cache pipeline step result
    # None = defer to cache option for each pipeline component in the pipeline definition
    enable_caching=False,

    # Local or GCS path to a compiled pipeline definition
    template_path="pipeline.json",

    # Dictionary containing input parameters for your pipeline
    parameter_values=parameter_values,

    # GCS path to act as the pipeline root
    pipeline_root=pipeline_root,
)

# Submit the Pipeline to Gemini Enterprise Agent Platform
pl.submit(
  # Email address of service account to use for the pipeline run
  # You must have iam.serviceAccounts.actAs permission on the service account to use it
  service_account=service_account,
)

Explainable AI: Get Metadata

To get metadata in dictionary format from TensorFlow 1 models:

from google.cloud.aiplatform.explain.metadata.tf.v1 import saved_model_metadata_builder

builder = saved_model_metadata_builder.SavedModelMetadataBuilder(
          'gs://python/to/my/model/dir', tags=[tf.saved_model.tag_constants.SERVING]
      )
generated_md = builder.get_metadata()

To get metadata in dictionary format from TensorFlow 2 models:

from google.cloud.aiplatform.explain.metadata.tf.v2 import saved_model_metadata_builder

builder = saved_model_metadata_builder.SavedModelMetadataBuilder('gs://python/to/my/model/dir')
generated_md = builder.get_metadata()

To use Explanation Metadata in endpoint deployment and model upload:

explanation_metadata = builder.get_metadata_protobuf()

# To deploy a model to an endpoint with explanation
model.deploy(..., explanation_metadata=explanation_metadata)

# To deploy a model to a created endpoint with explanation
endpoint.deploy(..., explanation_metadata=explanation_metadata)

# To upload a model with explanation
aiplatform.Model.upload(..., explanation_metadata=explanation_metadata)

Cloud Profiler

Cloud Profiler allows you to profile your remote Gemini Enterprise Agent Platform Training jobs on demand and visualize the results in Gemini Enterprise Agent Platform Tensorboard.

To start using the profiler with TensorFlow, update your training script to include the following:

from google.cloud.aiplatform.training_utils import cloud_profiler
...
cloud_profiler.init()

Next, run the job with with a Gemini Enterprise Agent Platform TensorBoard instance. For full details on how to do this, visit https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-overview

Finally, visit your TensorBoard in your Google Cloud Console, navigate to the "Profile" tab, and click the Capture Profile button. This will allow users to capture profiling statistics for the running jobs.

Next Steps

About

A Python SDK for Vertex AI, a fully managed, end-to-end platform for data science and machine learning.

Resources

License

Code of conduct

Contributing

Security policy

Stars

900 stars

Watchers

68 watching

Forks

Packages

 
 
 

Contributors

Languages