Skip to main content

A Python library that meters Google AI API usage to Revenium with support for both Gemini Developer API and Vertex AI.

Project description

Revenium Middleware for Google AI (Python)

PyPI version Python Versions Documentation License: MIT

A middleware library for metering and monitoring Google AI services usage in Python applications. Supports both Google AI SDK (Gemini Developer API) and Vertex AI SDK with flexible optional dependencies.

For complete examples and setup instructions, see examples/README.md

Features

  • Precise Usage Tracking: Monitor tokens, costs, and request counts for Google AI services
  • Seamless Integration: Drop-in middleware that works with minimal code changes
  • Dual SDK Support: Choose between Google AI SDK or Vertex AI SDK based on your needs
  • Optional Dependencies: Install only the SDK components you need
  • Streaming Support: Full support for streaming responses (both SDKs)
  • Enhanced Token Counting: Complete token tracking including embeddings (Vertex AI)
  • Flexible Configuration: Customize metering behavior to suit your application needs

What's Supported

Feature Google AI SDK Vertex AI SDK
Chat Completion Full support Full support
Streaming Full support Full support
Text Embeddings Basic support* Full support
Token Metering Chat/Streaming All operations
Metadata Tracking Full support Full support
Setup Complexity Simple (API key) Moderate (GCP project)

Note: *Google AI SDK embeddings don't return token counts due to API limitations, but requests are still tracked.

Installation

Choose the SDK variant that best fits your needs:

# Google AI SDK only (Gemini Developer API)
pip install "revenium-middleware-google[genai]"

# Vertex AI SDK only (recommended for production)
pip install "revenium-middleware-google[vertex]"

# Both SDKs (maximum flexibility)
pip install "revenium-middleware-google[all]"

Which SDK Should I Choose?

Use Case Recommended SDK Why
Quick prototyping Google AI SDK Simple API key setup, but does NOT support token counts on embeddings
Production applications Vertex AI SDK Full token counting, enterprise features
Embeddings-heavy workloads Vertex AI SDK Complete token tracking for embeddings
Enterprise/GCP environments Vertex AI SDK Advanced Google Cloud integration
Simple chat applications Either SDK Both provide full chat support

Recommendation: Use Vertex AI SDK for production applications that need comprehensive token counting and advanced features.

Quick Start

1. Create Project Directory

mkdir my-google-ai-project
cd my-google-ai-project

2. Create Virtual Environment

# Create virtual environment
python -m venv .venv

# Activate virtual environment
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

3. Install the Package

# Install with both SDKs (recommended)
pip install "revenium-middleware-google[all]" python-dotenv

4. Configure Environment Variables

Create a .env file in your project directory:

# Required
REVENIUM_METERING_API_KEY=hak_your_revenium_api_key

# For Google AI SDK (Gemini Developer API)
GOOGLE_API_KEY=your_google_api_key

# For Vertex AI SDK (Google Cloud)
GOOGLE_CLOUD_PROJECT=your_gcp_project_id
GOOGLE_CLOUD_LOCATION=us-central1

# Optional
# REVENIUM_LOG_LEVEL=DEBUG

For Vertex AI: Authenticate with Google Cloud:

gcloud auth application-default login

5. Run Your First Example

Download and run an example from the repository:

# For Google AI SDK
curl -O https://raw.githubusercontent.com/revenium/revenium-middleware-google-python/main/examples/getting_started_google_ai.py
python getting_started_google_ai.py

# For Vertex AI SDK
curl -O https://raw.githubusercontent.com/revenium/revenium-middleware-google-python/main/examples/getting_started_vertex_ai.py
python getting_started_vertex_ai.py

Or create a file test.py with this simple code:

from dotenv import load_dotenv
load_dotenv()  # Load environment variables from .env file

import revenium_middleware_google
from google import genai

client = genai.Client()
response = client.models.generate_content(
    model="gemini-2.0-flash-001",
    contents="Hello! Introduce yourself in one sentence."
)
print(response.text)

Run it:

python test.py

The middleware automatically tracks your Google AI usage and sends data to Revenium.

For complete examples and usage patterns, see examples/README.md.


Usage

Simply import the middleware and your Google AI calls will be metered automatically:

import revenium_middleware_google
from google import genai

client = genai.Client()
response = client.models.generate_content(
    model="gemini-2.0-flash-001",
    contents="Hello! Introduce yourself in one sentence."
)
print(response.text)

The middleware automatically meters all Google AI and Vertex AI calls. For production use, add metadata tracking to associate usage with organizations, users, or features.

For complete examples, see:

Configuration

For detailed configuration options, environment variables, and advanced setup, see the Configuration section in examples/README.md

Metadata Fields

Add business context to track usage by organization, user, task type, or custom fields. Pass a usage_metadata dictionary with any of these optional fields:

Field Description Use Case
trace_id Unique identifier for session or conversation tracking Link multiple API calls together for debugging, user session analytics, or distributed tracing across services
task_type Type of AI task being performed Categorize usage by workload (e.g., "chat", "code-generation", "doc-summary") for cost analysis and optimization
subscriber.id Unique user identifier Track individual user consumption for billing, rate limiting, or user analytics
subscriber.email User email address Identify users for support, compliance, or usage reports
subscriber.credential.name Authentication credential name Track which API key or service account made the request
subscriber.credential.value Authentication credential value Associate usage with specific credentials for security auditing
organizationName Organization or company name for lookup/auto-creation Multi-tenant cost allocation, usage quotas per organization (e.g., "AcmeCorp")
subscription_id Subscription plan identifier Track usage against subscription limits, identify plan upgrade opportunities
productName Your product or feature name for lookup/auto-creation Attribute AI costs to specific features in your application (e.g., "customer-chatbot")
agent AI agent or bot identifier Distinguish between multiple AI agents or automation workflows in your system
response_quality_score Custom quality rating (0.0-1.0) Track user satisfaction or automated quality metrics for model performance analysis

Note: organizationName and productName replace the deprecated field names. The legacy field names organizationId, organization_id, productId, and product_id are deprecated but still supported for backward compatibility. Use organizationName (or organization_name) and productName (or product_name) instead.

Example usage: See examples/getting_started_google_ai.py for complete metadata implementation.

API Reference: Complete metadata field documentation

Trace Visualization Fields (v0.2.0+)

Enhanced observability fields for distributed tracing and analytics. These can be set via environment variables or passed in usage_metadata:

Field Environment Variable Description Use Case
environment REVENIUM_ENVIRONMENT Deployment environment (e.g., "production", "staging") Track usage across different deployment environments; auto-detects from ENVIRONMENT, DEPLOYMENT_ENV
region REVENIUM_REGION Cloud region identifier (e.g., "us-east-1", "eastus") Multi-region deployment tracking; auto-detects from AWS_REGION, AZURE_REGION, GCP_REGION, GOOGLE_CLOUD_REGION
credential_alias REVENIUM_CREDENTIAL_ALIAS Human-readable API key name (e.g., "prod-google-key") Track which credential was used for credential rotation and security auditing
trace_type REVENIUM_TRACE_TYPE Workflow category identifier (max 128 chars) Group similar workflows (e.g., "customer-support", "data-analysis") for analytics
trace_name REVENIUM_TRACE_NAME Human-readable trace label (max 256 chars) Label trace instances (e.g., "Customer Support Chat", "Document Analysis")
parent_transaction_id REVENIUM_PARENT_TRANSACTION_ID Parent transaction ID for distributed tracing Link child operations to parent transactions across services
transaction_name REVENIUM_TRANSACTION_NAME Human-friendly operation name Label individual operations (e.g., "Generate Response", "Analyze Sentiment")

Note: operation_type and operation_subtype are automatically detected by the middleware based on the API method and request parameters.

Resources:

Tool Metering

The @meter_tool decorator lets you meter arbitrary tool/function calls (document parsers, data processors, etc.) alongside your automatic LLM API metering. Requires revenium-middleware>=0.4.2.

import os
from revenium_middleware import meter_tool, configure

# Configure the metering client for tool calls
configure(
    metering_url=os.getenv("REVENIUM_METERING_BASE_URL", "https://api.revenium.ai"),
    api_key=os.environ["REVENIUM_METERING_API_KEY"],  # Fail fast if not set
)

@meter_tool("document-parser", operation="parse", agent="doc-assistant")
def parse_document(file_path: str) -> dict:
    """Parse a document - timing and success/failure are automatically tracked."""
    # Your parsing logic here
    return {"pages": 12, "content": "..."}

# The decorator reports the tool call to Revenium automatically
result = parse_document("reports/q3_financials.pdf")

See examples/example_meter_tool.py for a complete example combining tool metering with Google AI LLM metering.

Terminal Summary Output

Display cost and metrics information in your terminal after each API request. Supports both human-readable and JSON formats.

Configuration:

Environment Variable Values Description
REVENIUM_PRINT_SUMMARY false (default), true, human, json Controls output format
REVENIUM_TEAM_ID Your team ID Required to fetch and display cost information

Enable in your .env file:

# Human-readable format
REVENIUM_PRINT_SUMMARY=human

# JSON format (for machine parsing)
REVENIUM_PRINT_SUMMARY=json

# Required for cost information
REVENIUM_TEAM_ID=your-team-id-here

Human-Readable Output Example:

REVENIUM USAGE SUMMARY
Model: gemini-2.0-flash-001
Provider: Google
Duration: 1.23s

Token Usage:
  Input Tokens:  150
  Output Tokens: 250
  Total Tokens:  400

Cost: $0.000045

Trace ID: abc-123

JSON Output Example:

{"model":"gemini-2.0-flash-001","provider":"Google","durationSeconds":1.23,"inputTokenCount":150,"outputTokenCount":250,"totalTokenCount":400,"cost":0.000045,"costStatus":"available","traceId":"abc-123"}

Notes:

  • Cost information requires REVENIUM_TEAM_ID to be set
  • If cost is not yet available, the summary will indicate "Pending" (human) or "costStatus":"unavailable" (JSON)
  • Summary output is fire-and-forget and will never block or fail your main API calls

Prompt Capture (v0.3.0+)

The middleware can capture and send prompts (system instructions, user inputs, and model outputs) to Revenium for monitoring, debugging, and compliance purposes.

Enabling Prompt Capture

Set the environment variable to enable prompt capture:

# Add to your .env file
REVENIUM_CAPTURE_PROMPTS=true

Or configure programmatically:

import os
os.environ['REVENIUM_CAPTURE_PROMPTS'] = 'true'

# Import middleware AFTER setting the environment variable
import revenium_middleware_google

What Gets Captured

When enabled, the following data is sent to Revenium:

  1. System Prompt: The system instruction provided to the model
  2. Input Messages: The user's prompt/question (contents parameter)
  3. Output Response: The model's complete response

Configuration Options

Environment Variable Default Description
REVENIUM_CAPTURE_PROMPTS false Enable/disable prompt capture
REVENIUM_MAX_PROMPT_LENGTH 10000 Maximum characters per field (truncated with [TRUNCATED] marker)

Example

import os

# Enable prompt capture
os.environ['REVENIUM_CAPTURE_PROMPTS'] = 'true'

import revenium_middleware_google
from google import genai

# Configure API
client = genai.Client(api_key=os.environ['GOOGLE_API_KEY'])

# Create model with system instruction
model_id = 'gemini-2.0-flash-exp'
config = genai.types.GenerateContentConfig(
    system_instruction="You are a helpful math tutor."
)

# This will capture:
# - System prompt: "You are a helpful math tutor."
# - Input message: "What is the Pythagorean theorem?"
# - Output response: The model's complete answer
response = client.models.generate_content(
    model=model_id,
    contents="What is the Pythagorean theorem?",
    config=config
)

Privacy and Security Considerations

  • Prompts may contain sensitive user data or proprietary information
  • Only enable prompt capture when you have appropriate data handling policies
  • Captured prompts are sent to Revenium's metering API
  • Consider using REVENIUM_MAX_PROMPT_LENGTH to limit data transmission
  • Review your organization's data privacy requirements before enabling

Supported SDKs

Prompt capture works with both:

  • ✅ Google AI SDK (Gemini Developer API)
  • ✅ Vertex AI SDK

Both streaming and non-streaming requests are supported.

Troubleshooting

Common Issues

Issue Solution
"No module named 'google.genai'" Install with Google AI support: pip install "revenium-middleware-google[genai]"
"No module named 'vertexai'" Install with Vertex AI support: pip install "revenium-middleware-google[vertex]"
Vertex AI authentication errors Verify Google Cloud credentials: gcloud auth application-default login
"Project not found" errors Ensure GOOGLE_CLOUD_PROJECT is set correctly
Embeddings showing 0 tokens Expected with Google AI SDK; use Vertex AI for full token counting
Requests not being tracked Ensure middleware is imported before Google AI/Vertex AI SDKs

Debug Mode

Enable debug logging to see provider detection and routing decisions:

# export REVENIUM_LOG_LEVEL=DEBUG
python your_script.py

Force Specific SDK

To ensure only one SDK is used:

# Use only Google AI SDK
pip install "revenium-middleware-google[genai]"

# Use only Vertex AI SDK
pip install "revenium-middleware-google[vertex]"

Google AI SDK Troubleshooting

Middleware not tracking requests:

  • Ensure middleware is imported before Google AI SDK
  • Check that environment variables are loaded correctly
  • Verify your REVENIUM_METERING_API_KEY is correct

Embeddings showing 0 tokens:

  • This is expected due to Google AI SDK limitations
  • Model name and metadata are still tracked correctly
  • Chat and streaming operations provide full token data

Vertex AI SDK Troubleshooting

Authentication issues:

  • Verify Google Cloud credentials: gcloud auth list
  • Check project access: gcloud projects describe YOUR_PROJECT_ID
  • Ensure service account has required permissions

Model not available errors:

  • Check if models are available in your region
  • Verify Vertex AI API is enabled in your project
  • Try a different model or region

Logging

Control logging with the REVENIUM_LOG_LEVEL environment variable in your .env file:

# Add to your .env file
REVENIUM_LOG_LEVEL=DEBUG

Available log levels:

  • DEBUG: Detailed debugging information (shows provider detection, routing decisions, token counts)
  • INFO: General information messages (default)
  • WARNING: Warning messages only
  • ERROR: Error messages only
  • CRITICAL: Critical error messages only

Compatibility

  • Python 3.8+
  • Google AI SDK (google-genai>=0.1.0) or Vertex AI SDK (google-cloud-aiplatform>=1.0.0)
  • Google Cloud Project (for Vertex AI SDK)

Supported Models

The middleware automatically works with all Google AI models available through both SDKs, including:

  • Gemini models (all versions and variants)
  • Text embedding models (all versions)
  • Future model releases (automatic support with no code changes required)

The middleware detects and meters any model supported by the Google AI SDK or Vertex AI SDK. When Google releases new models, they work immediately without updates to the middleware.

Documentation

For detailed documentation, visit docs.revenium.io

Contributing

See CONTRIBUTING.md

Code of Conduct

See CODE_OF_CONDUCT.md

Security

See SECURITY.md

License

This project is licensed under the MIT License - see the LICENSE file for details.

Support

For issues, feature requests, or contributions:


Built by Revenium

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

revenium_middleware_google-0.4.0.tar.gz (68.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

revenium_middleware_google-0.4.0-py3-none-any.whl (76.3 kB view details)

Uploaded Python 3

File details

Details for the file revenium_middleware_google-0.4.0.tar.gz.

File metadata

File hashes

Hashes for revenium_middleware_google-0.4.0.tar.gz
Algorithm Hash digest
SHA256 ddd16da7cf78febbcc816c274b60b6fc4bb221909acc9ad02f964380293d0d9c
MD5 7e548f50efbf55bc59e02d92635af02d
BLAKE2b-256 348e9039a99333fc75d3e0d4daa8b16015cb6fb740edcf05601a8f790d26eed1

See more details on using hashes here.

File details

Details for the file revenium_middleware_google-0.4.0-py3-none-any.whl.

File metadata

File hashes

Hashes for revenium_middleware_google-0.4.0-py3-none-any.whl
Algorithm Hash digest
SHA256 4b046ed30a21b138246937636820d27400617d627a402cbb6b4891d93680c365
MD5 808b59237f020572842596812c4d9173
BLAKE2b-256 028b2ebae096162389469377053e247481f5d6d2d0a306cde38b2afb37d72e1e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page