Metadata-Version: 2.4
Name: azureaicommunity-agent-token-guard
Version: 0.1.0
Summary: Token usage tracking and quota enforcement middleware for AI agent pipelines
Author-email: Vinoth Rajendran <r.vinoth@live.com>
License: MIT License
        
        Copyright (c) 2026 Azure AI Community
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/Azure-AI-Community/python-Agent-middleware
Project-URL: Repository, https://github.com/Azure-AI-Community/python-Agent-middleware
Project-URL: Issues, https://github.com/Azure-AI-Community/python-Agent-middleware/issues
Keywords: ai,tokens,quota,middleware,llm,agent,azure,azure-ai,community
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: agent-framework
Dynamic: license-file

<div align="center">

# 🔑 AzureAICommunity - Agent - Token Guard Middleware

Token usage tracking and quota enforcement middleware for AI agent applications built on the **Agent Framework**.

[![PyPI Version](https://img.shields.io/pypi/v/azureaicommunity-agent-token-guard)](https://pypi.org/project/azureaicommunity-agent-token-guard/)
[![Python Versions](https://img.shields.io/pypi/pyversions/azureaicommunity-agent-token-guard)](https://pypi.org/project/azureaicommunity-agent-token-guard/)
![PyPI Downloads](https://static.pepy.tech/badge/azureaicommunity-agent-token-guard)
[![License](https://img.shields.io/pypi/l/azureaicommunity-agent-token-guard)](https://pypi.org/project/azureaicommunity-agent-token-guard/)
[![PyPI Status](https://img.shields.io/pypi/status/azureaicommunity-agent-token-guard)](https://pypi.org/project/azureaicommunity-agent-token-guard/)



**Track every token, enforce every limit — supports all providers, both streaming and non-streaming.**

[Getting Started](#-installation) · [Configuration](#️-configuration) · [Usage](#-usage) · [Contributing](#-contributing)

</div>

---

## Overview

`azureaicommunity-agent-token-guard` is a plug-and-play token tracking and quota enforcement layer for AI agent pipelines built on `agent-framework`. It captures token usage per request, accumulates it against a period quota, and blocks future requests once the limit is hit — with zero changes to your existing agent code.

---

## ✨ Features

| | Feature |
|---|---|
| 📊 | **Track token usage** — captures `input_tokens`, `output_tokens`, `total_tokens`, model, and timestamp per request |
| 🚫 | **Enforce quotas** — blocks requests before they reach the LLM once a period limit is hit |
| 🔔 | **Quota alerts** — fires a callback when the limit is exceeded (log, notify, charge) |
| 🌊 | **Streaming support** — works with both `stream=True` and regular calls |
| 📅 | **Period-flexible** — built-in `month_key`, `week_key`, `day_key` or bring your own |
| 👥 | **Per-user quotas** — pluggable `user_id_getter` for multi-tenant apps |
| 🗄️ | **Pluggable storage** — implement `QuotaStore` protocol to use Redis, Postgres, etc. |
| 🔌 | **Provider-agnostic** — works with any `agent-framework` compatible LLM client |

---

## 📦 Installation

```bash
pip install azureaicommunity-agent-token-guard
```

---

## 🚀 Quick Start

```python
import asyncio, json
from agent_framework import Agent
from agent_framework.ollama import OllamaChatClient
from token_guard_middleware import TokenGuardMiddleware
from token_guard_middleware.token_tracker import InMemoryQuotaStore, QuotaExceededError

def save_usage(record):
    print(json.dumps(record, indent=2))

def quota_alert(payload):
    print("QUOTA EXCEEDED:", json.dumps(payload, indent=2))

quota_store = InMemoryQuotaStore()

middleware = TokenGuardMiddleware(
    on_usage=save_usage,
    on_quota_exceeded=quota_alert,
    quota_store=quota_store,
    quota_tokens=50,          # intentionally low to show quota enforcement
)

async def main():
    client = OllamaChatClient(model="gemma3:4b")
    agent = Agent(client)

    # First call — succeeds and records ~60 tokens (exceeds quota of 50)
    try:
        result = await agent.run("Hello!", middleware=[middleware])
        print(result.text)
    except QuotaExceededError as e:
        print(f"Blocked: {e}")

    # Second call — quota already exceeded, quota_alert fires and call is blocked
    try:
        result = await agent.run("How are you?", middleware=[middleware])
        print(result.text)
    except QuotaExceededError as e:
        print(f"Blocked: {e}")

asyncio.run(main())
```

---

## 🧑‍💻 Usage

### Usage Record

Every call to `on_usage` receives a dict:

```json
{
  "user_id": "anonymous",
  "period_key": "2026-04",
  "model": "gemma3:4b",
  "input_tokens": 11,
  "output_tokens": 52,
  "total_tokens": 63,
  "quota_tokens": 50,
  "used_tokens_after_call": 63,
  "timestamp_utc": "2026-04-14T11:46:09.698893+00:00",
  "streaming": false
}
```

### Quota Alert Payload

When the quota is exceeded `on_quota_exceeded` receives:

```json
{
  "user_id": "anonymous",
  "period_key": "2026-04",
  "used_tokens": 63,
  "quota_tokens": 50,
  "reason": "quota_exceeded_before_call"
}
```

---

## ⚙️ Configuration

### `TokenGuardMiddleware`

| Parameter | Type | Default | Description |
|---|---|---|---|
| `on_usage` | `Callable[[dict], Any]` | **required** | Called after every successful request with the usage record |
| `quota_store` | `QuotaStore` | **required** | Storage backend for accumulated token counts |
| `quota_tokens` | `int` | **required** | Max tokens allowed per period |
| `on_quota_exceeded` | `Callable[[dict], Any]` | `None` | Called when quota is hit (before raising) |
| `user_id_getter` | `Callable[[ChatContext], str]` | `default_user_id_getter` | Extracts user/tenant ID from context |
| `period_key_fn` | `Callable[[], str]` | `month_key` | Returns the current billing period key |

### Period key functions

```python
from token_guard_middleware.token_tracker import month_key, week_key, day_key

middleware = TokenGuardMiddleware(..., period_key_fn=month_key)   # Monthly (default)
middleware = TokenGuardMiddleware(..., period_key_fn=day_key)     # Daily
middleware = TokenGuardMiddleware(..., period_key_fn=week_key)    # Weekly

# Custom — e.g. per-user-per-day
middleware = TokenGuardMiddleware(
    ...,
    period_key_fn=lambda: f"{get_current_user_id()}-{day_key()}",
)
```

### Per-user quotas

```python
def get_user_id(context):
    return context.metadata.get("user_id", "anonymous")

middleware = TokenGuardMiddleware(
    ...,
    user_id_getter=get_user_id,
)
```

### Custom Storage Backend

Implement the `QuotaStore` protocol to persist usage in Redis, Postgres, or any other store:

```python
from token_guard_middleware.token_tracker import QuotaStore

class RedisQuotaStore:
    def get_usage(self, user_id: str, period_key: str) -> int:
        return int(redis.get(f"{user_id}:{period_key}") or 0)

    def add_usage(self, user_id: str, period_key: str, tokens: int) -> None:
        redis.incrby(f"{user_id}:{period_key}", tokens)

middleware = TokenGuardMiddleware(
    ...,
    quota_store=RedisQuotaStore(),
)
```

---

## ⚙️ How It Works

```
1. Intercept  →  middleware captures the outgoing agent request
2. Check      →  quota store is queried for current period usage
3. Block      →  if quota exceeded, raises QuotaExceededError before calling LLM
4. Forward    →  request proceeds to the LLM provider
5. Track      →  response token counts are extracted and written to quota store
6. Notify     →  on_usage callback fires with the full usage record
```

**Provider Compatibility:**

Works with any LLM client that implements the `agent-framework` `ChatClient` interface.

---

## 🤝 Contributing

Contributions are welcome! Please open an issue to discuss what you'd like to change before submitting a pull request.

1. Fork the repository
2. Create a feature branch (`git checkout -b feature/my-feature`)
3. Commit your changes (`git commit -m 'Add my feature'`)
4. Push to the branch (`git push origin feature/my-feature`)
5. Open a Pull Request

---

## 📄 License

MIT — see [LICENSE](LICENSE) for details.
