Metadata-Version: 2.4
Name: tokenvault-sdk
Version: 0.1.1
Summary: Python SDK for AI agents to track LLM usage and forward billing data to TokenVault
Author: TokenVault Team
License: MIT
Project-URL: Homepage, https://github.com/tokenvault/tokenvault-sdk
Project-URL: Documentation, https://docs.tokenvault.io/sdk
Project-URL: Repository, https://github.com/tokenvault/tokenvault-sdk
Project-URL: Issues, https://github.com/tokenvault/tokenvault-sdk/issues
Keywords: llm,usage-tracking,billing,ai-agents,groq
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.11
Description-Content-Type: text/markdown
Requires-Dist: redis>=5.0.0
Requires-Dist: httpx>=0.25.0
Requires-Dist: tenacity>=8.2.0
Requires-Dist: groq>=0.4.0
Provides-Extra: dev
Requires-Dist: pytest>=7.4.0; extra == "dev"
Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
Requires-Dist: hypothesis>=6.92.0; extra == "dev"
Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
Requires-Dist: pytest-mock>=3.12.0; extra == "dev"

# 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 forward billing data to TokenVault. Provides simple functions for balance checking and usage tracking without wrapping your LLM client.

## Features

- **Balance Checking**: Check organization balance before making LLM API calls
- **Usage Tracking**: Automatically track and forward LLM usage data to TokenVault
- **Redis Caching**: Fast balance lookups with configurable TTL
- **Fail-Open Mode**: Continue operations even when services are temporarily unavailable
- **Async Support**: Built with asyncio for high-performance concurrent operations
- **Retry Logic**: Automatic retries with exponential backoff for resilience
- **Groq Integration**: Native support for Groq API responses

## Installation

Install the SDK using pip:

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

## Requirements

- Python 3.11 or higher
- Redis server (for caching and usage record streaming)
- TokenVault Wallet API access

## Quick Start

### 1. Configuration

Set up your environment variables:

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

### 2. Basic Usage

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

async def main():
    # Initialize SDK configuration from environment variables
    config = SDKConfig.from_env()
    
    # Create usage tracker
    tracker = UsageTracker(
        redis_url=config.redis_url,
        wallet_api_url=config.wallet_api_url,
        balance_threshold=0.0,
        cache_ttl=60,
        fail_open=True,
    )
    
    # Initialize your Groq client
    groq_client = AsyncGroq(api_key="your-groq-api-key")
    
    try:
        # Check balance before making LLM call
        balance = await tracker.check_balance(
            organization_id="org_123"
        )
        print(f"Current balance: ${balance:.2f}")
        
        # Make your LLM API call
        response = await groq_client.chat.completions.create(
            model="llama-3.1-70b-versatile",
            messages=[
                {"role": "user", "content": "Hello, how are you?"}
            ]
        )
        
        # Track usage (fire-and-forget)
        asyncio.create_task(
            tracker.track_usage(
                groq_response=response.model_dump(),
                organization_id="org_123",
                profile_id="agent_456",
                user_id="user_789",
            )
        )
        
        # Use the response
        print(response.choices[0].message.content)
        
    finally:
        # Clean up connections
        await tracker.close()

if __name__ == "__main__":
    asyncio.run(main())
```

## Configuration Options

The SDK can be configured via environment variables or directly in code:

### Environment Variables

| Variable | Required | Default | Description |
|----------|----------|---------|-------------|
| `TOKENVAULT_REDIS_URL` | Yes | - | Redis connection URL |
| `TOKENVAULT_WALLET_API_URL` | Yes | - | TokenVault Wallet API base URL |
| `TOKENVAULT_STREAM_NAME` | No | `usage_records` | Redis stream name for usage records |
| `TOKENVAULT_BALANCE_THRESHOLD` | No | `0.0` | Minimum balance required for requests |
| `TOKENVAULT_CACHE_TTL` | No | `60` | Balance cache TTL in seconds |
| `TOKENVAULT_FAIL_OPEN` | No | `true` | Allow requests when services unavailable |
| `TOKENVAULT_MAX_RETRIES` | No | `3` | Maximum retry attempts |
| `TOKENVAULT_RETRY_BACKOFF_MS` | No | `100` | Initial retry backoff in milliseconds |

### Programmatic Configuration

```python
from tokenvault_sdk import SDKConfig, UsageTracker

# Create config from environment
config = SDKConfig.from_env()

# Or create config manually
config = SDKConfig(
    redis_url="redis://localhost:6379",
    wallet_api_url="https://api.tokenvault.io",
    stream_name="usage_records",
    balance_threshold=0.0,
    cache_ttl=60,
    fail_open=True,
    max_retries=3,
    retry_backoff_ms=100,
)

# Initialize tracker with config
tracker = UsageTracker(
    redis_url=config.redis_url,
    wallet_api_url=config.wallet_api_url,
    stream_name=config.stream_name,
    balance_threshold=config.balance_threshold,
    cache_ttl=config.cache_ttl,
    fail_open=config.fail_open,
    max_retries=config.max_retries,
    retry_backoff_ms=config.retry_backoff_ms,
)
```

## API Reference

### SDKConfig

Configuration class for TokenVault SDK.

**Methods:**
- `from_env(prefix="TOKENVAULT_")`: Create configuration from environment variables
- `validate()`: Validate configuration values

### UsageTracker

Core component for balance checking and usage tracking.

**Methods:**

#### `check_balance(organization_id: str, force_refresh: bool = False) -> float`

Check organization balance with Redis caching.

**Parameters:**
- `organization_id`: Organization identifier
- `force_refresh`: Force fetch from Wallet API, bypassing cache

**Returns:** Current balance as float

**Raises:** `InsufficientBalanceError` if balance is below threshold

#### `track_usage(groq_response: Dict, organization_id: str, profile_id: str, user_id: str = "", metadata: Optional[Dict] = None) -> None`

Track LLM usage and publish to Redis Stream.

**Parameters:**
- `groq_response`: Response from Groq API (as dictionary)
- `organization_id`: Organization identifier
- `profile_id`: Agent/profile identifier
- `user_id`: User identifier (optional)
- `metadata`: Additional metadata (optional)

**Note:** This method is designed to be called in a fire-and-forget manner using `asyncio.create_task()`.

#### `close() -> None`

Close all connections gracefully. Should be called when shutting down.

### Exceptions

- `SDKConfigurationError`: Raised when SDK configuration is invalid
- `InsufficientBalanceError`: Raised when balance is below threshold
- `TrackingError`: Raised when usage tracking fails (in fail-closed mode)

## Advanced Usage

### Fail-Open vs Fail-Closed Mode

The SDK supports two operational modes:

**Fail-Open Mode (default):**
```python
tracker = UsageTracker(
    redis_url=config.redis_url,
    wallet_api_url=config.wallet_api_url,
    fail_open=True,  # Allow requests when services unavailable
)
```

In fail-open mode, the SDK allows requests to proceed even when Redis or the Wallet API are temporarily unavailable. This ensures your agent remains operational during service disruptions.

**Fail-Closed Mode:**
```python
tracker = UsageTracker(
    redis_url=config.redis_url,
    wallet_api_url=config.wallet_api_url,
    fail_open=False,  # Reject requests when services unavailable
)
```

In fail-closed mode, the SDK raises exceptions when services are unavailable, ensuring strict balance enforcement.

### Custom Logging

The SDK uses Python's standard logging module. Configure logging to control output:

```python
from tokenvault_sdk import configure_logging

# Configure SDK logging
configure_logging(level="INFO")

# Or use standard logging configuration
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("tokenvault_sdk")
```

### 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,
    wallet_api_url=config.wallet_api_url,
)

try:
    balance = await tracker.check_balance("org_123")
except InsufficientBalanceError as e:
    print(f"Insufficient balance: {e}")
    # Handle insufficient balance (e.g., notify user, skip request)
```

## Development

### Running Tests

```bash
# Install development dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Run tests with coverage
pytest --cov=tokenvault_sdk --cov-report=html
```

### Building from Source

```bash
# Clone the repository
git clone https://github.com/tokenvault/tokenvault-sdk.git
cd tokenvault-sdk

# Install in development mode
pip install -e ".[dev]"
```

## Links

- **Homepage**: https://github.com/tokenvault/tokenvault-sdk
- **Documentation**: https://docs.tokenvault.io/sdk
- **PyPI**: https://pypi.org/project/tokenvault-sdk/
- **Issue Tracker**: https://github.com/tokenvault/tokenvault-sdk/issues

## License

This project is licensed under the MIT License - see the LICENSE file for details.

## Support

For questions, issues, or feature requests, please:
- Open an issue on [GitHub](https://github.com/tokenvault/tokenvault-sdk/issues)
- Check the [documentation](https://docs.tokenvault.io/sdk)
- Contact the TokenVault team

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.
