tokenvault-sdk 0.1.3

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 manage billing with TokenVault.

Features

• 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

pip install tokenvault-sdk

Requirements

• Python 3.11+
• Redis server
• TokenVault Wallet API access

Quick Start

Set environment variables:

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

Basic usage:

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,
        wallet_api_url=config.wallet_api_url,
    )
    
    groq_client = AsyncGroq(api_key="your-api-key")
    
    try:
        # Check balance
        balance = await tracker.check_balance(organization_id="org_123")
        print(f"Balance: ${balance:.2f}")
        
        # 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_WALLET_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:

from tokenvault_sdk import SDKConfig, UsageTracker

config = SDKConfig(
    redis_url="redis://localhost:6379",
    wallet_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_balance(organization_id: str, force_refresh: bool = False) -> float

• Check organization balance with optional cache bypass
• 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

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}")

Modes

Fail-Open (default): Allows requests when services are unavailable

tracker = UsageTracker(..., fail_open=True)

Fail-Closed: Rejects requests when services are unavailable

tracker = UsageTracker(..., fail_open=False)

License

MIT License - see LICENSE file for details