raindrop-vertex-ai

Raindrop integration for Google Vertex AI / Gen AI (Python). Automatically captures models.generate_content() and aio.models.generate_content() calls.

Installation

pip install raindrop-vertex-ai google-genai

google-genai is a required dependency.

Quick Start

from raindrop_vertex_ai import RaindropVertexAI
from google import genai

raindrop = RaindropVertexAI(api_key="your-write-key", user_id="user-123")

client = genai.Client(api_key="...")
wrapped = raindrop.wrap(client)

response = wrapped.models.generate_content(
    model="gemini-2.0-flash", contents="Hello!"
)
print(response.text)

raindrop.shutdown()

Omitting api_key disables telemetry shipping (a warning is emitted) but does not crash your application.

What Gets Tracked

• generate_content — input content, output text, model name
• Token usage — prompt_token_count and candidates_token_count from usage metadata
• Cached tokens — cached_content_token_count from usage metadata → ai.usage.cached_tokens
• Thinking tokens — thoughts_token_count from usage metadata (Gemini 2.5) → ai.usage.thoughts_tokens
• Finish reason — candidate.finish_reason (STOP, MAX_TOKENS, SAFETY, RECITATION) → vertex_ai.finish_reason
• Errors — captured with error type and message in properties, then re-raised
• Async support — both models.generate_content (sync) and aio.models.generate_content (async) are instrumented
• Double-wrap guard — calling wrap() on an already-wrapped client is a safe no-op

Configuration

raindrop = RaindropVertexAI(
    api_key="rk_...",                    # Optional: your Raindrop API key
    user_id="user-123",                  # Optional: associate events with a user
    convo_id="convo-456",                # Optional: conversation/thread ID
    project_id="support-prod",           # Optional: route events to a specific project (slug)
    tracing_enabled=True,                # Optional: enable Raindrop tracing (default: True)
    bypass_otel_for_tools=True,          # Optional: bypass OTEL for tools (default: True)
    debug=True,                          # Optional: enable verbose DEBUG logging
)

Library auto-instrumentation is opt-in

As of 0.0.4, disable_auto_instrument defaults to True: the integration no longer lets Traceloop monkey-patch every LLM client library it recognizes in your process (google.genai itself, Anthropic, botocore, etc.). The wrapper captures input/output, token usage, model name, and finish_reason directly from the wrapped client's responses, so no library patching is needed for full dashboards.

If you specifically want LLM-call-level spans from library instrumentation and have verified compatibility in your environment, opt back in with disable_auto_instrument=False.

Debug Logging

raindrop = RaindropVertexAI(api_key="rk_...", debug=True)

Async Usage

import asyncio
from raindrop_vertex_ai import RaindropVertexAI
from google import genai

raindrop = RaindropVertexAI(api_key="rk_...")
client = genai.Client(api_key="...")
wrapped = raindrop.wrap(client)

async def main():
    response = await wrapped.aio.models.generate_content(
        model="gemini-2.0-flash", contents="Hello!"
    )
    print(response.text)

asyncio.run(main())
raindrop.shutdown()

identify()

raindrop.identify(user_id="user-123", traits={"plan": "pro", "org": "acme"})

track_signal()

raindrop.track_signal(
    event_id="evt-abc",
    name="thumbs_up",
    signal_type="feedback",
    sentiment="POSITIVE",
    comment="Great response!",
)

Flushing and Shutdown

Always call shutdown() before your process exits to ensure all telemetry is shipped:

raindrop.flush()     # flush pending data
raindrop.shutdown()  # flush + release resources

Projects

Route events to a specific project by passing its slug as project_id:

raindrop = RaindropVertexAI(
    api_key="rk_...",
    project_id="support-prod",
)

project_id sets the X-Raindrop-Project-Id header on every event. Omit it (or pass "default") to use your org's default Production project, which is the existing behavior. The same option is accepted by the create_raindrop_vertex_ai(...) factory. Invalid slugs are ignored with a warning and no header is sent.

Factory Function

A create_raindrop_vertex_ai() factory is also available for convenience:

from raindrop_vertex_ai import create_raindrop_vertex_ai

raindrop = create_raindrop_vertex_ai(api_key="rk_...", user_id="user-123")

Full Documentation

See the Raindrop docs for complete API reference.

Testing

cd packages/vertex-ai-python
pip install -e ".[dev]"
python -m pytest tests/ -v

License

MIT