tokenvault-sdk 0.1.1

Python SDK for AI agents to track LLM usage and forward billing data to TokenVault

TokenVault SDK

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:

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:

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

2. Basic Usage

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

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):

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:

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:

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

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

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

# Run tests
pytest

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

Building from Source

# 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
• Check the documentation
• Contact the TokenVault team

Contributing

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