rex/llm_usage.py
Rama master · réplica de lectura del repositorio original.
"""LLM usage tracking — append-only JSON log with 10 MB rotation."""
from __future__ import annotations
import json
import os
from collections import defaultdict
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
from rex.logging_utils import get_logger
logger = get_logger(__name__)
_DEFAULT_PATH = Path("data/llm_usage.json")
_MAX_SIZE_BYTES = 10 * 1024 * 1024 # 10 MB
def _usage_path() -> Path:
return Path(os.environ.get("REX_LLM_USAGE_PATH", str(_DEFAULT_PATH)))
def record_usage(
*,
model: str,
prompt_tokens: int,
completion_tokens: int,
timestamp: str | None = None,
) -> None:
"""Append one usage record to the usage log file.
Args:
model: Model name (e.g. ``"llama3"``).
prompt_tokens: Number of prompt/input tokens.
completion_tokens: Number of completion/output tokens.
timestamp: ISO-8601 timestamp; defaults to current UTC time.
"""
path = _usage_path()
path.parent.mkdir(parents=True, exist_ok=True)
entry: dict[str, Any] = {
"model": model,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"timestamp": timestamp or datetime.now(UTC).isoformat(),
}
# Rotate if file exceeds 10 MB
if path.exists() and path.stat().st_size >= _MAX_SIZE_BYTES:
rotated = path.with_suffix(".json.1")
path.replace(rotated)
logger.info("LLM usage log rotated to %s", rotated)
try:
with path.open("a", encoding="utf-8") as fh:
fh.write(json.dumps(entry) + "\n")
except OSError as exc:
logger.warning("Failed to write LLM usage record: %s", exc)
def load_records(path: Path | None = None) -> list[dict[str, Any]]:
"""Return all usage records from the log file (empty list if missing)."""
target = path or _usage_path()
if not target.exists():
return []
records: list[dict[str, Any]] = []
try:
with target.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if line:
try:
records.append(json.loads(line))
except json.JSONDecodeError:
pass
except OSError as exc:
logger.warning("Failed to read LLM usage log: %s", exc)
return records
def summarise(path: Path | None = None) -> dict[str, Any]:
"""Return a usage summary dict.
Returns::
{
"total_requests": int,
"total_tokens": int,
"by_model": {
"<model>": {"requests": int, "prompt_tokens": int,
"completion_tokens": int, "total_tokens": int},
...
},
}
"""
records = load_records(path)
total_requests = len(records)
total_prompt = 0
total_completion = 0
by_model: dict[str, dict[str, int]] = defaultdict(
lambda: {"requests": 0, "prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
)
for rec in records:
prompt = int(rec.get("prompt_tokens", 0))
completion = int(rec.get("completion_tokens", 0))
model = str(rec.get("model", "unknown"))
total_prompt += prompt
total_completion += completion
by_model[model]["requests"] += 1
by_model[model]["prompt_tokens"] += prompt
by_model[model]["completion_tokens"] += completion
by_model[model]["total_tokens"] += prompt + completion
return {
"total_requests": total_requests,
"total_tokens": total_prompt + total_completion,
"by_model": dict(by_model),
}
# ---------------------------------------------------------------------------
# Cloud vs local classification
# ---------------------------------------------------------------------------
_CLOUD_PREFIXES = (
"gpt-",
"text-",
"o1",
"o3",
"babbage",
"davinci",
"curie",
"ada",
"claude-",
"gemini-",
"mistral-",
)
def _is_cloud_model(model: str) -> bool:
"""Return True if *model* is a cloud-hosted model (OpenAI, Anthropic, Google, etc.)."""
if not model:
return False
lower = model.lower()
for prefix in _CLOUD_PREFIXES:
if lower.startswith(prefix):
return True
return False
def _empty_bucket() -> dict[str, int]:
return {"requests": 0, "tokens": 0}
def _period_start(now: datetime, period: str) -> datetime:
if period == "today":
return now.replace(hour=0, minute=0, second=0, microsecond=0)
if period == "week":
return (now - timedelta(days=now.weekday())).replace(
hour=0, minute=0, second=0, microsecond=0
)
# month
return now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
def usage_api_summary(path: Path | None = None) -> dict[str, Any]:
"""Return usage split by local/cloud for today, this week, and this month.
Returns::
{
"local": {"requests": int, "tokens": int},
"cloud": {"requests": int, "tokens": int},
"by_period": {
"today": {
"local": {"requests": int, "tokens": int},
"cloud": {"requests": int, "tokens": int},
},
"week": {...},
"month": {...},
},
}
"""
records = load_records(path)
now = datetime.now(UTC)
totals: dict[str, dict[str, int]] = {
"local": _empty_bucket(),
"cloud": _empty_bucket(),
}
by_period: dict[str, dict[str, dict[str, int]]] = {
"today": {"local": _empty_bucket(), "cloud": _empty_bucket()},
"week": {"local": _empty_bucket(), "cloud": _empty_bucket()},
"month": {"local": _empty_bucket(), "cloud": _empty_bucket()},
}
period_starts = {p: _period_start(now, p) for p in ("today", "week", "month")}
for rec in records:
model = str(rec.get("model", "unknown"))
bucket = "cloud" if _is_cloud_model(model) else "local"
tokens = int(rec.get("prompt_tokens", 0)) + int(rec.get("completion_tokens", 0))
totals[bucket]["requests"] += 1
totals[bucket]["tokens"] += tokens
ts_raw = rec.get("timestamp", "")
try:
ts = datetime.fromisoformat(str(ts_raw))
if ts.tzinfo is None:
ts = ts.replace(tzinfo=UTC)
except (ValueError, TypeError):
ts = None
if ts is not None:
for period, start in period_starts.items():
if ts >= start:
by_period[period][bucket]["requests"] += 1
by_period[period][bucket]["tokens"] += tokens
return {
"local": totals["local"],
"cloud": totals["cloud"],
"by_period": by_period,
}