|
import json
|
|
import os
|
|
from datetime import datetime
|
|
from typing import Dict, Optional
|
|
from rp.core.logging import get_logger
|
|
|
|
logger = get_logger("usage")
|
|
USAGE_DB_FILE = os.path.expanduser("~/.assistant_usage.json")
|
|
EXCHANGE_RATE = 1.0
|
|
MODEL_COSTS = {
|
|
"x-ai/grok-code-fast-1": {"input": 0.0002, "output": 0.0015},
|
|
"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, {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}
|