src/token_meter.py

from __future__ import annotations

import hashlib
import json
import math
import os
import time
import uuid
from pathlib import Path
from typing import Any


APP_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_AUDIT_DIR = APP_ROOT / "data" / "token_audit"


def new_transaction_id(prefix: str = "totem") -> str:
    stamp = time.strftime("%Y%m%d-%H%M%S", time.gmtime())
    return f"{prefix}_{stamp}_{uuid.uuid4().hex[:10]}"


def stable_hash(value: Any) -> str:
    if isinstance(value, str):
        payload = value
    else:
        payload = json.dumps(value, ensure_ascii=False, sort_keys=True, default=str)
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def compact_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)


def estimate_text_tokens(text: str, model: str | None = None) -> dict[str, Any]:
    """
    Estimate tokens before an API call. Exact transaction totals should still come
    from the provider response usage object after the call completes.
    """
    text = text or ""
    try:
        import tiktoken  # type: ignore

        try:
            enc = tiktoken.encoding_for_model(model or "")
            method = f"tiktoken:{enc.name}:model"
        except Exception:
            enc = tiktoken.get_encoding("cl100k_base")
            method = "tiktoken:cl100k_base:fallback"
        return {
            "tokens": len(enc.encode(text)),
            "chars": len(text),
            "method": method,
        }
    except Exception:
        # Conservative no-dependency approximation. Most English prompt material
        # lands around 3.5-4 chars/token, but JSON punctuation pushes lower.
        return {
            "tokens": int(math.ceil(len(text) / 4)) if text else 0,
            "chars": len(text),
            "method": "heuristic:ceil(chars/4)",
        }


def estimate_components(components: dict[str, str], model: str | None = None) -> dict[str, Any]:
    by_component = {
        name: estimate_text_tokens(text, model=model)
        for name, text in components.items()
    }
    return {
        "components": by_component,
        "estimated_input_total": sum(item["tokens"] for item in by_component.values()),
        "methods": sorted({item["method"] for item in by_component.values()}),
    }


def extract_response_usage(response: Any) -> dict[str, int | None]:
    usage = getattr(response, "usage", None)
    if usage is None and isinstance(response, dict):
        usage = response.get("usage")
    if usage is None:
        return {
            "prompt_tokens": None,
            "completion_tokens": None,
            "total_tokens": None,
        }

    def get(name: str) -> int | None:
        value = getattr(usage, name, None)
        if value is None and isinstance(usage, dict):
            value = usage.get(name)
        try:
            return int(value) if value is not None else None
        except Exception:
            return None

    return {
        "prompt_tokens": get("prompt_tokens"),
        "completion_tokens": get("completion_tokens"),
        "total_tokens": get("total_tokens"),
    }


def audit_dir() -> Path:
    return Path(os.getenv("TOTEM_TOKEN_AUDIT_DIR") or DEFAULT_AUDIT_DIR)


def append_token_audit(record: dict[str, Any]) -> str:
    out_dir = audit_dir()
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / "totem_token_usage.jsonl"
    with path.open("a", encoding="utf-8") as fh:
        fh.write(compact_json(record) + "\n")
    return str(path)