https://github.com/cwest/ai-tokentrace
ai-tokentrace is a Python library for GenAI cost observability. It helps developers track token consumption in Google Generative AI applications to manage costs and optimize performance.
https://github.com/cwest/ai-tokentrace
adk-python ai cost-management firestore gemini genai google google-genai observability pubsub python telemetry token-tracing
Last synced: 6 months ago
JSON representation
ai-tokentrace is a Python library for GenAI cost observability. It helps developers track token consumption in Google Generative AI applications to manage costs and optimize performance.
- Host: GitHub
- URL: https://github.com/cwest/ai-tokentrace
- Owner: cwest
- License: apache-2.0
- Created: 2025-10-22T16:59:11.000Z (9 months ago)
- Default Branch: main
- Last Pushed: 2025-10-27T12:24:41.000Z (9 months ago)
- Last Synced: 2025-12-22T02:04:03.939Z (7 months ago)
- Topics: adk-python, ai, cost-management, firestore, gemini, genai, google, google-genai, observability, pubsub, python, telemetry, token-tracing
- Language: Python
- Homepage: https://pypi.org/project/ai-tokentrace/
- Size: 332 KB
- Stars: 1
- Watchers: 0
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- Contributing: docs/contributing.md
- License: LICENSE
- Code of conduct: docs/code-of-conduct.md
Awesome Lists containing this project
README
# ai-tokentrace
[](https://badge.fury.io/py/ai-tokentrace)
[](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.