File size: 2,145 Bytes
317eaf3 ce5b66d 317eaf3 ce5b66d 317eaf3 ce5b66d 317eaf3 | 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 | """
Token usage tracker – stores cumulative token counts on disk.
All values are in raw tokens (no monetary conversion).
"""
import json
import threading
from pathlib import Path
STATS_FILE = Path(__file__).parent.parent / "db" / "token_stats.json"
_lock = threading.Lock()
def _default_stats():
return {
"embedding_tokens": 0,
"prompt_tokens": 0,
"completion_tokens": 0,
"toc_analysis_tokens": 0,
}
def load_stats():
if STATS_FILE.exists():
try:
with open(STATS_FILE, "r") as f:
return json.load(f)
except Exception:
pass
return _default_stats()
def save_stats(stats):
STATS_FILE.parent.mkdir(exist_ok=True, parents=True)
with open(STATS_FILE, "w") as f:
json.dump(stats, f)
def add_embedding_tokens(count: int):
if count <= 0:
return
with _lock:
stats = load_stats()
stats["embedding_tokens"] = stats.get("embedding_tokens", 0) + count
save_stats(stats)
def add_chat_tokens(prompt_count: int, completion_count: int):
if prompt_count <= 0 and completion_count <= 0:
return
with _lock:
stats = load_stats()
stats["prompt_tokens"] = stats.get("prompt_tokens", 0) + prompt_count
stats["completion_tokens"] = stats.get("completion_tokens", 0) + completion_count
save_stats(stats)
def add_toc_analysis_tokens(count: int):
if count <= 0:
return
with _lock:
stats = load_stats()
stats["toc_analysis_tokens"] = stats.get("toc_analysis_tokens", 0) + count
save_stats(stats)
def get_token_summary():
"""Return raw token counts only – no monetary costs."""
with _lock:
stats = load_stats()
total = (
stats.get("embedding_tokens", 0)
+ stats.get("prompt_tokens", 0)
+ stats.get("completion_tokens", 0)
+ stats.get("toc_analysis_tokens", 0)
)
return {
"tokens": stats,
"total_tokens": total,
}
def clear_stats():
with _lock:
save_stats(_default_stats())
|