#!/usr/bin/env python3 """ CODITECT Token Usage Tracking System (ADR-075)
Provides unified token tracking across all LLM providers with:
- Per-model breakdowns
- Thread-safe aggregation
- Session export functionality
Based on RLM's UsageSummary pattern (submodules/rlm). """
from dataclasses import dataclass, field from typing import Dict, Optional from datetime import datetime from threading import Lock import json
@dataclass class ModelUsageSummary: """Per-model token usage tracking.
Tracks input, output, and cached tokens for a specific model/provider combination.
Attributes:
model: Model identifier (e.g., "claude-sonnet-4-5")
provider: Provider name (e.g., "anthropic", "openai")
num_calls: Number of API calls made
input_tokens: Total input tokens consumed
output_tokens: Total output tokens generated
cached_tokens: Tokens served from cache (Anthropic prompt caching)
"""
model: str
provider: str
num_calls: int = 0
input_tokens: int = 0
output_tokens: int = 0
cached_tokens: int = 0
@property
def total_tokens(self) -> int:
"""Total tokens (input + output)."""
return self.input_tokens + self.output_tokens
def to_dict(self) -> dict:
"""Serialize to dictionary for JSON export."""
return {
"model": self.model,
"provider": self.provider,
"num_calls": self.num_calls,
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"cached_tokens": self.cached_tokens,
"total_tokens": self.total_tokens
}
@dataclass class UsageSummary: """Aggregated token usage across all models.
Provides a session-level view of token consumption with per-model breakdowns.
Attributes:
session_id: Unique session identifier
started_at: Session start timestamp
model_usage: Dictionary of model usage keyed by "provider:model"
"""
session_id: str
started_at: datetime = field(default_factory=datetime.utcnow)
model_usage: Dict[str, ModelUsageSummary] = field(default_factory=dict)
@property
def total_calls(self) -> int:
"""Total API calls across all models."""
return sum(m.num_calls for m in self.model_usage.values())
@property
def total_input_tokens(self) -> int:
"""Total input tokens across all models."""
return sum(m.input_tokens for m in self.model_usage.values())
@property
def total_output_tokens(self) -> int:
"""Total output tokens across all models."""
return sum(m.output_tokens for m in self.model_usage.values())
@property
def total_tokens(self) -> int:
"""Total tokens (input + output)."""
return self.total_input_tokens + self.total_output_tokens
def add_usage(
self,
model: str,
provider: str,
input_tokens: int,
output_tokens: int,
cached_tokens: int = 0
) -> None:
"""Add usage for a model.
Creates new entry if model not seen, otherwise accumulates.
Args:
model: Model identifier
provider: Provider name
input_tokens: Input tokens for this call
output_tokens: Output tokens for this call
cached_tokens: Cached tokens (optional)
"""
key = f"{provider}:{model}"
if key not in self.model_usage:
self.model_usage[key] = ModelUsageSummary(model=model, provider=provider)
usage = self.model_usage[key]
usage.num_calls += 1
usage.input_tokens += input_tokens
usage.output_tokens += output_tokens
usage.cached_tokens += cached_tokens
def merge(self, other: "UsageSummary") -> "UsageSummary":
"""Merge with another summary (immutable, returns new).
Combines usage from two summaries without modifying either original.
Args:
other: UsageSummary to merge with
Returns:
New UsageSummary with combined usage
"""
merged = UsageSummary(session_id=self.session_id, started_at=self.started_at)
# Copy self
for key, usage in self.model_usage.items():
merged.model_usage[key] = ModelUsageSummary(
model=usage.model,
provider=usage.provider,
num_calls=usage.num_calls,
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
cached_tokens=usage.cached_tokens
)
# Merge other
for key, usage in other.model_usage.items():
if key in merged.model_usage:
m = merged.model_usage[key]
m.num_calls += usage.num_calls
m.input_tokens += usage.input_tokens
m.output_tokens += usage.output_tokens
m.cached_tokens += usage.cached_tokens
else:
merged.model_usage[key] = ModelUsageSummary(
model=usage.model,
provider=usage.provider,
num_calls=usage.num_calls,
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
cached_tokens=usage.cached_tokens
)
return merged
def to_dict(self) -> dict:
"""Serialize to dictionary for JSON export."""
return {
"session_id": self.session_id,
"started_at": self.started_at.isoformat(),
"total_calls": self.total_calls,
"total_input_tokens": self.total_input_tokens,
"total_output_tokens": self.total_output_tokens,
"total_tokens": self.total_tokens,
"model_usage": {k: v.to_dict() for k, v in self.model_usage.items()}
}
class UsageTracker: """Thread-safe usage tracking singleton.
Provides global access to token tracking with thread-safe recording.
Uses singleton pattern for process-wide tracking.
Example:
tracker = UsageTracker()
tracker.start_session("session-123")
tracker.record_usage("claude-sonnet-4-5", "anthropic", 100, 50)
summary = tracker.get_summary()
"""
_instance: Optional["UsageTracker"] = None
_lock: Lock = Lock()
def __new__(cls):
if cls._instance is None:
with cls._lock:
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._summary = None
cls._instance._usage_lock = Lock()
return cls._instance
def start_session(self, session_id: str) -> None:
"""Initialize tracking for a new session.
Args:
session_id: Unique identifier for the session
"""
with self._usage_lock:
self._summary = UsageSummary(session_id=session_id)
def record_usage(
self,
model: str,
provider: str,
input_tokens: int,
output_tokens: int,
cached_tokens: int = 0
) -> None:
"""Record usage from an LLM call (thread-safe).
Safe to call from multiple threads. No-op if session not started.
Args:
model: Model identifier
provider: Provider name
input_tokens: Input tokens consumed
output_tokens: Output tokens generated
cached_tokens: Cached tokens (optional)
"""
with self._usage_lock:
if self._summary:
self._summary.add_usage(
model=model,
provider=provider,
input_tokens=input_tokens,
output_tokens=output_tokens,
cached_tokens=cached_tokens
)
def get_summary(self) -> Optional[UsageSummary]:
"""Get current usage summary (snapshot).
Returns:
Current UsageSummary or None if no session started
"""
with self._usage_lock:
return self._summary
def reset(self) -> None:
"""Reset the tracker, clearing the current session."""
with self._usage_lock:
self._summary = None
def export_jsonl(self, path: str) -> None:
"""Export usage to JSONL file.
Appends current summary as a single JSON line.
Args:
path: Path to JSONL file
"""
with self._usage_lock:
if self._summary:
with open(path, 'a') as f:
f.write(json.dumps(self._summary.to_dict()) + '\n')
Convenience function for quick access
def get_tracker() -> UsageTracker: """Get the global UsageTracker instance.""" return UsageTracker()