Google Vertex AI

Python auto-instrumentation for Google Vertex AI

Python auto-instrumentation library for VertexAI SDK and the Google Cloud AI Platform. Traces are fully OpenTelemetry compatible and can be sent to any OpenTelemetry collector for viewing, such as Arize Phoenix.

Installation

pip install -Uqqq openinference-instrumentation-vertexai

Quickstart

Install dependencies:

pip install -U \
    openinference-instrumentation-vertexai \
    arize-phoenix \
    opentelemetry-sdk \
    opentelemetry-exporter-otlp \
    "opentelemetry-proto>=1.12.0"

Start Phoenix server

Start the phoenix server so that it is ready to receive traces. The phoenix server runs entirely on your machine and does not send data over the internet.

python -m phoenix.server.main serve

Setup instrumentation

Instrumenting vertexai is simple:

from openinference.instrumentation.vertexai import VertexAIInstrumentor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor

endpoint = "http://127.0.0.1:4317"
tracer_provider = TracerProvider()
tracer_provider.add_span_processor(SimpleSpanProcessor(OTLPSpanExporter(endpoint)))

VertexAIInstrumentor().instrument(tracer_provider=tracer_provider)

Use Vertex AI

Now, all calls by generative_models are instrumented and can be viewed in the phoenix UI:

import vertexai
from vertexai.generative_models import GenerativeModel

vertexai.init(location="us-central1")
model = GenerativeModel("gemini-1.5-flash")

print(model.generate_content("Why is sky blue?"))

Visit the Phoenix app at http://localhost:6006 to see your traces.

More Info

• OpenInference and Phoenix documentation
• How to customize spans to track sessions, metadata, etc.
• How to account for private information and span payload customization
• PyPI package