TokenGuard

Production-ready token tracking, policy evaluation engines, budget limits, and alerts for LLM applications.

Quick Start

1. Direct Token Tracking & Policy Enforcement

Track exact token usage reported by LLM API responses and enforce rate-limiting policies:

from token_guard import TokenGuard, SlidingWindowPolicy

# Configure a Sliding Window Policy: max 50,000 tokens per hour
policy = SlidingWindowPolicy(limit=50_000, window=3600)
guard = TokenGuard(policy=policy)

# Record exact token usage for a user request
result = guard.track_usage("user_123", input_tokens=420, output_tokens=150)

print(result.total_tokens)                    # 570
print(result.limit_exceeded)                  # False
print(result.cumulative_usage.total_tokens)   # 570

2. Async Tracking for Web Frameworks (FastAPI / Starlette)

Use non-blocking async execution in web servers:

import asyncio
from token_guard import AsyncTokenGuard, AsyncTokenBucketPolicy

async def main():
    policy = AsyncTokenBucketPolicy(capacity=10_000, refill_rate=100.0)
    guard = AsyncTokenGuard(policy=policy)
    
    result = await guard.track_usage("user_456", input_tokens=80, output_tokens=20)
    print(result.total_tokens)   # 100
    print(result.limit_exceeded) # False

asyncio.run(main())

Why TokenGuard?

LLM API calls are billed per token (inputs + outputs). Uncontrolled application requests can quickly trigger unexpected cost spikes, upstream rate limit errors, or user abuse.

TokenGuard operates as a lightweight, thread-safe, and event-loop-safe middleware layer:

• Flexible Policy Engine: Enforce Sliding Window, Token Bucket, Fixed Window, Leaky Bucket, Cost ($/day), Quota, or Role-based rate limits.
• Cost & Budget Protection: Enforce usage caps per user, model, tier, or session.
• Unified Multi-Provider Support: Track token metrics across OpenAI, Groq, OpenRouter, and AWS Bedrock under a single unified API.
• Pluggable Enterprise Storage: Keep state in-memory during development or swap to Redis, SQLite, PostgreSQL, or AWS DynamoDB in production with one config change.
• Real-time Alerts: Dispatch warnings and webhook notifications (Slack, console, webhooks) the moment limit thresholds are crossed.

Features

Installation

# Core package (includes OpenAI/tiktoken local counting, policies, and memory storage)
pip install llm-token-guard

# Install optional storage backends & provider extras:
pip install "llm-token-guard[redis]"         # Redis storage support
pip install "llm-token-guard[postgres]"      # PostgreSQL support (psycopg, asyncpg)
pip install "llm-token-guard[dynamodb]"      # AWS DynamoDB support (boto3, aioboto3)
pip install "llm-token-guard[sqlite-async]"  # Async SQLite (aiosqlite) support
pip install "llm-token-guard[groq]"          # Groq HuggingFace tokenizers
pip install "llm-token-guard[bedrock]"       # AWS Bedrock CountTokens API
pip install "llm-token-guard[all]"           # All optional dependencies

Essential Usage Examples

1. Combining Multiple Policies

Combine rate limits, daily cost caps, and monthly quotas into a single TokenGuard instance:

from token_guard import TokenGuard, SlidingWindowPolicy, CostPolicy, QuotaPolicy

policies = [
    SlidingWindowPolicy(limit=10_000, window=3600),     # Max 10k tokens / hour
    CostPolicy(daily_limit_usd=5.0),                    # Max $5.00 cost / day
    QuotaPolicy(monthly_tokens=500_000),                # Max 500k tokens / month
]

guard = TokenGuard(policies=policies)
result = guard.track_usage("user_789", input_tokens=500, output_tokens=200)

if result.limit_exceeded:
    print(f"Request blocked: {result.policy_result.reason}")
    print(f"Retry after: {result.policy_result.retry_after}s")

2. Configuring Production Storage Backends

Initialize storage backends using StorageFactory or environment variables:

from token_guard import TokenGuard, StorageFactory, SlidingWindowPolicy

# Option A: Create storage from a DSN connection string (Postgres / Redis / SQLite)
storage = StorageFactory.create("postgres", connection_string="postgresql://user:pass@localhost:5432/db")

# Option B: Auto-create storage from environment variables (e.g. TOKEN_GUARD_STORAGE=redis)
# storage = StorageFactory.from_env()

guard = TokenGuard(
    policy=SlidingWindowPolicy(limit=50_000, window=3600),
    storage=storage,
)

3. Text Token Counting & Auto Provider Detection

Count tokens from text strings before executing LLM requests:

from token_guard import TokenGuard, CounterFactory

# Auto-detect token counter based on model name
counter = CounterFactory.auto("meta-llama/llama-3-70b-instruct")
guard = TokenGuard(counter=counter)

# Track usage from prompt and response text
result = guard.track(
    user_id="user_123",
    input_text="Explain quantum computing in simple terms.",
    output_text="Quantum computing is a field of computing focused on quantum physics principles..."
)

print(f"Input Tokens: {result.input_tokens}")
print(f"Output Tokens: {result.output_tokens}")

4. Setting Up Alert Handlers

Trigger warning alerts when users breach limits:

from token_guard import TokenGuard, BaseAlertHandler

class SlackWebhookAlertHandler(BaseAlertHandler):
    def alert(self, user_id: str, usage, max_tokens: int) -> None:
        print(f"[ALERT] User {user_id} exceeded token limit! Usage: {usage.total_tokens}")

guard = TokenGuard(
    max_tokens=1_000,
    alert_handlers=[SlackWebhookAlertHandler()]
)

Architecture & Execution Pipeline

TokenGuard follows a modular Strategy Pattern architecture, cleanly decoupling token counting, policy evaluation, storage persistence, and alert dispatching into independent layers:

graph TD
    classDef app fill:#1e1e2e,stroke:#74c7ec,stroke-width:2px,color:#cdd6f4
    classDef core fill:#313244,stroke:#cba6f7,stroke-width:2px,color:#cdd6f4
    classDef engine fill:#45475a,stroke:#f9e2af,stroke-width:2px,color:#cdd6f4
    classDef storage fill:#313244,stroke:#a6e3a1,stroke-width:2px,color:#cdd6f4
    classDef alert fill:#313244,stroke:#f38ba8,stroke-width:2px,color:#cdd6f4
    classDef decision fill:#181825,stroke:#fab387,stroke-width:2px,color:#cdd6f4

    App["💻 Application Server"]:::app --> TG["🛡️ TokenGuard / AsyncTokenGuard"]:::core
    
    subgraph Execution Pipeline
        TG --> Counter["1. Token Counter<br/>(OpenAI, Groq, Bedrock, Direct)"]:::core
        Counter --> Policy["2. Policy Evaluator<br/>(Sliding Window, Token Bucket, Cost, Quota, Role)"]:::engine
        Policy --> Decision{"Allowed?"}:::decision
    end

    Decision -- "YES (Allowed)" --> Storage["3. Storage Backend<br/>(Memory, Redis, SQLite, Postgres, DynamoDB)"]:::storage
    Decision -- "NO (Rejected)" --> Alerts["3. Alert Manager<br/>(Console, Slack, Webhooks)"]:::alert
    
    Storage --> Result["📦 TrackResult<br/>(Cumulative Usage, Policy Metadata)"]:::core
    Alerts --> Result
    Result --> App

1. Token Calculation: TokenGuard computes exact or estimated prompt and response tokens via the configured Token Counter.
2. Policy Evaluation: The request context (PolicyContext) is passed to the Policy Engine. Active policies are evaluated in order.
3. Decision & Execution:
   • Allowed: Token usage is atomically persisted in the Storage Backend, and an approved TrackResult is returned.
   • Rejected (Short-Circuit): Storage modification is skipped, configured Alert Handlers are triggered, and a rejected TrackResult (limit_exceeded=True) is returned with retry guidance.

Provider Compatibility

Documentation & Advanced Guides

For detailed setup instructions, configuration options, and advanced integrations, refer to the documentation guides:

• 📘 Policy Engine Guide — In-depth configurations for Sliding Window, Token Bucket, Cost, Quota, and Role-based policies.
• 📘 Token Counting & Providers Guide — Tokenizer selection and accuracy comparisons for OpenAI, Groq, OpenRouter, and AWS Bedrock.
• 📘 Storage Backends Guide — Production setup for Redis, SQLite, PostgreSQL, and AWS DynamoDB.
  • 📘 PostgreSQL Setup Guide — Sync (psycopg) & Async (asyncpg) PostgreSQL storage setup.
  • 📘 AWS DynamoDB Setup Guide — Serverless AWS DynamoDB storage setup.
• 📘 FastAPI Integration Guide — Adding AsyncTokenGuard middleware and non-blocking route tracking.
• 📘 Async Support Guide — Event loop integration and async storage drivers.
• 📘 Custom Backends Guide — Implementing custom counters, storage drivers, and policies.

Project Structure

token_guard/
├── docs/                 # Detailed guides and reference docs
│   └── storage/          # Storage-specific guides (PostgreSQL, DynamoDB)
├── token_guard/          # Core library source code
│   ├── counters/         # Token counters (OpenAI, Groq, Bedrock, etc.)
│   ├── engine/           # Policy evaluators and execution pipelines
│   ├── policies/         # Rate limiting, cost, quota, and role policies
│   └── storage/          # Storage backends (Memory, Redis, SQLite, Postgres, DynamoDB)
├── tests/                # Test suite (sync & async)
├── example_fastapi.py    # FastAPI integration demo
└── pyproject.toml        # Build configuration and dependencies

Runnable Examples

• All-Features Demo: Runnable test suite demonstrating sync/async features, storage drivers, and policies.
• FastAPI Integration Demo: Web server integration with async token limits and route handling.
• Multi-Provider Demo: Mapping multiple counter and storage backends.

Running Tests

pip install -e ".[dev]"

# Run all offline sync and async tests (no API keys required)
pytest tests/ -v

# Run provider integration tests (requires respective provider API keys)
export GROQ_API_KEY=gsk_...
pytest tests/test_groq_integration.py -v -s

export OPENAI_API_KEY=sk-...
pytest tests/test_openai_integration.py -v -s

export OPENROUTER_API_KEY=sk-or-v1-...
pytest tests/test_openrouter_integration.py -v -s

export AWS_ACCESS_KEY_ID=AKIA... AWS_SECRET_ACCESS_KEY=... AWS_DEFAULT_REGION=us-east-1
pytest tests/test_bedrock_integration.py -v -s

Roadmap

Completed Milestones

• Multi-provider token counting — OpenAI, Groq, OpenRouter, AWS Bedrock ✅
• Auto-detect provider — CounterFactory.auto() ✅
• Pluggable storage — Memory, Redis, SQLite ✅
• StorageFactory — from_env(), from_url(), from_config() ✅
• Redis connection pooling + TTL + from_url() + ping() ✅
• GitHub Actions CI/CD — auto-publish on version tag ✅
• Exact token tracking — track_usage() with API-reported counts ✅
• Async support — async def track(...) for non-blocking execution ✅
• Policy Engine (v0.5.0) — Sliding Window, Token Bucket, Fixed Window, Leaky Bucket, Cost, Quota, Role policies ✅
• PostgreSQL & DynamoDB Drivers (v0.6.0) — Built-in enterprise storage drivers ✅
• Hardening & Performance (v0.6.1) — Deadlock fixes, thread safety, memory eviction, double query optimization, custom exceptions ✅
• Provider Integrations & Python 3.13/3.14 (v0.6.2) — Integration tests for OpenAI/OpenRouter/Bedrock, Python 3.13 & 3.14 CI matrix, DX README refactor ✅

Upcoming Enterprise Roadmap

Phase 1: Distributed Storage-Backed Policies

• Redis Lua & SQL Window Counters — Execute rate-limiting counters directly in storage (Redis Lua scripts & Postgres SQL window queries) for cluster-wide rate limiting across pod workers.
• Budget Warning Thresholds — Fire warning alerts at configurable percentages (e.g. 80%) before hard limit rejection.

Phase 2: Enterprise Middleware & Webhook Handlers

• FastAPI & Starlette Middleware — Native drop-in TokenGuardMiddleware with automatic API key extraction, policy enforcement, and standardized HTTP 429 JSON responses with Retry-After headers.
• Enterprise Alert Handlers — Built-in SlackAlertHandler, WebhookAlertHandler, and PagerDutyAlertHandler.

Phase 3: Observability & Provider Expansion

• Prometheus Metrics Exporter — Expose token_guard_tokens_total counters and token_guard_policy_rejections_total histograms.
• OpenTelemetry Tracing — Automatic trace spans across counter calculation and policy evaluation.
• Vertex AI & Cohere Counters — Native exact token counting backends.

Contributing

Contributions are welcome! Please follow these basic guidelines:

1. Fork the repository and create a feature branch.
2. Ensure the full test suite passes locally before submitting your PR:

   pytest tests/ -v
   

3. Follow PEP 8 style standards.

License

MIT ©Abhijit Gunjal — see LICENSE for details.