| import threading |
| import pandas as pd |
| from dataclasses import dataclass |
|
|
| from ..core.logging import logger |
| from ..core.decorators import atomic_method |
| from ..core.callbacks import suppress_cost_logs |
| from ..core.registry import MODEL_REGISTRY |
| from .model_configs import LLMConfig |
| from ..models.base_model import BaseLLM |
|
|
| def get_openai_model_cost() -> dict: |
| import json |
| from importlib.resources import files |
| |
| |
| |
| json_path = files('litellm') / 'model_prices_and_context_window_backup.json' |
| model_cost = json.loads(json_path.read_text(encoding="utf-8")) |
| return model_cost |
|
|
| def infer_litellm_company_from_model(model: str) -> str: |
|
|
| if "/" in model: |
| company = model.split("/")[0] |
| else: |
| if "claude" in model or "anthropic" in model: |
| company = "anthropic" |
| elif "gemini" in model: |
| company = "gemini" |
| elif "deepseek" in model: |
| company = "deepseek" |
| elif "openrouter" in model: |
| company = "openrouter" |
| elif "azure" in model.lower(): |
| company = "azure" |
| else: |
| company = "openai" |
| return company |
|
|
|
|
| @dataclass |
| class Cost: |
| input_tokens: int |
| output_tokens: int |
| input_cost: float |
| output_cost: float |
|
|
|
|
| class CostManager: |
|
|
| def __init__(self): |
|
|
| self.total_input_tokens = {} |
| self.total_output_tokens = {} |
| self.total_tokens = {} |
|
|
| self.total_input_cost = {} |
| self.total_output_cost = {} |
| self.total_cost = {} |
|
|
| self._lock = threading.Lock() |
|
|
| def compute_total_cost(self): |
| total_tokens, total_cost = 0, 0.0 |
| for _, value in self.total_tokens.items(): |
| total_tokens += value |
| for _, value in self.total_cost.items(): |
| total_cost += value |
| return total_tokens, total_cost |
|
|
| @atomic_method |
| def update_cost(self, cost: Cost, model: str): |
|
|
| self.total_input_tokens[model] = self.total_input_tokens.get(model, 0) + cost.input_tokens |
| self.total_output_tokens[model] = self.total_output_tokens.get(model, 0) + cost.output_tokens |
| current_total_tokens = cost.input_tokens + cost.output_tokens |
| self.total_tokens[model] = self.total_tokens.get(model, 0) + current_total_tokens |
|
|
| self.total_input_cost[model] = self.total_input_cost.get(model, 0.0) + cost.input_cost |
| self.total_output_cost[model] = self.total_output_cost.get(model, 0.0) + cost.output_cost |
| current_total_cost = cost.input_cost + cost.output_cost |
| self.total_cost[model] = self.total_cost.get(model, 0.0) + current_total_cost |
| |
| total_tokens, total_cost = self.compute_total_cost() |
| if not suppress_cost_logs.get(): |
| logger.info(f"Total cost: ${total_cost:.3f} | Total tokens: {total_tokens} | Current cost: ${current_total_cost:.3f} | Current tokens: {current_total_tokens}") |
|
|
| def display_cost(self): |
|
|
| data = { |
| "Model": [], |
| "Total Cost (USD)": [], |
| "Total Input Cost (USD)": [], |
| "Total Output Cost (USD)": [], |
| "Total Tokens": [], |
| "Total Input Tokens": [], |
| "Total Output Tokens": [], |
| } |
|
|
| for model in self.total_tokens.keys(): |
|
|
| data["Model"].append(model) |
| data["Total Cost (USD)"].append(round(self.total_cost[model], 4)) |
| data["Total Input Cost (USD)"].append(round(self.total_input_cost[model], 4)) |
| data["Total Output Cost (USD)"].append(round(self.total_output_cost[model], 4)) |
|
|
| data["Total Tokens"].append(self.total_tokens[model]) |
| data["Total Input Tokens"].append(self.total_input_tokens[model]) |
| data["Total Output Tokens"].append(self.total_output_tokens[model]) |
| |
| |
| df = pd.DataFrame(data) |
| if len(df) > 1: |
| summary = { |
| "Model": "TOTAL", |
| "Total Cost (USD)": df["Total Cost (USD)"].sum(), |
| "Total Input Cost (USD)": df["Total Input Cost (USD)"].sum(), |
| "Total Output Cost (USD)": df["Total Output Cost (USD)"].sum(), |
| "Total Tokens": df["Total Tokens"].sum(), |
| "Total Input Tokens": df["Total Input Tokens"].sum(), |
| "Total Output Tokens": df["Total Output Tokens"].sum(), |
| } |
| df = df._append(summary, ignore_index=True) |
| |
| print(df.to_string(index=False)) |
|
|
| def get_total_cost(self): |
| |
| total_cost = 0.0 |
| for model in self.total_cost.keys(): |
| total_cost += self.total_cost[model] |
| return total_cost |
|
|
|
|
| cost_manager = CostManager() |
|
|
|
|
| def create_llm_instance(llm_config: LLMConfig) -> BaseLLM: |
|
|
| llm_cls = MODEL_REGISTRY.get_model(llm_config.llm_type) |
| llm = llm_cls(config=llm_config) |
| return llm |
|
|
