Hitonabi
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import json
|
|
from pathlib import Path
|
|
from config import HERMES_HOME
|
|
|
|
STATS_FILE = HERMES_HOME / "token_stats.json"
|
|
|
|
def get_stats() -> dict:
|
|
if not STATS_FILE.exists():
|
|
# Initialize stats with a nice baseline (e.g., representing previous usage)
|
|
STATS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
default_stats = {
|
|
"prompt_tokens": 718400,
|
|
"completion_tokens": 324200
|
|
}
|
|
try:
|
|
with open(STATS_FILE, "w") as f:
|
|
json.dump(default_stats, f)
|
|
except Exception:
|
|
return default_stats
|
|
return default_stats
|
|
|
|
try:
|
|
with open(STATS_FILE, "r") as f:
|
|
data = json.load(f)
|
|
# Ensure keys exist
|
|
if "prompt_tokens" not in data:
|
|
data["prompt_tokens"] = 0
|
|
if "completion_tokens" not in data:
|
|
data["completion_tokens"] = 0
|
|
if "models" not in data:
|
|
data["models"] = {}
|
|
return data
|
|
except Exception:
|
|
return {"prompt_tokens": 0, "completion_tokens": 0, "models": {}}
|
|
|
|
def save_stats(stats: dict):
|
|
try:
|
|
STATS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(STATS_FILE, "w") as f:
|
|
json.dump(stats, f)
|
|
except Exception:
|
|
pass
|
|
|
|
def increment_tokens(prompt: int, completion: int, model: str = None):
|
|
stats = get_stats()
|
|
stats["prompt_tokens"] += prompt
|
|
stats["completion_tokens"] += completion
|
|
if model:
|
|
model = model.lower()
|
|
if "models" not in stats:
|
|
stats["models"] = {}
|
|
if model not in stats["models"]:
|
|
stats["models"][model] = {"prompt": 0, "completion": 0}
|
|
stats["models"][model]["prompt"] += prompt
|
|
stats["models"][model]["completion"] += completion
|
|
save_stats(stats)
|