diff --git a/backend/igny8_core/ai/tracker.py b/backend/igny8_core/ai/tracker.py index cab51748..8eb0b5a2 100644 --- a/backend/igny8_core/ai/tracker.py +++ b/backend/igny8_core/ai/tracker.py @@ -5,6 +5,7 @@ import time import logging from typing import List, Dict, Any, Optional, Callable from datetime import datetime +from decimal import Decimal from igny8_core.ai.constants import DEBUG_MODE logger = logging.getLogger(__name__) @@ -195,24 +196,35 @@ class CostTracker: """Tracks API costs and token usage""" def __init__(self): - self.total_cost = 0.0 + self.total_cost = Decimal('0.0') self.total_tokens = 0 self.operations = [] - def record(self, function_name: str, cost: float, tokens: int, model: str = None): - """Record an API call cost""" + def record(self, function_name: str, cost, tokens: int, model: str = None): + """Record an API call cost + + Args: + function_name: Name of the AI function + cost: Cost value (can be float or Decimal) + tokens: Number of tokens used + model: Model name + """ + # Convert cost to Decimal if it's a float to avoid type mixing + if not isinstance(cost, Decimal): + cost = Decimal(str(cost)) + self.total_cost += cost self.total_tokens += tokens self.operations.append({ 'function': function_name, - 'cost': cost, + 'cost': float(cost), # Store as float for JSON serialization 'tokens': tokens, 'model': model }) - def get_total(self) -> float: - """Get total cost""" - return self.total_cost + def get_total(self): + """Get total cost (returns float for JSON serialization)""" + return float(self.total_cost) def get_total_tokens(self) -> int: """Get total tokens"""