raindrop-bedrock 0.0.1__tar.gz

raindrop_bedrock-0.0.1/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.1
+Name: raindrop-bedrock
+Version: 0.0.1
+Summary: Raindrop integration for AWS Bedrock
+License: MIT
+Author: Raindrop AI
+Author-email: sdk@raindrop.ai
+Requires-Python: >=3.10,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: boto3 (>=1.34.0)
+Requires-Dist: raindrop-ai (>=0.0.42)
+Description-Content-Type: text/markdown
+
+# raindrop-bedrock
+
+Raindrop 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 boto3
+```
+
+## Usage
+
+```python
+import boto3
+from raindrop_bedrock import create_raindrop_bedrock
+
+raindrop = create_raindrop_bedrock(
+    api_key="rk_...",
+    user_id="user-123",
+)
+
+client = boto3.client("bedrock-runtime", region_name="us-east-1")
+wrapped = raindrop["wrap"](client)
+
+response = wrapped.converse(
+    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
+    messages=[
+        {"role": "user", "content": [{"text": "Hello!"}]},
+    ],
+)
+
+raindrop["flush"]()
+```
+
+## What gets captured
+
+- **converse()**: input messages, output text, model ID, token usage (inputTokens/outputTokens)
+- **invoke_model()**: raw request/response bodies, model ID, token usage (Claude, Titan, and Llama formats)
+- **Errors**: tracked and re-raised to the caller
+
+## Options
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `api_key` | `str` | required | Raindrop API key |
+| `user_id` | `str` | `None` | Associate all events with a user |
+| `convo_id` | `str` | `None` | Group events into a conversation |
+
+## Testing
+
+```bash
+cd packages/bedrock-python
+pip install -e .
+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.
+- **No `events.*`, `users.*`, or `signals.*` APIs** — Python SDK limitation. Use `raindrop.analytics` directly for these features.
+

raindrop_bedrock-0.0.1/README.md

@@ -0,0 +1,60 @@
+# raindrop-bedrock
+
+Raindrop 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 boto3
+```
+
+## Usage
+
+```python
+import boto3
+from raindrop_bedrock import create_raindrop_bedrock
+
+raindrop = create_raindrop_bedrock(
+    api_key="rk_...",
+    user_id="user-123",
+)
+
+client = boto3.client("bedrock-runtime", region_name="us-east-1")
+wrapped = raindrop["wrap"](client)
+
+response = wrapped.converse(
+    modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
+    messages=[
+        {"role": "user", "content": [{"text": "Hello!"}]},
+    ],
+)
+
+raindrop["flush"]()
+```
+
+## What gets captured
+
+- **converse()**: input messages, output text, model ID, token usage (inputTokens/outputTokens)
+- **invoke_model()**: raw request/response bodies, model ID, token usage (Claude, Titan, and Llama formats)
+- **Errors**: tracked and re-raised to the caller
+
+## Options
+
+| Option | Type | Default | Description |
+|--------|------|---------|-------------|
+| `api_key` | `str` | required | Raindrop API key |
+| `user_id` | `str` | `None` | Associate all events with a user |
+| `convo_id` | `str` | `None` | Group events into a conversation |
+
+## Testing
+
+```bash
+cd packages/bedrock-python
+pip install -e .
+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.
+- **No `events.*`, `users.*`, or `signals.*` APIs** — Python SDK limitation. Use `raindrop.analytics` directly for these features.

raindrop_bedrock-0.0.1/pyproject.toml

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "raindrop-bedrock"
+version = "0.0.1"
+description = "Raindrop integration for AWS Bedrock"
+authors = ["Raindrop AI <sdk@raindrop.ai>"]
+license = "MIT"
+readme = "README.md"
+packages = [{ include = "raindrop_bedrock" }]
+
+[tool.poetry.dependencies]
+python = ">=3.10,<4.0"
+raindrop-ai = ">=0.0.42"
+boto3 = ">=1.34.0"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.0"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = ["-v", "--tb=short"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

raindrop_bedrock-0.0.1/raindrop_bedrock/__init__.py

@@ -0,0 +1,3 @@
+from .wrapper import create_raindrop_bedrock
+
+__all__ = ["create_raindrop_bedrock"]

raindrop_bedrock-0.0.1/raindrop_bedrock/wrapper.py

@@ -0,0 +1,292 @@
+"""
+AWS Bedrock client wrapper for Raindrop observability.
+
+Wraps a boto3 bedrock-runtime client to automatically capture
+converse() and invoke_model() calls and ship them to Raindrop.
+
+Usage:
+    from raindrop_bedrock import create_raindrop_bedrock
+
+    raindrop = create_raindrop_bedrock(api_key="rk_...")
+    bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
+    wrapped = raindrop["wrap"](bedrock)
+
+    response = wrapped.converse(modelId="anthropic.claude-3-5-sonnet-20241022-v2:0", ...)
+    raindrop["flush"]()
+"""
+
+from __future__ import annotations
+
+import json
+import uuid
+from typing import Any, Dict, Optional
+
+import raindrop.analytics as raindrop
+
+def create_raindrop_bedrock(
+    api_key: str,
+    user_id: Optional[str] = None,
+    convo_id: Optional[str] = None,
+) -> dict:
+    """
+    Create a Raindrop-instrumented Bedrock wrapper.
+
+    Args:
+        api_key: Raindrop API key (rk_...).
+        user_id: Associate all events with this user.
+        convo_id: Conversation ID to group related events.
+
+    Returns:
+        Dict with 'wrap', 'flush', and 'shutdown'.
+
+    Usage:
+        raindrop = create_raindrop_bedrock(api_key="rk_...")
+        bedrock = boto3.client("bedrock-runtime", region_name="us-east-1")
+        wrapped = raindrop["wrap"](bedrock)
+        response = wrapped.converse(modelId="...", messages=[...])
+        raindrop["flush"]()
+    """
+    raindrop.init(api_key=api_key)
+
+    def wrap(client: Any) -> Any:
+        return wrap_bedrock_client(client, user_id=user_id, convo_id=convo_id)
+
+    return {
+        "wrap": wrap,
+        "flush": raindrop.flush,
+        "shutdown": raindrop.shutdown,
+    }
+
+
+def wrap_bedrock_client(
+    client: Any,
+    user_id: Optional[str] = None,
+    convo_id: Optional[str] = None,
+) -> Any:
+    """Wrap a boto3 bedrock-runtime client to capture LLM calls."""
+
+    original_converse = getattr(client, "converse", None)
+    original_invoke_model = getattr(client, "invoke_model", None)
+
+    if original_converse:
+        def wrapped_converse(**kwargs: Any) -> Any:
+            return _instrument_converse(original_converse, kwargs, user_id, convo_id)
+        client.converse = wrapped_converse
+
+    if original_invoke_model:
+        def wrapped_invoke_model(**kwargs: Any) -> Any:
+            return _instrument_invoke_model(original_invoke_model, kwargs, user_id, convo_id)
+        client.invoke_model = wrapped_invoke_model
+
+    return client
+
+
+def _instrument_converse(
+    original_fn: Any,
+    kwargs: Dict[str, Any],
+    user_id: Optional[str],
+    convo_id: Optional[str],
+) -> Any:
+    try:
+        model_id = kwargs.get("modelId")
+        input_text = _format_converse_input(kwargs.get("messages"))
+        event_id = str(uuid.uuid4())
+    except Exception:
+        return original_fn(**kwargs)
+
+    try:
+        response = original_fn(**kwargs)
+    except Exception:
+        try:
+            raindrop.track_ai(
+                user_id=user_id or "unknown",
+                event="ai_generation",
+                event_id=event_id,
+                input=input_text,
+                model=model_id,
+                convo_id=convo_id,
+            )
+        except Exception:
+            pass
+        raise
+
+    try:
+        output_text = _extract_converse_output(response)
+        usage = response.get("usage", {})
+        prompt_tokens = usage.get("inputTokens")
+        completion_tokens = usage.get("outputTokens")
+
+        properties: Dict[str, Any] = {}
+        if prompt_tokens is not None:
+            properties["ai.usage.prompt_tokens"] = prompt_tokens
+        if completion_tokens is not None:
+            properties["ai.usage.completion_tokens"] = completion_tokens
+
+        raindrop.track_ai(
+            user_id=user_id or "unknown",
+            event="ai_generation",
+            event_id=event_id,
+            input=input_text,
+            output=output_text,
+            model=model_id,
+            properties=properties or None,
+            convo_id=convo_id,
+        )
+    except Exception:
+        pass
+
+    return response
+
+
+def _instrument_invoke_model(
+    original_fn: Any,
+    kwargs: Dict[str, Any],
+    user_id: Optional[str],
+    convo_id: Optional[str],
+) -> Any:
+    try:
+        model_id = kwargs.get("modelId")
+        input_text = _parse_body(kwargs.get("body"))
+        event_id = str(uuid.uuid4())
+    except Exception:
+        return original_fn(**kwargs)
+
+    try:
+        response = original_fn(**kwargs)
+    except Exception:
+        try:
+            raindrop.track_ai(
+                user_id=user_id or "unknown",
+                event="ai_generation",
+                event_id=event_id,
+                input=input_text,
+                model=model_id,
+                convo_id=convo_id,
+            )
+        except Exception:
+            pass
+        raise
+
+    try:
+        output_text, prompt_tokens, completion_tokens = _parse_invoke_response(response)
+
+        properties: Dict[str, Any] = {}
+        if prompt_tokens is not None:
+            properties["ai.usage.prompt_tokens"] = prompt_tokens
+        if completion_tokens is not None:
+            properties["ai.usage.completion_tokens"] = completion_tokens
+
+        raindrop.track_ai(
+            user_id=user_id or "unknown",
+            event="ai_generation",
+            event_id=event_id,
+            input=input_text,
+            output=output_text,
+            model=model_id,
+            properties=properties or None,
+            convo_id=convo_id,
+        )
+    except Exception:
+        pass
+
+    return response
+
+
+def _format_converse_input(messages: Any) -> Optional[str]:
+    if not messages:
+        return None
+    user_msgs = [m for m in messages if isinstance(m, dict) and m.get("role") == "user"]
+    to_extract = user_msgs if user_msgs else messages[-1:]
+    parts = []
+    for msg in to_extract:
+        content = msg.get("content", []) if isinstance(msg, dict) else []
+        text_parts = []
+        for block in (content if isinstance(content, list) else []):
+            if isinstance(block, dict) and block.get("text"):
+                text_parts.append(block["text"])
+        if text_parts:
+            parts.append(" ".join(text_parts))
+    return "\n".join(parts) if parts else None
+
+
+def _extract_converse_output(response: Dict[str, Any]) -> Optional[str]:
+    output = response.get("output", {})
+    message = output.get("message", {}) if isinstance(output, dict) else {}
+    content = message.get("content", []) if isinstance(message, dict) else []
+    texts = []
+    for block in (content if isinstance(content, list) else []):
+        if isinstance(block, dict) and "text" in block:
+            texts.append(block["text"])
+    return "\n".join(texts) if texts else None
+
+
+def _parse_body(body: Any) -> Optional[str]:
+    if body is None:
+        return None
+    try:
+        if isinstance(body, (bytes, bytearray)):
+            return body.decode("utf-8")
+        if isinstance(body, str):
+            return body
+        if hasattr(body, "read"):
+            data = body.read()
+            if hasattr(body, "seek"):
+                body.seek(0)
+            return data.decode("utf-8") if isinstance(data, bytes) else str(data)
+        return str(body)
+    except Exception:
+        return None
+
+
+def _parse_invoke_response(
+    response: Dict[str, Any],
+) -> tuple:
+    """Returns (output_text, prompt_tokens, completion_tokens).
+
+    If the body is a streaming object (like botocore.response.StreamingBody),
+    it is read and replaced with a BytesIO so the caller can still call .read().
+    """
+    body = response.get("body")
+    if not body:
+        return None, None, None
+
+    try:
+        if hasattr(body, "read"):
+            raw = body.read()
+            import io
+            response["body"] = io.BytesIO(raw)
+        elif isinstance(body, (bytes, bytearray)):
+            raw = body
+        else:
+            raw = str(body).encode("utf-8")
+
+        parsed = json.loads(raw)
+
+        # Anthropic Claude format
+        content = parsed.get("content")
+        if isinstance(content, list):
+            output = "\n".join(
+                c.get("text", "") for c in content if isinstance(c, dict) and c.get("type") == "text"
+            ) or None
+            usage = parsed.get("usage", {})
+            return output, usage.get("input_tokens"), usage.get("output_tokens")
+
+        # Amazon Titan format
+        results = parsed.get("results")
+        if isinstance(results, list) and results:
+            output = "\n".join(
+                r.get("outputText", "") for r in results if isinstance(r, dict)
+            ) or None
+            return output, parsed.get("inputTextTokenCount"), results[0].get("tokenCount") if isinstance(results[0], dict) else None
+
+        # Llama format
+        generation = parsed.get("generation")
+        if generation:
+            return generation, parsed.get("prompt_token_count"), parsed.get("generation_token_count")
+
+        # Fallback
+        output = parsed.get("output") or parsed.get("completion") or parsed.get("generated_text")
+        return output, None, None
+
+    except Exception:
+        return None, None, None