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.
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.
