# Revenium Universal Griptape Driver

[![PyPI version](https://badge.fury.io/py/revenium-griptape.svg)](https://badge.fury.io/py/revenium-griptape)
[![Python Versions](https://img.shields.io/pypi/pyversions/revenium-griptape.svg)](https://pypi.org/project/revenium-griptape/)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)

Add automatic AI usage metering and cost tracking to your Griptape applications with **zero code changes**. Works with **any LLM provider** that Griptape supports - simply wrap your existing driver and get detailed analytics via Revenium.

## What This Does

This package provides a universal wrapper for any Griptape prompt driver that automatically:
- **Meters all AI usage** across any provider and sends data to Revenium for analytics
- **Tracks costs and token usage** with provider-specific optimizations  
- **Maintains full compatibility** with existing Griptape code
- **Supports rich metadata** for detailed reporting and cost allocation
- **Auto-detects providers** and uses the best available middleware

## Provider Support Architecture

### Tier 1: Direct Native Support ⚡
*Best performance with optimized provider-specific middleware*
- **OpenAI** → `revenium-middleware-openai`
- **Anthropic** → `revenium-middleware-anthropic` 
- **Ollama** → `revenium-middleware-ollama`

### Tier 2: Universal LiteLLM Support 🌐
*Supports 100+ providers through LiteLLM integration*
- **Google/Gemini** → via `revenium-middleware-litellm`
- **Cohere** → via `revenium-middleware-litellm`
- **Azure OpenAI** → via `revenium-middleware-litellm`
- **AWS Bedrock** → via `revenium-middleware-litellm`
- **LiteLLM Proxy** → via `revenium-middleware-litellm` (centralized proxy server)
- **And 100+ more** → [See LiteLLM docs](https://docs.litellm.ai/docs/providers)

## How It Works

1. Wrap any Griptape driver with `ReveniumDriver`
2. Add optional metadata for tracking
3. Your AI calls automatically get metered and reported to Revenium
4. View analytics and cost data in your Revenium dashboard

## Installation

```bash
# Base package
pip install revenium-griptape

# Install middleware for your providers
pip install revenium-middleware-openai      # For OpenAI (Tier 1)
pip install revenium-middleware-anthropic   # For Anthropic (Tier 1)  
pip install revenium-middleware-ollama      # For Ollama (Tier 1)
pip install revenium-middleware-litellm     # For all other providers (Tier 2)
```

## Environment Setup

Create a `.env` file in your project:

```bash
# Required: Your Revenium API key for metering
REVENIUM_METERING_API_KEY=your_revenium_api_key_here

# Optional: Custom Revenium endpoint (defaults to production)
REVENIUM_METERING_BASE_URL=https://api.revenium.io/meter

# Tier 1 Provider API Keys (Direct Support)
OPENAI_API_KEY=your_openai_api_key_here
ANTHROPIC_API_KEY=your_anthropic_api_key_here
# OLLAMA_HOST=http://localhost:11434  # For local Ollama

# Tier 2 Provider API Keys (via LiteLLM - as needed)
GEMINI_API_KEY=your_google_gemini_api_key_here
COHERE_API_KEY=your_cohere_api_key_here
AWS_ACCESS_KEY_ID=your_aws_access_key_here        # For Bedrock
AWS_SECRET_ACCESS_KEY=your_aws_secret_key_here    # For Bedrock
AZURE_API_KEY=your_azure_openai_api_key_here      # For Azure OpenAI

# LiteLLM Proxy Configuration (optional - for centralized proxy server)
LITELLM_PROXY_URL=http://localhost:4000/chat/completions
LITELLM_API_KEY=sk-1234

# See LiteLLM docs for other provider environment variables
```

> **Tier 2 Environment Variables**: For Tier 2 providers, environment variable names follow [LiteLLM conventions](https://docs.litellm.ai/docs/providers). Each provider has specific variable names (e.g., `GEMINI_API_KEY` for Google, `COHERE_API_KEY` for Cohere). See the [LiteLLM provider documentation](https://docs.litellm.ai/docs/providers) for the complete list.

## Quick Start

### Universal Driver (Recommended)

```python
import os
from dotenv import load_dotenv

# ✅ CRITICAL: Load environment variables BEFORE importing
load_dotenv()

from griptape.structures import Agent
from griptape.drivers.prompt.openai_chat_prompt_driver import OpenAiChatPromptDriver

# Universal driver works with ANY Griptape prompt driver
from revenium_griptape import ReveniumDriver

# Method 1: Wrap an existing Griptape driver
base_driver = OpenAiChatPromptDriver(
    model="gpt-4o-mini",
    api_key=os.getenv("OPENAI_API_KEY")
)

driver = ReveniumDriver(
    base_driver=base_driver,
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "user_api_key",
                "value": "user_key_value"
            }
        },
        "task_type": "qa"
    }
)

# Method 2: Auto-detect from model name
driver = ReveniumDriver(
    model="gpt-4o-mini",  # Auto-detects OpenAI
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "user_api_key",
                "value": "user_key_value"
            }
        },
        "task_type": "qa"
    }
)

# Use exactly like any Griptape driver - metering is automatic!
agent = Agent(prompt_driver=driver)
result = agent.run("What is the capital of France?")
print(result.output.value)
```

### Direct Driver Usage

```python
from revenium_griptape import ReveniumOpenAiDriver, ReveniumAnthropicDriver

# OpenAI (Direct)
openai_driver = ReveniumOpenAiDriver(
    model="gpt-4o-mini",
    api_key=os.getenv("OPENAI_API_KEY"),
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "openai_api_key",
                "value": "openai_key_value"
            }
        }
    }
)

# Anthropic (Direct)
anthropic_driver = ReveniumAnthropicDriver(
    model="claude-3-haiku-20240307",
    api_key=os.getenv("ANTHROPIC_API_KEY"),
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "anthropic_api_key",
                "value": "anthropic_key_value"
            }
        }
    }
)

# LiteLLM for 100+ providers
from revenium_griptape import ReveniumLiteLLMDriver

# Google Gemini via LiteLLM
gemini_driver = ReveniumLiteLLMDriver(
    model="gemini-pro",
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "gemini_api_key",
                "value": "gemini_key_value"
            }
        }
    }
)

# LiteLLM Proxy
proxy_driver = ReveniumLiteLLMDriver(
    model="gpt-4",  # Model available through your proxy
    proxy_url=os.getenv("LITELLM_PROXY_URL"),
    proxy_api_key=os.getenv("LITELLM_API_KEY"),
    usage_metadata={
        "subscriber": {
            "id": "user-123",
            "email": "user@company.com",
            "credential": {
                "name": "proxy_api_key",
                "value": "proxy_key_value"
            }
        }
    }
)
```

## Supported Metadata Fields

Add any of these optional fields to track usage across different dimensions:

| Field | Description | Example |
|-------|-------------|---------|
| `trace_id` | Session or conversation ID | `"chat-session-123"` |
| `task_type` | Type of AI task | `"summarization"`, `"qa"` |
| `subscriber` | User information object | `{"id": "user-456", "email": "user@company.com"}` |
| `subscriber.id` | User ID | `"user-456"` |
| `subscriber.email` | User email | `"user@company.com"` |
| `subscriber.credential` | API credential info | `{"name": "api_key", "value": "key_value"}` |
| `organization_id` | Team or department | `"sales-team"` |
| `subscription_id` | Billing plan | `"enterprise-plan"` |
| `product_id` | Product or feature | `"ai-assistant"` |
| `agent` | AI agent identifier | `"support-bot-v2"` |


## Examples

**[View Complete Examples →](examples/)**

| Example | Description |
|---------|-------------|
| **[Universal Driver](examples/universal_example.py)** | **Recommended** - Multi-provider with auto-detection |
| **[OpenAI Direct](examples/openai_example.py)** | Direct OpenAI integration |
| **[Anthropic Direct](examples/anthropic_example.py)** | Direct Anthropic/Claude integration |
| **[LiteLLM Client](examples/litellm_direct_example.py)** | LiteLLM direct client integration |
| **[LiteLLM Proxy](examples/litellm_proxy_example.py)** | LiteLLM proxy server integration |

Each example is self-contained with environment setup, error handling, and detailed explanations.

**Quick Start**: Run `python examples/universal_example.py` after setting up your environment variables.
