219 lines
7.1 KiB
Python
219 lines
7.1 KiB
Python
"""
|
|
Centralized LLM cost tracking utility.
|
|
|
|
Responsibilities:
|
|
1. Calculate token cost per model.
|
|
2. Persist cost logs to logs/{project_id}/cost_{model_name}.log.
|
|
3. Persist structured JSONL records to logs/{project_id}/cost_{model_name}.jsonl.
|
|
4. Provide a single wrapped OpenAI chat-completion call so all components can reuse
|
|
one cost-accounting flow.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
from datetime import datetime
|
|
from decimal import Decimal, ROUND_HALF_UP
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
MODEL_COSTS_PER_1M_TOKENS: Dict[str, Dict[str, float]] = {
|
|
"Qwen/Qwen3.5-27B": {"input": 0.5, "output": 3.0},
|
|
"gemini-3-flash": {"input": 0.5, "output": 3.0},
|
|
"gemini-3.1-flash-lite": {"input": 0.25, "output": 1.5},
|
|
}
|
|
|
|
|
|
_COUNTER_LOCK = threading.Lock()
|
|
_COUNTERS: Dict[str, Dict[str, Decimal]] = {}
|
|
|
|
|
|
def _decimal(value: Any) -> Decimal:
|
|
return Decimal(str(value))
|
|
|
|
|
|
def _quantize_8(value: Decimal) -> Decimal:
|
|
return value.quantize(Decimal("0.00000001"), rounding=ROUND_HALF_UP)
|
|
|
|
|
|
def _safe_model_name(model_name: str) -> str:
|
|
if not model_name:
|
|
return "unknown_model"
|
|
safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", model_name)
|
|
return safe or "unknown_model"
|
|
|
|
|
|
def _resolve_project_id(metadata: Optional[Dict[str, Any]]) -> str:
|
|
metadata = metadata or {}
|
|
value = metadata.get("project_id")
|
|
if value is not None and str(value).strip():
|
|
return str(value).strip()
|
|
return "global"
|
|
|
|
|
|
def _resolve_component(metadata: Optional[Dict[str, Any]]) -> str:
|
|
metadata = metadata or {}
|
|
component = metadata.get("component")
|
|
if component is None:
|
|
return "unknown_component"
|
|
component = str(component).strip()
|
|
return component if component else "unknown_component"
|
|
|
|
|
|
def _normalize_usage(usage: Any) -> Dict[str, int]:
|
|
if usage is None:
|
|
return {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
|
|
|
|
prompt_tokens = int(getattr(usage, "prompt_tokens", 0) or 0)
|
|
completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0)
|
|
total_tokens = int(getattr(usage, "total_tokens", prompt_tokens + completion_tokens) or 0)
|
|
|
|
return {
|
|
"input_tokens": prompt_tokens,
|
|
"output_tokens": completion_tokens,
|
|
"total_tokens": total_tokens,
|
|
}
|
|
|
|
|
|
def _resolve_model_rates(model_name: str) -> Dict[str, Decimal]:
|
|
env_in = os.environ.get("LLM_COST_INPUT_PER_1M")
|
|
env_out = os.environ.get("LLM_COST_OUTPUT_PER_1M")
|
|
if env_in is not None and env_out is not None:
|
|
return {"input": _decimal(env_in), "output": _decimal(env_out)}
|
|
|
|
rates = MODEL_COSTS_PER_1M_TOKENS.get(model_name, {"input": 0.0, "output": 0.0})
|
|
return {"input": _decimal(rates.get("input", 0.0)), "output": _decimal(rates.get("output", 0.0))}
|
|
|
|
|
|
def _get_logs_root() -> str:
|
|
# backend/app/utils -> backend/logs
|
|
return os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "logs")
|
|
|
|
|
|
def _ensure_project_log_dir(project_id: str) -> str:
|
|
root = _get_logs_root()
|
|
project_dir = os.path.join(root, project_id)
|
|
os.makedirs(project_dir, exist_ok=True)
|
|
return project_dir
|
|
|
|
|
|
def _build_log_paths(project_id: str, model_name: str) -> Dict[str, str]:
|
|
project_dir = _ensure_project_log_dir(project_id)
|
|
model_safe = _safe_model_name(model_name)
|
|
return {
|
|
"log": os.path.join(project_dir, f"cost_{model_safe}.log"),
|
|
"jsonl": os.path.join(project_dir, f"cost_{model_safe}.jsonl"),
|
|
}
|
|
|
|
|
|
def _update_counters(counter_key: str, request_cost: Decimal) -> Dict[str, Decimal]:
|
|
with _COUNTER_LOCK:
|
|
current = _COUNTERS.get(counter_key)
|
|
if current is None:
|
|
current = {
|
|
"requests": Decimal("0"),
|
|
"total_cost": Decimal("0"),
|
|
}
|
|
current["requests"] += Decimal("1")
|
|
current["total_cost"] += request_cost
|
|
_COUNTERS[counter_key] = current
|
|
return {
|
|
"requests": current["requests"],
|
|
"total_cost": current["total_cost"],
|
|
}
|
|
|
|
|
|
def record_llm_cost(
|
|
*,
|
|
model_name: str,
|
|
usage: Any,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Calculate and persist LLM usage cost for one request.
|
|
|
|
Returns a normalized cost payload that is also written into JSONL.
|
|
"""
|
|
metadata = metadata or {}
|
|
timestamp = datetime.now().isoformat()
|
|
project_id = _resolve_project_id(metadata)
|
|
component = _resolve_component(metadata)
|
|
usage_dict = _normalize_usage(usage)
|
|
|
|
rates = _resolve_model_rates(model_name)
|
|
input_cost = _quantize_8(_decimal(usage_dict["input_tokens"]) * rates["input"] / _decimal(1_000_000))
|
|
output_cost = _quantize_8(_decimal(usage_dict["output_tokens"]) * rates["output"] / _decimal(1_000_000))
|
|
total_cost = _quantize_8(input_cost + output_cost)
|
|
|
|
record = {
|
|
"timestamp": timestamp,
|
|
"model": model_name,
|
|
"input_tokens": usage_dict["input_tokens"],
|
|
"output_tokens": usage_dict["output_tokens"],
|
|
"total_tokens": usage_dict["total_tokens"],
|
|
"input_cost_usd": float(input_cost),
|
|
"output_cost_usd": float(output_cost),
|
|
"total_cost_usd": float(total_cost),
|
|
"metadata": {
|
|
"component": component,
|
|
"simulation_id": metadata.get("simulation_id"),
|
|
"platform": metadata.get("platform"),
|
|
"phase": metadata.get("phase"),
|
|
"project_id": project_id,
|
|
"report_id": metadata.get("report_id"),
|
|
},
|
|
}
|
|
|
|
paths = _build_log_paths(project_id=project_id, model_name=model_name)
|
|
counter_key = f"{project_id}::{model_name}"
|
|
counter = _update_counters(counter_key, total_cost)
|
|
request_no = int(counter["requests"])
|
|
cumulative_cost = _quantize_8(counter["total_cost"])
|
|
|
|
line = (
|
|
f"[{timestamp}] [Request {request_no}] [{component}] "
|
|
f"Called model: {model_name}, "
|
|
f"input_tokens: {usage_dict['input_tokens']} | "
|
|
f"output_tokens: {usage_dict['output_tokens']} | "
|
|
f"total_tokens: {usage_dict['total_tokens']} | "
|
|
f"input_cost_usd: {float(input_cost):.8f} | "
|
|
f"output_cost_usd: {float(output_cost):.8f} | "
|
|
f"total_cost_usd: {float(total_cost):.8f} | "
|
|
f"cumulative_total_cost_usd: {float(cumulative_cost):.8f}"
|
|
)
|
|
|
|
with open(paths["log"], "a", encoding="utf-8") as log_f:
|
|
log_f.write(line + "\n")
|
|
|
|
with open(paths["jsonl"], "a", encoding="utf-8") as jsonl_f:
|
|
jsonl_f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
return record
|
|
|
|
|
|
def create_tracked_chat_completion(
|
|
*,
|
|
client: Any,
|
|
model: str,
|
|
messages: Any,
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""
|
|
Single entry-point for OpenAI-compatible chat completion + cost logging.
|
|
|
|
All modules should call this wrapper instead of calling
|
|
client.chat.completions.create directly.
|
|
"""
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
**kwargs,
|
|
)
|
|
record_llm_cost(model_name=model, usage=getattr(response, "usage", None), metadata=metadata)
|
|
return response
|
|
|