VertexAI Tracing

Instrument LLM calls made with the Vertex AI SDK (e.g., Gemini models) using OpenInference and view traces in Arize.

The Vertex AI SDK, particularly for models like Gemini, can be instrumented using the openinference-instrumentation-vertexai package. Traces can be sent to Arize for observability.

Launch Arize

To get started, sign up for a free Arize account and get your Space ID and API Key.

Install

pip install openinference-instrumentation-vertexai google-cloud-aiplatform arize-otel

Note: google-cloud-aiplatform includes the vertexai SDK.

Environment and Authentication Setup

Before running your code, ensure your Google Cloud environment is authenticated and configured for Vertex AI. This typically involves:

Authenticating via the Google Cloud CLI: gcloud auth application-default login
Setting your Google Cloud Project ID. You can do this via an environment variable:
```
export GOOGLE_CLOUD_PROJECT='your-gcp-project-id'
```
The vertexai.init(project='your-gcp-project-id', location='your-region') call in your Python code will also use this project ID if not set by other means.

Refer to Google's official Vertex AI documentation for the most current and detailed setup instructions.

Setup Tracing

Use the register function to connect your application to Arize and instrument the Vertex AI client.

# Import open-telemetry dependencies
from arize.otel import register

# Setup OTel via our convenience function
tracer_provider = register(
    space_id = "your-space-id", # in app space settings page
    api_key = "your-api-key", # in app space settings page
    project_name = "my-vertexai-app", # Or your preferred project name
)

# Import the instrumentor from OpenInference
from openinference.instrumentation.vertexai import VertexAIInstrumentor

# Instrument the Vertex AI client
VertexAIInstrumentor().instrument(tracer_provider=tracer_provider)

Run VertexAI

Initialize Vertex AI and use a generative model. Ensure your project and location are correctly configured.

import vertexai
from vertexai.generative_models import GenerativeModel

# Initialize Vertex AI. Project ID can be picked from GOOGLE_CLOUD_PROJECT env var
# or explicitly passed: vertexai.init(project="your-gcp-project-id", location="us-central1")
vertexai.init(location="us-central1") # Ensure your project is implicitly set or passed here

# Load a Gemini model
model = GenerativeModel("gemini-1.5-flash-001") # Or gemini-1.0-pro, etc.

response = model.generate_content("Why is the sky blue?")

# Accessing the text content from the response
if response and response.candidates and response.candidates[0].content.parts:
    print(response.candidates[0].content.parts[0].text)
elif hasattr(response, 'text'): # Fallback for older response structures or simpler responses
    print(response.text)
else:
    print("Could not extract text from response or response is empty.")
    print(f"Full response: {response}")

Observe

Now that you have tracing setup, all instrumented calls to Vertex AI models (like Gemini generate_content) will be streamed to your Arize account for observability and evaluation.

Resources

Last updated 1 month ago

Was this helpful?