[bot] Add Vertex AI Python SDK (google-cloud-aiplatform) integration for GenerativeModel, TextEmbeddingModel, and Agent Builder instrumentation

[braintrust-bot]

Summary

The Vertex AI Python SDK (google-cloud-aiplatform, imported as vertexai) is Google Cloud's enterprise Python SDK for generative AI. It is actively maintained with frequent weekly releases (v1.159.0 as of July 2026) and provides GenerativeModel.generate_content() for chat/text generation, TextEmbeddingModel.get_embeddings() for embeddings, and Vertex AI Agent Builder via vertexai.preview.reasoning_engines for agent execution. This repository has zero instrumentation for any google-cloud-aiplatform execution surface — no integration directory, no wrapper, no patcher, no auto_instrument() support.

This is a distinct package from google-generativeai/google-genai, which is already instrumented as GoogleGenAIIntegration. The existing integration covers the google-genai SDK (including its optional Vertex AI backend mode via genai.Client(vertexai=True)), but google-cloud-aiplatform is a completely separate Python package with different APIs, different authentication (Google Cloud IAM/ADC vs. Google API keys), and a different class hierarchy.

Enterprise Google Cloud users who pip install google-cloud-aiplatform and use vertexai.generative_models.GenerativeModel directly get zero Braintrust spans today.

Comparable provider SDKs with dedicated integrations in this repo: anthropic, openai, mistral, cohere, google_genai, huggingface_hub, bedrock_runtime.

What needs to be instrumented

The google-cloud-aiplatform package exposes these execution surfaces, none of which are instrumented:

Generative AI — text and chat generation (highest priority)

Class / Method	Description	Return type
GenerativeModel.generate_content(contents, ...)	Synchronous text/chat generation — core execution surface	GenerationResponse
GenerativeModel.generate_content_async(contents, ...)	Async variant	Coroutine[GenerationResponse]
GenerativeModel.generate_content(..., stream=True)	Streaming generation	Iterable[GenerationResponse]
ChatSession.send_message(content, ...)	Multi-turn chat, maintains history	GenerationResponse
ChatSession.send_message_async(content, ...)	Async multi-turn chat	Coroutine[GenerationResponse]

GenerationResponse includes text (generated content), usage_metadata (prompt_token_count, candidates_token_count, total_token_count), candidates[0].finish_reason, and candidates[0].safety_ratings. Token counts are directly available for span metrics.

Embeddings

Class / Method	Description	Return type
TextEmbeddingModel.get_embeddings(texts, ...)	Generate dense vector embeddings for a list of texts	list[TextEmbedding]
TextEmbeddingModel.get_embeddings_async(texts, ...)	Async embeddings	Coroutine[list[TextEmbedding]]

TextEmbedding includes values (embedding vector), statistics.token_count. The model name is available from the instantiated TextEmbeddingModel.

Vertex AI Agent Builder (agent execution)

Class / Method	Description	Return type
reasoning_engine.query(input, ...)	Execute a deployed Vertex AI Agent Builder agent	varies
AdkApp.query(input, ...)	Execute an Agent Development Kit (ADK) agent deployed on Vertex	varies

vertexai.preview.reasoning_engines.ReasoningEngine.query() and .query_async() are the primary entry points for running agents deployed via Vertex AI Agent Builder. Each call may invoke multiple LLM calls and tool uses internally.

Key distinction from existing google_genai integration

	google-genai (COVERED)	google-cloud-aiplatform (NOT COVERED)
PyPI package	google-generativeai or google-genai	google-cloud-aiplatform
Import	import google.genai	import vertexai or from google.cloud import aiplatform
Auth	Google API key or Vertex AI via vertexai=True flag	Google Cloud ADC / service account / IAM
Primary use case	Google AI Studio, rapid prototyping	Enterprise GCP, Vertex AI platform
Client class	genai.Client()	vertexai.generative_models.GenerativeModel(model_name)
Installation	pip install google-generativeai	pip install google-cloud-aiplatform

The existing google_genai integration's vertexai flag check (in tracing.py) handles serialization differences for google-genai used with vertexai=True — it does NOT instrument google-cloud-aiplatform.

No coverage in any instrumentation layer

• No integration directory (py/src/braintrust/integrations/vertex_ai/)
• No wrapper function (e.g. wrap_vertex_ai())
• No patcher in any existing integration
• No nox test session (test_vertex_ai)
• No version entry in py/src/braintrust/integrations/versioning.py
• No mention in py/src/braintrust/integrations/__init__.py

A grep for google.cloud.aiplatform, from vertexai, and import vertexai across py/src/braintrust/ returns zero matches. The only occurrences of "vertex" in the SDK source (google_genai/tracing.py:171,175) are handling serialization for the google-genai SDK's Vertex AI mode — not the google-cloud-aiplatform package.

Braintrust docs status

not_found — google-cloud-aiplatform / Vertex AI Python SDK is not listed on the Braintrust integrations directory or the tracing guide. The documented "Gemini" integration maps to the existing google_genai integration. There is no auto_instrument() reference for vertexai, no wrap_vertex_ai() function, and no Vertex AI SDK setup documentation anywhere in Braintrust docs.

Upstream references

• google-cloud-aiplatform on PyPI: https://pypi.org/project/google-cloud-aiplatform/ (v1.159.0)
• google-cloud-aiplatform on GitHub: https://github.com/googleapis/python-aiplatform
• Vertex AI Python SDK quickstart: https://cloud.google.com/vertex-ai/generative-ai/docs/start/quickstarts/quickstart-multimodal
• GenerativeModel API reference: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/python/latest/vertexai.generative_models.GenerativeModel
• TextEmbeddingModel reference: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/python/latest/vertexai.language_models.TextEmbeddingModel
• Vertex AI Agent Builder / Reasoning Engines: https://cloud.google.com/vertex-ai/generative-ai/docs/reasoning-engine/overview
• Vertex AI Generative AI docs: https://cloud.google.com/vertex-ai/generative-ai/docs/overview

Local repo files inspected

• py/src/braintrust/integrations/ — no vertex_ai/ directory exists on main
• py/src/braintrust/wrappers/ — no Vertex AI wrapper
• py/noxfile.py — no test_vertex_ai session
• py/pyproject.toml [tool.braintrust.matrix] — no google-cloud-aiplatform entry
• py/src/braintrust/integrations/__init__.py — Vertex AI not listed in integration registry
• py/src/braintrust/integrations/versioning.py — no Vertex AI version matrix
• py/src/braintrust/integrations/google_genai/tracing.py — vertexai flag handling is for google-genai SDK only
• Full repo grep for google.cloud.aiplatform, from vertexai, import vertexai — zero matches in SDK source