Metadata-Version: 2.4
Name: raindrop-bedrock
Version: 0.0.3
Summary: Raindrop observability integration for AWS Bedrock
Author-email: Raindrop AI <sdk@raindrop.ai>
License-Expression: MIT
Project-URL: Homepage, https://raindrop.ai
Project-URL: Documentation, https://docs.raindrop.ai
Project-URL: Repository, https://github.com/invisible-tools/dawn/tree/main/packages/bedrock-python
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: raindrop-ai>=0.0.42
Requires-Dist: boto3>=1.34.0
Provides-Extra: async
Requires-Dist: aioboto3>=12.0.0; extra == "async"
Provides-Extra: dev
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: requests>=2.28.0; extra == "dev"
Dynamic: license-file

# raindrop-bedrock

[![PyPI version](https://img.shields.io/pypi/v/raindrop-bedrock)](https://pypi.org/project/raindrop-bedrock/)
[![Python](https://img.shields.io/pypi/pyversions/raindrop-bedrock)](https://pypi.org/project/raindrop-bedrock/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

[Raindrop](https://raindrop.ai) observability integration for **AWS Bedrock** (Python). Automatically captures `converse()` and `invoke_model()` calls by wrapping the boto3 bedrock-runtime client.

## Installation

```bash
pip install raindrop-bedrock
```

For async support with `aioboto3`:

```bash
pip install raindrop-bedrock[async]
```

## Quick Start

```python
import boto3
from raindrop_bedrock import RaindropBedrock

rb = RaindropBedrock(api_key="your-write-key", user_id="user-123")

client = boto3.client("bedrock-runtime", region_name="us-east-1")
rb.wrap(client)

response = client.converse(
    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
    messages=[{"role": "user", "content": [{"text": "Hello!"}]}],
)

rb.flush()
```

### Debug Mode

Enable verbose logging with the `debug` flag:

```python
rb = RaindropBedrock(api_key="your-write-key", user_id="user-123", debug=True)
```

### Async Usage

```python
import aioboto3
from raindrop_bedrock import RaindropBedrock

rb = RaindropBedrock(api_key="rk_...", user_id="user-123")

session = aioboto3.Session()
async with session.client("bedrock-runtime", region_name="us-east-1") as client:
    rb.async_wrap(client)
    response = await client.converse(
        modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
        messages=[{"role": "user", "content": [{"text": "Hello!"}]}],
    )

rb.flush()
```

## identify()

Associate a user with optional traits:

```python
rb.identify(user_id="user-123", traits={"plan": "pro", "company": "Acme"})
```

## track_signal()

Track feedback, edits, or custom signals:

```python
rb.track_signal(
    event_id="evt_abc123",
    name="thumbs_up",
    signal_type="feedback",
    sentiment="POSITIVE",
    comment="Great answer!",
)
```

## flush() / shutdown()

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

```python
rb.flush()      # flush pending data
rb.shutdown()   # flush + release resources
```

### Legacy Factory Function

The `create_raindrop_bedrock()` factory function is still supported for backwards compatibility:

```python
from raindrop_bedrock import create_raindrop_bedrock

raindrop = create_raindrop_bedrock(api_key="your-write-key", user_id="user-123")
client = boto3.client("bedrock-runtime", region_name="us-east-1")
raindrop.wrap(client)
```

## What Gets Captured

| Method | Captured Data |
|--------|--------------|
| `converse()` | Input messages, output text, model ID, token usage (`inputTokens`/`outputTokens`), stop reason (`stopReason`), cached tokens (`cacheReadInputTokenCount`, `cacheWriteInputTokenCount`), conversation ID |
| `invoke_model()` | Raw request/response bodies, model ID, token usage (Claude, Titan, and Llama formats), stop reason (Claude: `stop_reason`, Llama: `stop_reason`), cached tokens (Claude: `cache_read_input_tokens`) |
| **Errors** | Error type and message are captured in event properties, then the exception is re-raised |

### Captured Properties

| Property Key | Source | Description |
|---|---|---|
| `ai.usage.prompt_tokens` | Both APIs | Input/prompt token count |
| `ai.usage.completion_tokens` | Both APIs | Output/completion token count |
| `ai.usage.cached_tokens` | Converse: `cacheReadInputTokenCount`; Claude InvokeModel: `cache_read_input_tokens` | Tokens read from cache |
| `ai.usage.cache_write_tokens` | Converse: `cacheWriteInputTokenCount` | Tokens written to cache |
| `bedrock.finish_reason` | Converse: `stopReason`; InvokeModel: varies by model | Why the model stopped generating |

## API Reference

### `RaindropBedrock(api_key=None, user_id=None, convo_id=None, tracing_enabled=True, bypass_otel_for_tools=True, debug=False)`

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `api_key` | `str \| None` | `None` | Raindrop API key. Warns if not provided. |
| `user_id` | `str \| None` | `None` | Default user ID for events (falls back to `"unknown"`) |
| `convo_id` | `str \| None` | `None` | Group events into a conversation |
| `tracing_enabled` | `bool` | `True` | Enable Raindrop tracing |
| `bypass_otel_for_tools` | `bool` | `True` | Bypass OpenTelemetry for tool-level instrumentation |
| `debug` | `bool` | `False` | Enable verbose DEBUG-level logging |

### Methods

| Method | Description |
|--------|-------------|
| `wrap(client)` | Instrument a sync boto3 bedrock-runtime client |
| `async_wrap(client)` | Instrument an async aioboto3 bedrock-runtime client |
| `identify(user_id, traits=None)` | Identify a user with optional traits |
| `track_signal(event_id, name, ...)` | Track a signal event |
| `flush()` | Flush pending events |
| `shutdown()` | Flush and shut down |

## Testing

```bash
cd packages/bedrock-python
pip install -e ".[async]"
pip install pytest
python -m pytest tests/ -v
```

## Known Limitations

- **InvokeModel body replacement**: After consuming the response body stream, it's replaced with a `BytesIO` object. Callers using `StreamingBody.read()` will get the same bytes, but the original `StreamingBody` API is not preserved.
- **Async support** requires the `[async]` extra (`aioboto3>=12.0.0`).

## Full Documentation

[docs.raindrop.ai/integrations/bedrock](https://docs.raindrop.ai/integrations/bedrock)

## License

MIT
