Spaces:
Sleeping
Sleeping
File size: 3,556 Bytes
bef6364 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | 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)
|