import redis import json import os import logging from typing import Dict, Tuple # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') class BudgetExceededError(Exception): """Custom exception raised when budget is exceeded.""" pass class EnhancedCostTracker: """ Tracks cost per provider in real-time and enforces budget limits. Persists budget data to Redis for cross-session tracking. """ def __init__(self, config_path="config.json", redis_host="localhost", redis_port=6379, redis_db=0): """ Initializes the cost tracker. Args: config_path (str): Path to the configuration file containing budget information. redis_host (str): Redis host address. redis_port (int): Redis port number. redis_db (int): Redis database number. """ self.redis_client = redis.Redis(host=redis_host, port=redis_port, db=redis_db, decode_responses=True) self.config_path = config_path self.budgets = self._load_budgets() self._initialize_redis() def _load_budgets(self) -> Dict[str, float]: """Loads budget information from the configuration file.""" try: with open(self.config_path, 'r') as f: config = json.load(f) budgets = config.get('budgets', {}) logging.info(f"Budgets loaded from {self.config_path}: {budgets}") return budgets except FileNotFoundError: logging.error(f"Config file not found: {self.config_path}") return {} except json.JSONDecodeError: logging.error(f"Error decoding JSON from {self.config_path}") return {} except Exception as e: logging.error(f"Error loading budgets: {e}") return {} def _initialize_redis(self): """Initializes Redis with budget values if they don't exist.""" redis_key = "genesis:budgets" if not self.redis_client.exists(redis_key): logging.info("Initializing Redis with budget values.") self.redis_client.hset(redis_key, mapping=self.budgets) # Store as strings else: logging.info("Loading budget values from Redis.") # Ensure values loaded from Redis are floats loaded_budgets = self.redis_client.hgetall(redis_key) self.budgets = {k: float(v) for k, v in loaded_budgets.items()} #Convert loaded budgets to float logging.info(f"Loaded budgets from Redis: {self.budgets}") def decrement_budget(self, provider: str, cost: float) -> None: """ Decrements the budget for a given provider. Args: provider (str): The name of the provider. cost (float): The cost to decrement from the budget. Raises: BudgetExceededError: If the budget is exceeded. """ redis_key = "genesis:budgets" current_budget = float(self.redis_client.hget(redis_key, provider)) if current_budget is None: raise ValueError(f"Provider {provider} not found in budgets.") if current_budget - cost < 0: logging.error(f"Budget exceeded for {provider}. Current budget: {current_budget}, cost: {cost}") raise BudgetExceededError(f"Budget exceeded for {provider}") new_budget = current_budget - cost self.redis_client.hset(redis_key, provider, str(new_budget)) # Store as string self.budgets[provider] = new_budget #Update in memory representation self._check_thresholds(provider) logging.info(f"Budget decremented for {provider}. New budget: {new_budget}") def get_current_budget(self, provider: str) -> float: """ Retrieves the current budget for a given provider. Args: provider (str): The name of the provider. Returns: float: The current budget for the provider. """ redis_key = "genesis:budgets" budget = self.redis_client.hget(redis_key, provider) if budget is None: return None return float(budget) def _check_thresholds(self, provider: str) -> None: """ Checks if the budget has reached certain thresholds and emits alerts. Args: provider (str): The name of the provider. """ current_budget = self.get_current_budget(provider) initial_budget = self.budgets[provider] # Get the initial budget from the loaded values. if initial_budget == 0: logging.warning("Initial budget is 0, skipping threshold check.") return remaining_percentage = (current_budget / initial_budget) * 100 if remaining_percentage <= 50 and remaining_percentage > 25: self._emit_alert(provider, "50% threshold reached") elif remaining_percentage <= 25 and remaining_percentage > 10: self._emit_alert(provider, "75% threshold reached") elif remaining_percentage <= 10 and remaining_percentage > 0: self._emit_alert(provider, "90% threshold reached") def _emit_alert(self, provider: str, message: str) -> None: """ Emits an alert message. This is a placeholder; in a real system, this would integrate with an alerting/monitoring service. Args: provider (str): The name of the provider. message (str): The alert message. """ logging.warning(f"ALERT: {message} for provider {provider}. Current budget: {self.get_current_budget(provider)}") def reset_budgets(self) -> None: """Resets all budgets to their initial values from the config file.""" redis_key = "genesis:budgets" self.redis_client.delete(redis_key) self._initialize_redis() # Reinitialize from the config self.budgets = self._load_budgets() # Reload into memory logging.info("All budgets reset to initial values.") if __name__ == '__main__': # Example usage # Create a config.json file with budget information: # { # "budgets": { # "aws": 1000.0, # "azure": 500.0, # "gcp": 750.0 # } # } try: tracker = EnhancedCostTracker(config_path="config.json") # Ensure config.json exists! # Example Usage print("Initial AWS budget:", tracker.get_current_budget("aws")) tracker.decrement_budget("aws", 100.0) print("AWS budget after decrement:", tracker.get_current_budget("aws")) tracker.decrement_budget("azure", 400.0) print("Azure budget after decrement:", tracker.get_current_budget("azure")) # Trigger a threshold alert for i in range(8): try: tracker.decrement_budget("aws", 50.0) except BudgetExceededError as e: print(f"Budget Exceeded: {e}") break print("AWS budget after decrement:", tracker.get_current_budget("aws")) # Try exceeding the budget: try: tracker.decrement_budget("gcp", 800.0) # Exceeds the budget except BudgetExceededError as e: print(f"Budget Exceeded: {e}") # Reset budgets tracker.reset_budgets() print("Budgets reset. Initial AWS budget:", tracker.get_current_budget("aws")) except FileNotFoundError as e: print(f"Error: Config file not found: {e}") except redis.exceptions.ConnectionError as e: print(f"Error: Could not connect to Redis: {e}") except Exception as e: print(f"An unexpected error occurred: {e}")