ai-tokentrace 0.2.1
Add your description here
Requires Python: >=3.9

# ai-tokentrace

[![PyPI version](https://badge.fury.io/py/ai-tokentrace.svg)](https://badge.fury.io/py/ai-tokentrace)
[![CI](https://github.com/cwest/ai-tokentrace/actions/workflows/ci.yml/badge.svg)](https://github.com/cwest/ai-tokentrace/actions/workflows/ci.yml)

**GenAI Cost Observability for Google's Generative AI.**

`ai-tokentrace` provides a transparent and easy way to track token consumption in your GenAI applications. Whether you're using the standard `google-genai` SDK or building complex agents with the Google Agent Development Kit (ADK), this library helps you manage costs, optimize performance, and gain deep insights into your model usage.

## Features

*   **🔍 Automatic Tracking:** Seamlessly integrates with `google-genai` to capture token usage from every API call.
*   **🤖 ADK Support:** Includes a plugin for the Google Agent Development Kit for effortless agent monitoring.
*   **🔌 Multiple Backends:** Export data to where you need it:
    *   **Logging:** Simple standard output for development.
    *   **JSONL:** Structured local files for easy analysis.
    *   **Google Cloud Firestore:** Scalable, queryable cloud storage.
    *   **Google Cloud Pub/Sub:** Event-driven pipelines for real-time analytics.
*   **⚡ Async Native:** Fully non-blocking to keep your applications fast.
*   **📊 Rich Metrics:** Tracks input/output tokens, thinking tokens, cached content, tool usage, and more.

## Installation

Install using `pip` or `uv` (recommended).

### Basic Installation

For standard logging or JSONL export:

```bash
pip install ai-tokentrace
# or
uv pip install ai-tokentrace
```

### With Extra Backends

Install with specific extras for Cloud integrations or ADK support:

```bash
# For Google Cloud Firestore
uv pip install "ai-tokentrace[firestore]"

# For Google Cloud Pub/Sub
uv pip install "ai-tokentrace[pubsub]"

# For Google ADK support
uv pip install "ai-tokentrace[adk]"

# Install everything
uv pip install "ai-tokentrace[firestore,pubsub,adk]"
```

## Quick Start

### 1. Using with `google-genai` SDK

Simply wrap your client with `TrackedGenaiClient`. It works exactly like the standard client but logs all token usage.

```python
import os
from google import genai
from ai_tokentrace import TrackedGenaiClient

# 1. Initialize standard client
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

# 2. Wrap with tracking (uses logging by default)
tracked_client = TrackedGenaiClient(client=client)

# 3. Use as normal!
response = tracked_client.models.generate_content(
    model="gemini-2.5-flash",
    contents="Explain quantum computing in 5 words."
)
print(response.text)
# Output: "Complex superposition processes information fast."
# Log: {"timestamp": "...", "model_name": "gemini-2.5-flash", "total_tokens": 15, ...}
```

### 2. Using with Google ADK

Add the `TokenTrackingPlugin` to your ADK app.

```python
from google.adk.agents import LlmAgent
from google.adk.apps.app import App
from ai_tokentrace.adk import TokenTrackingPlugin

agent = LlmAgent(model="gemini-2.5-flash", ...)

app = App(
    name="my_app",
    root_agent=agent,
    plugins=[TokenTrackingPlugin()]  # Tracks all agent interactions
)
```

## Advanced Usage

### Configuring Backends

You can configure different backends for storing your token usage data.

**Firestore Example:**

```python
from ai_tokentrace import TrackedGenaiClient
from ai_tokentrace.services import FirestoreTokenUsageService

service = FirestoreTokenUsageService(collection_name="genai_usage_logs")
tracked_client = TrackedGenaiClient(client=client, service=service)
```

**Pub/Sub Example:**

```python
from ai_tokentrace import TrackedGenaiClient
from ai_tokentrace.services import PubSubTokenUsageService

service = PubSubTokenUsageService(topic_id="my-usage-topic", project_id="my-project")
tracked_client = TrackedGenaiClient(client=client, service=service)
```

### Self-Inspection for Agents

Give your agents the ability to see their own token usage!

```python
from ai_tokentrace.services import FirestoreTokenUsageService

service = FirestoreTokenUsageService(...)

# Add the inspection tool to your agent
agent = LlmAgent(
    ...,
    tools=[service.get_inspection_tool()]
)
```

## Examples

Check out the `examples/` directory for complete, runnable projects:

*   **[google-genai/](examples/google-genai/)**: Scripts demonstrating sync/async usage, streaming, and different backends.
*   **[adk/](examples/adk/)**: Full ADK applications showing multi-agent tracking, multimodal capabilities, and self-inspection.

## Contributing

Contributions are welcome! Please see [CONTRIBUTING.md](docs/contributing.md) for guidelines.

## License

Apache 2.0 - See [LICENSE](LICENSE) for more details.