Skip to main content
GitHub

Vertex AI

Auto-instrument Google Vertex AI.

Risicare automatically instruments Google Vertex AI for Gemini and other models.

Installation

pip install risicare google-cloud-aiplatform

Auto-Instrumentation

import risicare
import vertexai
from vertexai.generative_models import GenerativeModel
 
risicare.init()
 
vertexai.init(project="your-project", location="us-central1")
 
model = GenerativeModel("gemini-1.5-pro")
 
# Automatically traced
response = model.generate_content("Hello!")

Captured Attributes

AttributeDescription
gen_ai.systemvertexai
gen_ai.request.modelModel name
gen_ai.request.streamWhether streaming was requested
gen_ai.usage.prompt_tokensInput tokens (from prompt_token_count)
gen_ai.usage.completion_tokensOutput tokens (from candidates_token_count)
gen_ai.usage.total_tokensTotal tokens (from total_token_count)
gen_ai.completion.finish_reasonStop reason
gen_ai.latency_msRequest latency in milliseconds
gen_ai.response.candidatesNumber of response candidates
gen_ai.prompt.partsNumber of prompt parts

Streaming

responses = model.generate_content(
    "Write a story",
    stream=True
)
 
for response in responses:
    print(response.text, end="")

Chat Sessions

Not Instrumented

ChatSession.send_message is not instrumented. Only GenerativeModel.generate_content calls are traced. If you need tracing for multi-turn conversations, call generate_content directly with the full message history instead of using ChatSession.

Supported Models

ModelDescription
gemini-1.5-proGemini 1.5 Pro
gemini-1.5-flashGemini 1.5 Flash
gemini-1.0-proGemini 1.0 Pro
text-bisonPaLM 2 Text
chat-bisonPaLM 2 Chat

Function Calling

from vertexai.generative_models import FunctionDeclaration, Tool
 
get_weather = FunctionDeclaration(
    name="get_weather",
    description="Get weather for a location",
    parameters={
        "type": "object",
        "properties": {"location": {"type": "string"}},
        "required": ["location"]
    }
)
 
tools = Tool(function_declarations=[get_weather])
model = GenerativeModel("gemini-1.5-pro", tools=[tools])
 
response = model.generate_content("What's the weather in Paris?")
# Function calls are captured as child spans

Next Steps