Metadata-Version: 2.4
Name: tokenwarden
Version: 0.1.0
Summary: Drop-in spend governor for Python LLM workloads — cost attribution, budget enforcement, and OpenTelemetry signals via a single httpx transport.
Author-email: Venkat <venkat.polur@gmail.com>
License: MIT
License-File: LICENSE
Keywords: anthropic,budget,cost,llm,openai,opentelemetry
Requires-Python: >=3.10
Requires-Dist: httpx>=0.25.0
Provides-Extra: all
Requires-Dist: opentelemetry-api>=1.20.0; extra == 'all'
Requires-Dist: redis>=5.0.0; extra == 'all'
Provides-Extra: dev
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: respx>=0.21; extra == 'dev'
Provides-Extra: otel
Requires-Dist: opentelemetry-api>=1.20.0; extra == 'otel'
Provides-Extra: redis
Requires-Dist: redis>=5.0.0; extra == 'redis'
Description-Content-Type: text/markdown

# TokenWarden

**Drop-in spend governor for Python LLM workloads.**  
Real-time cost attribution, budget enforcement, and OpenTelemetry signals — via a single `httpx` transport.

[![PyPI](https://img.shields.io/pypi/v/tokenwarden)](https://pypi.org/project/tokenwarden)
[![Python](https://img.shields.io/pypi/pyversions/tokenwarden)](https://pypi.org/project/tokenwarden)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)

---

## What it does

Every HTTP call your application makes to an LLM provider is intercepted by TokenWarden's `httpx` transport. For each request it:

1. Detects the provider (OpenAI, Anthropic, Gemini, xAI, DeepSeek)
2. Resolves attribution context (tenant, user, feature) from Python context variables
3. Checks current budget spend — returns a `429` response if an **Enforce** budget is exceeded
4. Forwards the request
5. Parses token usage from the response (including streaming SSE responses)
6. Calculates cost using a live price catalog
7. Increments budget counters
8. Emits OpenTelemetry metrics

Because it works at the `httpx` transport layer, it is compatible with the official OpenAI, Anthropic, and any other SDK that accepts a custom `http_client`.

---

## Installation

```bash
pip install tokenwarden
```

With optional extras:

```bash
pip install tokenwarden[redis]   # Redis-backed distributed budget store
pip install tokenwarden[otel]    # OpenTelemetry signals
pip install tokenwarden[all]     # Everything
```

---

## Quick start

```python
import httpx
from openai import OpenAI
from tokenwarden import TokenWardenTransport
from tokenwarden.budget._store import BudgetConfig, BudgetMode

transport = TokenWardenTransport(
    budgets=[
        BudgetConfig(
            key="global-daily",
            limit=50.00,           # $50/day
            window_seconds=86_400,
            mode=BudgetMode.ENFORCE,
        )
    ]
)

client = OpenAI(http_client=httpx.Client(transport=transport))

response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "Hello!"}],
)
print(response.choices[0].message.content)
```

When the daily budget is exhausted, the next request receives a `429` before it ever leaves your process — no API call is made and no cost is incurred.

---

## Provider examples

### OpenAI

```python
import httpx
from openai import OpenAI
from tokenwarden import TokenWardenTransport

transport = TokenWardenTransport()
client = OpenAI(http_client=httpx.Client(transport=transport))

response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "Explain quantum computing simply."}],
)
print(response.choices[0].message.content)
```

### Anthropic

```python
import anthropic
import httpx
from tokenwarden import TokenWardenTransport

transport = TokenWardenTransport()
client = anthropic.Anthropic(http_client=httpx.Client(transport=transport))

message = client.messages.create(
    model="claude-sonnet-4-6",
    max_tokens=1024,
    messages=[{"role": "user", "content": "Write a haiku about the ocean."}],
)
print(message.content[0].text)
```

### Google Gemini

```python
import os
import httpx
from tokenwarden import TokenWardenTransport

transport = TokenWardenTransport()
api_key = os.environ["GEMINI_API_KEY"]

with httpx.Client(transport=transport) as client:
    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={api_key}"
    response = client.post(url, json={
        "contents": [{"parts": [{"text": "What is the capital of France?"}]}]
    })
    data = response.json()
    print(data["candidates"][0]["content"]["parts"][0]["text"])
```

### xAI (Grok)

```python
import os
import httpx
from openai import OpenAI
from tokenwarden import TokenWardenTransport

transport = TokenWardenTransport()
client = OpenAI(
    api_key=os.environ["XAI_API_KEY"],
    base_url="https://api.x.ai/v1",
    http_client=httpx.Client(transport=transport),
)

response = client.chat.completions.create(
    model="grok-3",
    messages=[{"role": "user", "content": "What is the meaning of life?"}],
)
print(response.choices[0].message.content)
```

### DeepSeek

```python
import os
import httpx
from openai import OpenAI
from tokenwarden import TokenWardenTransport

transport = TokenWardenTransport()
client = OpenAI(
    api_key=os.environ["DEEPSEEK_API_KEY"],
    base_url="https://api.deepseek.com/v1",
    http_client=httpx.Client(transport=transport),
)

response = client.chat.completions.create(
    model="deepseek-chat",
    messages=[{"role": "user", "content": "Summarise the theory of relativity."}],
)
print(response.choices[0].message.content)
```

---

## Attribution

Attribution tells TokenWarden *who* made a call so budget limits can be applied per tenant, per user, or per feature. Set it using context variables — safe across threads and async tasks.

```python
from tokenwarden import set_attribution, clear_attribution

set_attribution(
    tenant_id="acme-corp",
    user_id="user-42",
    feature="report-generation",
)

# LLM calls made here are attributed automatically

clear_attribution()  # optional cleanup
```

### In a FastAPI middleware

```python
from fastapi import Request
from tokenwarden import set_attribution

@app.middleware("http")
async def attach_attribution(request: Request, call_next):
    set_attribution(
        tenant_id=request.headers.get("X-Tenant-Id"),
        user_id=request.headers.get("X-User-Id"),
    )
    return await call_next(request)
```

### In a Django view

```python
from tokenwarden import set_attribution

def my_view(request):
    set_attribution(
        tenant_id=request.user.profile.tenant_id,
        user_id=str(request.user.id),
    )
    # ... call LLM
```

---

## Budget configuration

### Budget modes

| Mode | Behaviour |
|---|---|
| `BudgetMode.ENFORCE` | Returns `429 application/problem+json` before the request is sent |
| `BudgetMode.WARN` | Allows the request but emits an OTel breach counter |

### Budget dimensions

| Dimension | Tracks spend per |
|---|---|
| `"global"` (default) | All requests combined |
| `"tenant"` | `tenant_id` set via `set_attribution()` |
| `"user"` | `user_id` set via `set_attribution()` |
| `"feature"` | `feature` set via `set_attribution()` |

### Examples

```python
from tokenwarden.budget._store import BudgetConfig, BudgetMode

# $50/day global — hard block
BudgetConfig(key="global-daily", limit=50.00, window_seconds=86_400)

# $1,000/month per tenant — hard block
BudgetConfig(key="tenant-monthly", limit=1_000.00, window_seconds=2_592_000,
             mode=BudgetMode.ENFORCE, dimension="tenant")

# $5/day per user — warn only
BudgetConfig(key="user-daily", limit=5.00, window_seconds=86_400,
             mode=BudgetMode.WARN, dimension="user")
```

Multiple budgets applied at once:

```python
from tokenwarden import TokenWardenTransport, set_attribution
from tokenwarden.budget._store import BudgetConfig, BudgetMode
import httpx
from openai import OpenAI

transport = TokenWardenTransport(
    budgets=[
        BudgetConfig(key="tenant-monthly", limit=1_000.00,
                     window_seconds=2_592_000, dimension="tenant"),
        BudgetConfig(key="user-daily",     limit=5.00,
                     window_seconds=86_400,   dimension="user",
                     mode=BudgetMode.WARN),
    ]
)

set_attribution(tenant_id="acme-corp", user_id="user-42")
client = OpenAI(http_client=httpx.Client(transport=transport))
```

---

## Async usage

Use `AsyncTokenWardenTransport` with `httpx.AsyncClient`:

```python
import asyncio
import httpx
from openai import AsyncOpenAI
from tokenwarden import AsyncTokenWardenTransport
from tokenwarden.budget._store import BudgetConfig

transport = AsyncTokenWardenTransport(
    budgets=[BudgetConfig(key="global-daily", limit=50.00, window_seconds=86_400)]
)

async def main():
    client = AsyncOpenAI(http_client=httpx.AsyncClient(transport=transport))
    response = await client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hello!"}],
    )
    print(response.choices[0].message.content)

asyncio.run(main())
```

---

## Redis budget store

For multi-instance deployments where budget state must be shared across processes:

```bash
pip install tokenwarden[redis]
```

```python
from tokenwarden import TokenWardenTransport, RedisBudgetStore
from tokenwarden.budget._store import BudgetConfig
import httpx
from openai import OpenAI

store = RedisBudgetStore("redis://localhost:6379")

transport = TokenWardenTransport(
    store=store,
    budgets=[
        BudgetConfig(key="tenant-monthly", limit=1_000.00,
                     window_seconds=2_592_000, dimension="tenant")
    ]
)

client = OpenAI(http_client=httpx.Client(transport=transport))
```

Uses an atomic Lua script (`INCRBYFLOAT` + `EXPIRE`) — concurrent requests across multiple instances are safe.

---

## Live price catalog

TokenWarden ships with an embedded price catalog. To keep prices current without reinstalling, start a background refresh thread:

```python
from tokenwarden import TokenWardenTransport, CatalogRefreshThread
from tokenwarden.catalog._catalog import PriceCatalog
import httpx
from openai import OpenAI

catalog = PriceCatalog()
refresher = CatalogRefreshThread(catalog, interval_seconds=21_600)  # every 6 hours
refresher.start()  # daemon thread — exits with the process

transport = TokenWardenTransport()
client = OpenAI(http_client=httpx.Client(transport=transport))
```

On any fetch failure the last-known-good catalog is retained.

---

## OpenTelemetry

```bash
pip install tokenwarden[otel]
```

Wire up your exporter as normal — TokenWarden emits automatically:

```python
from opentelemetry import metrics
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import ConsoleMetricExporter, PeriodicExportingMetricReader

metrics.set_meter_provider(MeterProvider(
    metric_readers=[PeriodicExportingMetricReader(ConsoleMetricExporter())]
))

from tokenwarden import TokenWardenTransport
transport = TokenWardenTransport()  # signals flow automatically
```

### Signals emitted

| Signal | Type | Tags |
|---|---|---|
| `tokenwarden.tokens_consumed` | Counter | `gen_ai.system`, `gen_ai.request.model` |
| `tokenwarden.cost_usd` | Counter | `gen_ai.system`, `gen_ai.request.model` |
| `tokenwarden.request_latency_ms` | Histogram | `gen_ai.system`, `gen_ai.request.model` |
| `tokenwarden.budget_breaches` | Counter | `tokenwarden.budget_key`, `tokenwarden.budget_mode` |

Attribution tags (`tokenwarden.tenant_id`, `tokenwarden.user_id`, `tokenwarden.feature`) are added when set via `set_attribution()`.

---

## Handling budget exceeded responses

When an Enforce budget is breached, TokenWarden returns a `429` with a JSON body:

```json
{
  "type": "https://tokenwarden.dev/errors/budget-exceeded",
  "title": "Budget Exceeded",
  "status": 429,
  "detail": "Budget 'tenant-monthly' has been exceeded."
}
```

The OpenAI and Anthropic SDKs raise `RateLimitError` on `429`, so catch it normally:

```python
import openai

try:
    response = client.chat.completions.create(...)
except openai.RateLimitError as e:
    print("Budget exhausted — try again later")
```

---

## Supported providers

| Provider | Host | Notes |
|---|---|---|
| OpenAI | `api.openai.com` | Streaming: final chunk usage |
| Anthropic | `api.anthropic.com` | Streaming: `message_start` + `message_delta` |
| Google Gemini | `generativelanguage.googleapis.com` | `usageMetadata` in response |
| xAI (Grok) | `api.x.ai` | OpenAI-compatible |
| DeepSeek | `api.deepseek.com` | OpenAI-compatible |

---

## License

MIT
