Vertex AI Python SDK Integration

Use Vertex AI’s Python SDK to integrate with Helicone to log your Vertex AI usage.

Proxy Integration

Python SDK

1

Create an account + Generate an API Key

Log into helicone or create an account. Once you have an account, you can generate an API key.

2

Set API keys as environment variables

export HELICONE_API_KEY=<your API key>
export PROJECT_ID=<your Google Cloud project ID>
export LOCATION=<your location>

3

Install required packages

Ensure you have the necessary packages installed in your Python environment:

pip install google-cloud-aiplatform

4

Import libraries

from vertexai.generative_models import GenerativeModel
import vertexai
import os

5

Initialize Vertex AI with Helicone

HELICONE_API_KEY = os.environ.get("HELICONE_API_KEY")
PROJECT_ID = os.environ.get("PROJECT_ID")
LOCATION = os.environ.get("LOCATION")

vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    api_endpoint="gateway.helicone.ai",
    api_transport="rest",  # Must be 'rest' or else it will not work
    request_metadata=[
        ('helicone-target-url', f'https://{LOCATION}-aiplatform.googleapis.com'),
        ('helicone-auth', f'Bearer {HELICONE_API_KEY}')
    ]
)

6

Initialize the model and generate content

model = GenerativeModel("gemini-1.5-flash-001")
response = model.generate_content("Tell me a fun fact about space.")
print(response.text)

Manual Logging with Log Builder

For more control over logging, especially with asynchronous and streaming responses, you can use the Helicone Manual Logger:

1

Install additional required packages

pip install python-dotenv helicone

2

Create a log builder example

This example demonstrates using the Helicone log builder with both async and streaming responses:

import os
import asyncio
from dotenv import load_dotenv
import vertexai
from vertexai.generative_models import GenerativeModel
from helicone_helpers import HeliconeManualLogger
import warnings
import sys

# Load environment variables
load_dotenv()

HELICONE_API_KEY = os.getenv("HELICONE_API_KEY")
PROJECT_ID = os.getenv("PROJECT_ID")
LOCATION = os.getenv("LOCATION")

helicone_logger = HeliconeManualLogger(api_key=HELICONE_API_KEY)

async def generate_content_async(model, prompt="Tell me a joke"):
    generate_kwargs = {
        "contents": [prompt],
    }

    log_builder = helicone_logger.new_builder(
        request=generate_kwargs
    )
    log_builder.add_model("gemini-1.5-pro")
    try:
        response = await model.generate_content_async(**generate_kwargs)
        log_builder.add_response(response.to_dict())
        print(f"Response: {response.text}")
        return response.text
    except Exception as e:
        print(f"Error: {e}")
        return str(e)
    finally:
        await log_builder.send_log()

async def generate_content_streaming(model: GenerativeModel, prompt="Tell me a joke"):
    generate_kwargs = {
        "contents": [prompt],
        "stream": True,
    }

    log_builder = helicone_logger.new_builder(
        request=generate_kwargs
    )
    log_builder.add_model("gemini-1.5-pro")
    try:
        responses = await model.generate_content_async(**generate_kwargs)

        full_response = []
        print("Streaming response:")
        async for response in responses:
            log_builder.add_chunk(response.to_dict())
            chunk = response.text
            print(chunk, end="", flush=True)
            full_response.append(chunk)
        print("\n")
        return "".join(full_response)
    except Exception as e:
        print(f"Error in streaming: {e}")
        return str(e)
    finally:
        await log_builder.send_log()

async def main():
    print("Vertex AI Gemini with Helicone Example")

    vertexai.init(
        project=PROJECT_ID,
        location=LOCATION,
    )

    model = GenerativeModel(
        "gemini-1.5-pro",
        generation_config={"response_mime_type": "application/json"}
    )

    print("\n=== Regular Async Response ===")
    await generate_content_async(model, "Tell me a joke about programming")

    print("\n=== Streaming Response ===")
    await generate_content_streaming(model, "Tell me a short 2 sentence story about a programmer")

if __name__ == "__main__":
    if not sys.warnoptions:
        warnings.filterwarnings("ignore", category=ResourceWarning)

    asyncio.run(main())

    sys.exit(0)

3

Create the HeliconeManualLogger helper

Create a file named helicone_helpers.py with the following content:

import json
import os
import aiohttp
import requests

class HeliconeManualLogger:
    def __init__(self, api_key=None):
        self.api_key = api_key or os.environ.get("HELICONE_API_KEY")
        if not self.api_key:
            raise ValueError("Helicone API key is required")

        self.base_url = "https://api.helicone.ai/v1/log"

    def new_builder(self, request):
        return LogBuilder(self.api_key, request)

class LogBuilder:
    def __init__(self, api_key, request):
        self.api_key = api_key
        self.request = request
        self.model = None
        self.response = None
        self.chunks = []

    def add_model(self, model):
        self.model = model
        return self

    def add_response(self, response):
        self.response = response
        return self

    def add_chunk(self, chunk):
        self.chunks.append(chunk)
        return self

    async def send_log(self):
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "provider": "vertex",
            "model": self.model,
            "request": self.request
        }

        if self.response:
            payload["response"] = self.response

        if self.chunks:
            payload["response_chunks"] = self.chunks

        async with aiohttp.ClientSession() as session:
            async with session.post(
                "https://api.helicone.ai/v1/log",
                headers=headers,
                json=payload
            ) as response:
                if response.status != 200:
                    print(f"Failed to log to Helicone: {response.status}")
                    text = await response.text()
                    print(f"Response: {text}")
                return response.status == 200

The log builder approach provides several advantages:

• Precise control over what gets logged
• Support for both async and streaming responses
• Ability to add custom properties and metadata
• Manual control over when logs are sent

This is particularly useful for complex applications where you need more control over the logging process.