rp/pr/core/usage_tracker.py

import json
import os
from datetime import datetime
from typing import Dict, Optional

from pr.core.logging import get_logger

logger = get_logger("usage")

USAGE_DB_FILE = os.path.expanduser("~/.assistant_usage.json")

EXCHANGE_RATE = 1.0  # Keep in USD

MODEL_COSTS = {
    "x-ai/grok-code-fast-1": {"input": 0.0002, "output": 0.0015},  # per 1000 tokens in USD
    "gpt-4": {"input": 0.03, "output": 0.06},
    "gpt-4-turbo": {"input": 0.01, "output": 0.03},
    "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
    "claude-3-opus": {"input": 0.015, "output": 0.075},
    "claude-3-sonnet": {"input": 0.003, "output": 0.015},
    "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
}


class UsageTracker:

    def __init__(self):
        self.session_usage = {
            "requests": 0,
            "total_tokens": 0,
            "input_tokens": 0,
            "output_tokens": 0,
            "estimated_cost": 0.0,
            "models_used": {},
        }

    def track_request(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        total_tokens: Optional[int] = None,
    ):
        if total_tokens is None:
            total_tokens = input_tokens + output_tokens

        self.session_usage["requests"] += 1
        self.session_usage["total_tokens"] += total_tokens
        self.session_usage["input_tokens"] += input_tokens
        self.session_usage["output_tokens"] += output_tokens

        if model not in self.session_usage["models_used"]:
            self.session_usage["models_used"][model] = {
                "requests": 0,
                "tokens": 0,
                "cost": 0.0,
            }

        model_usage = self.session_usage["models_used"][model]
        model_usage["requests"] += 1
        model_usage["tokens"] += total_tokens

        cost = self._calculate_cost(model, input_tokens, output_tokens)
        model_usage["cost"] += cost
        self.session_usage["estimated_cost"] += cost

        self._save_to_history(model, input_tokens, output_tokens, cost)

        logger.debug(f"Tracked request: {model}, tokens: {total_tokens}, cost: €{cost:.4f}")

    @staticmethod
    def _calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
        if model not in MODEL_COSTS:
            base_model = model.split("/")[0] if "/" in model else model
            if base_model not in MODEL_COSTS:
                logger.warning(f"Unknown model for cost calculation: {model}")
                return 0.0
            costs = MODEL_COSTS[base_model]
        else:
            costs = MODEL_COSTS[model]

        input_cost = (input_tokens / 1000) * costs["input"] * EXCHANGE_RATE
        output_cost = (output_tokens / 1000) * costs["output"] * EXCHANGE_RATE

        return input_cost + output_cost

    def _save_to_history(self, model: str, input_tokens: int, output_tokens: int, cost: float):
        try:
            history = []
            if os.path.exists(USAGE_DB_FILE):
                with open(USAGE_DB_FILE) as f:
                    history = json.load(f)

            history.append(
                {
                    "timestamp": datetime.now().isoformat(),
                    "model": model,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "total_tokens": input_tokens + output_tokens,
                    "cost": cost,
                }
            )

            if len(history) > 10000:
                history = history[-10000:]

            with open(USAGE_DB_FILE, "w") as f:
                json.dump(history, f, indent=2)

        except Exception as e:
            logger.error(f"Error saving usage history: {e}")

    def get_session_summary(self) -> Dict:
        return self.session_usage.copy()

    def get_formatted_summary(self) -> str:
        usage = self.session_usage
        lines = [
            "\n=== Session Usage Summary ===",
            f"Total Requests: {usage['requests']}",
            f"Total Tokens: {usage['total_tokens']:,}",
            f"  Input: {usage['input_tokens']:,}",
            f"  Output: {usage['output_tokens']:,}",
            f"Estimated Cost: ${usage['estimated_cost']:.4f}",
        ]

        if usage["models_used"]:
            lines.append("\nModels Used:")
            for model, stats in usage["models_used"].items():
                lines.append(
                    f"  {model}: {stats['requests']} requests, "
                    f"{stats['tokens']:,} tokens, ${stats['cost']:.4f}"
                )

        return "\n".join(lines)

    @staticmethod
    def get_total_usage() -> Dict:
        if not os.path.exists(USAGE_DB_FILE):
            return {"total_requests": 0, "total_tokens": 0, "total_cost": 0.0}

        try:
            with open(USAGE_DB_FILE) as f:
                history = json.load(f)

            total_tokens = sum(entry["total_tokens"] for entry in history)
            total_cost = sum(entry["cost"] for entry in history)

            return {
                "total_requests": len(history),
                "total_tokens": total_tokens,
                "total_cost": total_cost,
            }

        except Exception as e:
            logger.error(f"Error loading usage history: {e}")
            return {"total_requests": 0, "total_tokens": 0, "total_cost": 0.0}