ContextQMD
Back to Opik

Using Opik with VertexAI

apps/opik-documentation/documentation/docs/cookbook/vertexai.ipynb

2.0.24-52623.0 KB
Original Source

Using Opik with VertexAI

Opik integrates with VertexAI to provide a simple way to log traces for all VertexAI LLM calls. This works for all the supported models.

Creating an account on Comet.com

Comet provides a hosted version of the Opik platform, simply create an account and grab your API Key.

You can also run the Opik platform locally, see the installation guide for more information.

python
%pip install --upgrade opik google-genai -q
python
import opik

opik.configure(use_local=False)

Preparing our environment

First, we will set up our Google GenAI client with VertexAI credentials.

python
from google import genai

PROJECT_ID = "[your-project-id]"
LOCATION = "us-central1"

if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    raise ValueError("Please set your PROJECT_ID")

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Logging traces

In order to log traces to Opik, we need to wrap calls made via VertexAI with the track_genai function:

python
import os
from opik.integrations.genai import track_genai

os.environ["OPIK_PROJECT_NAME"] = "vertexai-integration-demo"
vertexai_client = track_genai(client)


prompt = """
Write a short two sentence story about Opik.
"""

response = vertexai_client.models.generate_content(
    model="gemini-2.0-flash-001", contents=prompt
)
print(response.text)

The prompt and response messages are automatically logged to Opik and can be viewed in the UI.

Using it with the track decorator

If you have multiple steps in your LLM pipeline, you can use the track decorator to log the traces for each step. If Gemini model is called within one of these steps, the LLM call with be associated with that corresponding step:

python
from opik import track


@track
def generate_story(prompt):
    response = vertexai_client.models.generate_content(
        model="gemini-2.0-flash-001", contents=prompt
    )
    return response.text


@track
def generate_topic():
    prompt = "Generate a topic for a story about Opik."
    response = vertexai_client.models.generate_content(
        model="gemini-2.0-flash-001", contents=prompt
    )
    return response.text


@track
def generate_opik_story():
    topic = generate_topic()
    story = generate_story(topic)
    return story


generate_opik_story()

The trace can now be viewed in the UI: