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. The list lands on the wrapper's own model span (see Model spans; on by default with tracing) and, with disable_auto_instrument=False, on the instrumentor's span. 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",                  # Required: events without a user id are not sent
    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.

Model spans need no library instrumentation: the wrapper records its own (see Model spans). If you specifically want the instrumentor's spans instead and have verified compatibility in your environment, opt back in with disable_auto_instrument=False; the wrapper then adds no span for calls that instrumentation already records.

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()

Per-call properties and identity

generate_content has no channel for your own event data, so wrap the call in event_context. Its properties merge into every event captured inside the block, and user_id / convo_id replace the constructor values for those events:

from raindrop_vertex_ai import event_context

with event_context(user_id="user-123", convo_id="conv-9", properties={"plan": "pro"}):
    wrapped.models.generate_content(model="gemini-2.0-flash", contents="Hello!")

Updating the last event

last_event_id is the event_id of the most recent event this wrapper captured. Use it to attach a signal, or reopen the event with resume_interaction() (defaults to last_event_id) and update it through the core SDK's Interaction:

interaction = raindrop.resume_interaction()
if interaction is not None:
    interaction.set_properties({"reviewed": True})
    interaction.finish()

Key clashes. If one of your properties has the same key as one the wrapper derives itself, your value wins, whichever update carries it (a debug-level log line names the key). The wrapper's own keys are ai.* (for example ai.usage.prompt_tokens), error/error.* and vertex_ai.*; avoid them unless you mean to replace the wrapper's value.

Property-only updates need raindrop-ai 0.0.77 or newer; older core versions drop an update that carries no input or output.

last_event_id is per wrapper instance; under concurrent calls it is whichever event was captured last, so pass event_id= when you need to correlate concurrent calls.

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")

Standard options and flush results

When the wrapper builds its own client, it accepts the core SDK's standard options by the same names and passes them on:

raindrop = RaindropVertexAI(
    api_key="rk_...",
    user_id="user-123",
    redact_pii=True,        # local PII redaction of event input/output
    flush_timeout=5,        # seconds; deadline for flush() and shutdown()
    debug=True,             # also raises the core SDK's log verbosity
    max_queue_size=5_000,   # any other option `Raindrop(...)` declares
)

result = raindrop.flush()   # the core's FlushResult; safe to ignore

• An option the installed raindrop-ai does not support is ignored with a debug log, never an error. flush_timeout and the FlushResult return value need raindrop-ai 0.0.77 or later; on older releases flush() and shutdown() return None.
• With client= (a Raindrop you built), configure these options on that client; the wrapper ignores them (debug log).
• flush() and shutdown() never raise. The result is truthy when everything was delivered, and reports pending, dropped and timed_out.

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.

User id

Every event needs a real user id: set user_id= on the wrapper or per call with event_context(user_id=...). Without one (missing, blank, or a placeholder such as "unknown" or "anonymous"), the event is not sent and the wrapper logs one warning; it never raises. Updating an event you resumed with resume_interaction(event_id=...) needs no user id, since ingest keeps the one it stored.

Model spans

With tracing on (the default), the wrapper records each generate_content call as a span on the Raindrop client's tracing. The span carries:

• the model: gen_ai.provider.name (google-vertex for a Vertex AI client, google for a Gemini Developer API client), gen_ai.request.model and gen_ai.response.model;
• the token usage the provider reports: gen_ai.usage.input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens and reasoning_tokens;
• the event it belongs to: ai.telemetry.metadata.raindrop.eventId and the traceloop.association.properties.* keys.

Token usage also stays on the event as ai.usage.* properties. When an OpenTelemetry instrumentation that sees the call already records it (the provider's own instrumentation, for example with auto_instrument=True on the client), the wrapper adds no span of its own, so a call is never counted twice. An instrumentation for another provider or library does not suppress the wrapper's span.

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