src/cost_tracker.py

"""
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())