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

Tool catalog capture (ai.prompt.tools)

Every model span produced under a wrapped generate_content() / generate_content_stream() call (sync and aio) records the request's config.tools as ai.prompt.tools: each function_declarations entry becomes a function tool (parameters or parameters_json_schema recorded as given as inputSchema) and other keys (google_search, code_execution, retrieval, url_context, ...) become provider-defined tools. Dict configs and google.genai.types.GenerateContentConfig / types.Tool objects are both read. Capture is exact: the wrapper reads the list off the request it is about to send (the OpenLLMetry vertexai / google_generativeai instrumentors emit no tool definitions themselves). Each tool is one JSON document in an OTLP string array: {"type":"function","name","description","inputSchema"} for function tools (the JSON Schema is recorded as given — properties, enums, required and per-field descriptions intact) or {"type":"provider-defined","id","name","args"} for provider-owned tools, which get no invented schema. A call without a visible tool list records no attribute (this includes a tools list containing bare Python callables, which the SDK turns into declarations only on the wire); an explicit empty list records [].

Pass tools= to RaindropVertexAI(...) / create_raindrop_vertex_ai(...) to record a different catalog instead: it replaces the inferred list for every call through that wrapper (never merges), and tools=[] records an empty catalog. The wrapper has no per-call metadata channel, so there is no per-call override. Tool definitions are prompt content and follow the same switch as prompts: with TRACELOOP_TRACE_CONTENT=false they are never recorded. Model spans only exist when library instrumentation is enabled (disable_auto_instrument=False), and capture requires raindrop-ai>=0.0.70 (older cores log at debug and record nothing).

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.

Application Git metadata

RaindropVertexAI(...) and create_raindrop_vertex_ai(...) accept the keyword-only app_git option. It defaults to True: explicit Raindrop Git environment or deployment context is applied immediately, and the base SDK may perform one bounded background local-Git lookup from the process working directory. Event capture, flush, and shutdown never wait for that lookup. Pass False to disable enrichment, or pass an AppGitOptions mapping with commit_sha, commit_dirty, branch, source_directory, detect_branch, and/or auto_detect. Automatic branch discovery remains opt-in through detect_branch=True (or RAINDROP_GIT_DETECT_BRANCH=true).

For an ordinary in-process application, the process working directory is treated as the application-under-test checkout. A remote, coding, workflow, or observer process must not rely on its own checkout: pass app_git=False, provide explicit revision values, or set source_directory to the actual application checkout. Canonical per-operation properties remain authoritative. When supplying client=, configure app_git while constructing that Raindrop client; the supplied client is authoritative and the wrapper's app_git argument does not reconfigure it.

Release order is deliberate: first publish the base SDK feature, then publish the wrapper feature release with its minimum dependency coordinated to that base release. The existing raindrop-ai lower bound remains compatible, but application Git metadata is unavailable on an older core and must not be claimed complete until the base is upgraded. Until coordination assigns a released version, the wrapper checks for an explicit base app_git parameter and omits the option when unsupported. Explicit non-default configuration is debug-logged and omitted. Unsupported app_git is determined by signature inspection before construction, not by retrying initialization after a TypeError; Git configuration adds no initialization attempts and does not change any existing framework-specific initialization fallback.

Testing

cd packages/vertex-ai-python
pip install -e ".[dev]"
python -m pytest tests/ -v   # unit tests (no external services)

End-to-end behavior is verified by the cross-SDK conformance harness. This package ships a thin conformance driver at conformance/driver.py that maps the shared scenario corpus onto the wrapper's public API; known gaps are tracked as ticket-linked entries in conformance/failures.txt. The fault lane runs on every PR touching packages/*-python/** (.github/workflows/conformance-wrappers-python.yml) against a local capture server; the prod lane verifies delivery by reading back through the public Query API. The harness is pinned by commit SHA (HARNESS_REF). See the harness docs: HOW-IT-WORKS · AGENTS · README.

License

MIT