Google Vertex AI Integration

pip install turingpulse[vertexai]

pip install turingpulse[all]

import vertexai
from vertexai.generative_models import GenerativeModel
from turingpulse import init
from turingpulse.integrations.vertexai import instrument_vertexai

# Initialize TuringPulse
init(
    api_key="sk_live_...",
    project_id="my-project"
)

# Instrument Vertex AI - this wraps all Gemini calls
instrument_vertexai()

# Initialize Vertex AI
vertexai.init(project="your-gcp-project", location="us-central1")

# Your code works exactly the same - now with full tracing!
model = GenerativeModel("gemini-1.5-pro")
response = model.generate_content("Explain quantum computing in simple terms")
print(response.text)

from turingpulse.integrations.vertexai import instrument_vertexai

instrument_vertexai(
    # Capture full prompts and responses (default: True)
    capture_content=True,
    
    # Track token counts and costs (default: True)
    capture_tokens=True,
    
    # Add custom metadata to all traces
    default_metadata={
        "environment": "production",
        "team": "ml-platform"
    },
    
    # Filter which models to trace (default: all)
    model_filter=lambda model: "gemini" in model.lower(),
)

from vertexai.generative_models import GenerativeModel

model = GenerativeModel("gemini-1.5-flash")

# Streaming is automatically instrumented
response = model.generate_content(
    "Write a haiku about AI",
    stream=True
)

for chunk in response:
    print(chunk.text, end="", flush=True)

# Full trace is captured including:
# - Time to first token
# - Total streaming duration
# - Complete token counts

from vertexai.generative_models import GenerativeModel, Part

model = GenerativeModel("gemini-1.5-pro")

# Image input is tracked (image content stored as reference)
response = model.generate_content([
    Part.from_uri("gs://your-bucket/image.jpg", mime_type="image/jpeg"),
    "Describe what you see in this image"
])

# Trace captures:
# - Image URI reference
# - Text prompt
# - Model response
# - Multi-modal token counts

from vertexai.generative_models import GenerativeModel, FunctionDeclaration, Tool

# Define your functions
get_weather = FunctionDeclaration(
    name="get_weather",
    description="Get the current weather for a location",
    parameters={
        "type": "object",
        "properties": {
            "location": {"type": "string", "description": "City name"}
        },
        "required": ["location"]
    }
)

weather_tool = Tool(function_declarations=[get_weather])
model = GenerativeModel("gemini-1.5-pro", tools=[weather_tool])

# Function calls are automatically traced
response = model.generate_content("What's the weather in San Francisco?")

# Trace captures:
# - Function call request
# - Function arguments
# - Function response (when provided)

from vertexai.generative_models import GenerativeModel

model = GenerativeModel("gemini-1.5-pro")
chat = model.start_chat()

# Each message in the conversation is traced
response1 = chat.send_message("Hi! I'm building an AI agent.")
response2 = chat.send_message("What frameworks should I consider?")
response3 = chat.send_message("Tell me more about LangGraph")

# Full conversation history is tracked with:
# - Session ID linking all messages
# - Turn-by-turn latency and tokens
# - Cumulative conversation metrics

from vertexai.generative_models import GenerativeModel
from google.api_core.exceptions import ResourceExhausted

model = GenerativeModel("gemini-1.5-pro")

try:
    response = model.generate_content("...")
except ResourceExhausted as e:
    # Error is automatically captured in trace with:
    # - Error type and message
    # - Stack trace
    # - Request that caused the error
    raise

# TuringPulse configuration
export TURINGPULSE_API_KEY="sk_live_..."
export TURINGPULSE_PROJECT_ID="my-project"

# Vertex AI configuration  
export GOOGLE_CLOUD_PROJECT="your-gcp-project"
export GOOGLE_CLOUD_REGION="us-central1"