98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
"""Sum token usage across compressed and uncompressed run directories."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Annotated
|
|
|
|
import typer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class UsageTotals:
|
|
"""Aggregate usage counters for a directory of run records."""
|
|
|
|
files: int = 0
|
|
errors: int = 0
|
|
prompt_tokens: int = 0
|
|
cached_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
reasoning_tokens: int = 0
|
|
total_tokens: int = 0
|
|
per_file: list[tuple[str, int, int, int]] = field(default_factory=list)
|
|
|
|
|
|
def tally_directory(directory: Path) -> UsageTotals:
|
|
"""Return aggregated usage stats for every JSON record in a directory."""
|
|
totals = UsageTotals()
|
|
decoder = json.JSONDecoder()
|
|
for path in sorted(directory.glob("*.json")):
|
|
text = path.read_text().lstrip()
|
|
record, _ = decoder.raw_decode(text)
|
|
totals.files += 1
|
|
usage = record.get("usage")
|
|
if not usage:
|
|
totals.errors += 1
|
|
continue
|
|
prompt_tokens = usage.get("prompt_tokens", 0)
|
|
completion_tokens = usage.get("completion_tokens", 0)
|
|
total_tokens = usage.get("total_tokens", 0)
|
|
cached_tokens = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0)
|
|
reasoning_tokens = (usage.get("completion_tokens_details") or {}).get("reasoning_tokens", 0)
|
|
totals.prompt_tokens += prompt_tokens
|
|
totals.completion_tokens += completion_tokens
|
|
totals.total_tokens += total_tokens
|
|
totals.cached_tokens += cached_tokens
|
|
totals.reasoning_tokens += reasoning_tokens
|
|
totals.per_file.append((path.name, prompt_tokens, completion_tokens, total_tokens))
|
|
return totals
|
|
|
|
|
|
def log_totals(label: str, totals: UsageTotals) -> None:
|
|
"""Log a one-block summary for a directory."""
|
|
counted = totals.files - totals.errors
|
|
average_total = totals.total_tokens / counted if counted else 0
|
|
logger.info("[%s]", label)
|
|
logger.info(" files : %d (with usage: %d, errors: %d)", totals.files, counted, totals.errors)
|
|
logger.info(" prompt tokens : %d", totals.prompt_tokens)
|
|
logger.info(" cached tokens : %d", totals.cached_tokens)
|
|
logger.info(" completion tok : %d", totals.completion_tokens)
|
|
logger.info(" reasoning tok : %d", totals.reasoning_tokens)
|
|
logger.info(" total tokens : %d", totals.total_tokens)
|
|
logger.info(" avg total/file : %.1f", average_total)
|
|
|
|
|
|
def main(
|
|
runs_dir: Annotated[Path, typer.Option("--runs-dir")] = Path("output/openai_runs_temp_1"),
|
|
log_level: Annotated[str, typer.Option("--log-level")] = "INFO",
|
|
) -> None:
|
|
"""Print token usage totals for the compressed and uncompressed run directories."""
|
|
logging.basicConfig(level=log_level, format="%(message)s")
|
|
|
|
grand = UsageTotals()
|
|
for label in ("compressed", "uncompressed"):
|
|
directory = runs_dir / label
|
|
if not directory.is_dir():
|
|
logger.warning("%s: directory not found at %s", label, directory)
|
|
continue
|
|
totals = tally_directory(directory)
|
|
log_totals(label, totals)
|
|
grand.files += totals.files
|
|
grand.errors += totals.errors
|
|
grand.prompt_tokens += totals.prompt_tokens
|
|
grand.cached_tokens += totals.cached_tokens
|
|
grand.completion_tokens += totals.completion_tokens
|
|
grand.reasoning_tokens += totals.reasoning_tokens
|
|
grand.total_tokens += totals.total_tokens
|
|
|
|
log_totals("grand total", grand)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|