# TokenVault SDK

[![PyPI version](https://badge.fury.io/py/tokenvault-sdk.svg)](https://badge.fury.io/py/tokenvault-sdk)
[![Python versions](https://img.shields.io/pypi/pyversions/tokenvault-sdk.svg)](https://pypi.org/project/tokenvault-sdk/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

A lightweight Python SDK for AI agents to track LLM usage and manage billing with TokenVault.

## Features

- Credit balance checking with Redis caching
- Usage tracking and forwarding to TokenVault
- Async/await support with automatic retries
- Fail-open mode for resilience
- Groq API integration

## Installation

```bash
pip install tokenvault-sdk
```

## Requirements

- Python 3.11+
- Redis server
- TokenVault API access

## Quick Start

Set environment variables:

```bash
export TOKENVAULT_REDIS_URL="redis://localhost:6379"
export TOKENVAULT_API_URL="https://api.tokenvault.io"
```

Basic usage:

```python
import asyncio
from tokenvault_sdk import SDKConfig, UsageTracker
from groq import AsyncGroq

async def main():
    config = SDKConfig.from_env()
    tracker = UsageTracker(
        redis_url=config.redis_url,
        tokenvault_api_url=config.tokenvault_api_url,
    )
    
    groq_client = AsyncGroq(api_key="your-api-key")
    
    try:
        # Check credit balance
        credit_balance = await tracker.check_credit_balance(organization_id="org_123")
        print(f"Credit balance: {credit_balance:.2f} credits")
        
        # Make LLM call
        response = await groq_client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[{"role": "user", "content": "Hello"}]
        )
        
        # Track usage (fire-and-forget)
        asyncio.create_task(
            tracker.track_usage(
                groq_response=response.model_dump(),
                organization_id="org_123",
                profile_id="agent_456",
            )
        )
        
        print(response.choices[0].message.content)
    finally:
        await tracker.close()

asyncio.run(main())
```

## Configuration

Environment variables:

| Variable | Required | Default |
|----------|----------|---------|
| `TOKENVAULT_REDIS_URL` | Yes | - |
| `TOKENVAULT_API_URL` | Yes | - |
| `TOKENVAULT_STREAM_NAME` | No | `usage_records` |
| `TOKENVAULT_BALANCE_THRESHOLD` | No | `0.0` |
| `TOKENVAULT_CACHE_TTL` | No | `60` |
| `TOKENVAULT_FAIL_OPEN` | No | `true` |
| `TOKENVAULT_MAX_RETRIES` | No | `3` |
| `TOKENVAULT_RETRY_BACKOFF_MS` | No | `100` |

Or configure programmatically:

```python
from tokenvault_sdk import SDKConfig, UsageTracker

config = SDKConfig(
    redis_url="redis://localhost:6379",
    tokenvault_api_url="https://api.tokenvault.io",
    balance_threshold=0.0,
    cache_ttl=60,
    fail_open=True,
)

tracker = UsageTracker(
    redis_url=config.redis_url,
    wallet_api_url=config.wallet_api_url,
    balance_threshold=config.balance_threshold,
    cache_ttl=config.cache_ttl,
    fail_open=config.fail_open,
)
```

## API Reference

### UsageTracker

**Methods:**

`check_credit_balance(organization_id: str, force_refresh: bool = False) -> float`
- Check organization credit balance with optional cache bypass
- Returns credit balance as float
- Raises `InsufficientBalanceError` if below threshold

`track_usage(groq_response: Dict, organization_id: str, profile_id: str, user_id: str = "", metadata: Optional[Dict] = None)`
- Track LLM usage and publish to Redis Stream
- Designed for fire-and-forget usage with `asyncio.create_task()`

`close()`
- Close all connections gracefully

### Exceptions

- `SDKConfigurationError`: Invalid configuration
- `InsufficientBalanceError`: Balance below threshold
- `TrackingError`: Usage tracking failed (fail-closed mode)

## Error Handling

```python
from tokenvault_sdk import (
    UsageTracker,
    InsufficientBalanceError,
    SDKConfigurationError,
)

try:
    config = SDKConfig.from_env()
except SDKConfigurationError as e:
    print(f"Configuration error: {e}")
    exit(1)

tracker = UsageTracker(
    redis_url=config.redis_url,
    tokenvault_api_url=config.tokenvault_api_url,
)

try:
    credit_balance = await tracker.check_credit_balance("org_123")
except InsufficientBalanceError as e:
    print(f"Insufficient credit balance: {e}")
```

## Modes

**Fail-Open (default):** Allows requests when services are unavailable
```python
tracker = UsageTracker(..., fail_open=True)
```

**Fail-Closed:** Rejects requests when services are unavailable
```python
tracker = UsageTracker(..., fail_open=False)
```

## License

MIT License - see LICENSE file for details
