Coverage for llm_dataset_engine/utils/budget_controller.py: 93%

1""" 

2Budget control and enforcement for LLM costs. 

3 

4Implements cost monitoring with threshold warnings and hard limits. 

5""" 

6 

7from decimal import Decimal 

8from typing import Optional 

9 

10import structlog 

11 

12logger = structlog.get_logger(__name__) 

13 

14 

15class BudgetExceededError(Exception): 

16 """Raised when budget limit is exceeded.""" 

17 

18 pass 

19 

20 

21class BudgetController: 

22 """ 

23 Controls and enforces budget limits during execution. 

24  

25 Follows Single Responsibility: only handles budget management. 

26 """ 

27 

28 def __init__( 

29 self, 

30 max_budget: Optional[Decimal] = None, 

31 warn_at_75: bool = True, 

32 warn_at_90: bool = True, 

33 fail_on_exceed: bool = True, 

34 ): 

35 """ 

36 Initialize budget controller. 

37 

38 Args: 

39 max_budget: Maximum allowed budget in USD 

40 warn_at_75: Warn at 75% of budget 

41 warn_at_90: Warn at 90% of budget 

42 fail_on_exceed: Raise error if budget exceeded 

43 """ 

44 self.max_budget = max_budget 

45 self.warn_at_75 = warn_at_75 

46 self.warn_at_90 = warn_at_90 

47 self.fail_on_exceed = fail_on_exceed 

48 

49 self._warned_75 = False 

50 self._warned_90 = False 

51 

52 def check_budget(self, current_cost: Decimal) -> None: 

53 """ 

54 Check if budget is within limits. 

55 

56 Args: 

57 current_cost: Current accumulated cost 

58 

59 Raises: 

60 BudgetExceededError: If budget exceeded and fail_on_exceed=True 

61 """ 

62 if self.max_budget is None: 

63 return 

64 

65 usage_ratio = float(current_cost / self.max_budget) 

66 

67 # 75% warning 

68 if ( 

69 self.warn_at_75 

70 and not self._warned_75 

71 and usage_ratio >= 0.75 

72 ): 

73 logger.warning( 

74 f"Budget warning: 75% used " 

75 f"(${current_cost:.4f} / ${self.max_budget:.2f})" 

76 ) 

77 self._warned_75 = True 

78 

79 # 90% warning 

80 if ( 

81 self.warn_at_90 

82 and not self._warned_90 

83 and usage_ratio >= 0.90 

84 ): 

85 logger.warning( 

86 f"Budget warning: 90% used " 

87 f"(${current_cost:.4f} / ${self.max_budget:.2f})" 

88 ) 

89 self._warned_90 = True 

90 

91 # Budget exceeded 

92 if current_cost > self.max_budget: 

93 error_msg = ( 

94 f"Budget exceeded: ${current_cost:.4f} > ${self.max_budget:.2f}" 

95 ) 

96 logger.error(error_msg) 

97 

98 if self.fail_on_exceed: 

99 raise BudgetExceededError(error_msg) 

100 

101 def get_remaining(self, current_cost: Decimal) -> Optional[Decimal]: 

102 """ 

103 Get remaining budget. 

104 

105 Args: 

106 current_cost: Current accumulated cost 

107 

108 Returns: 

109 Remaining budget or None if no limit 

110 """ 

111 if self.max_budget is None: 

112 return None 

113 return self.max_budget - current_cost 

114 

115 def get_usage_percentage(self, current_cost: Decimal) -> Optional[float]: 

116 """ 

117 Get budget usage as percentage. 

118 

119 Args: 

120 current_cost: Current accumulated cost 

121 

122 Returns: 

123 Usage percentage or None if no limit 

124 """ 

125 if self.max_budget is None: 

126 return None 

127 return float(current_cost / self.max_budget) * 100 

128 

129 def can_afford( 

130 self, estimated_cost: Decimal, current_cost: Decimal 

131 ) -> bool: 

132 """ 

133 Check if estimated additional cost is within budget. 

134 

135 Args: 

136 estimated_cost: Estimated cost for next operation 

137 current_cost: Current accumulated cost 

138 

139 Returns: 

140 True if within budget 

141 """ 

142 if self.max_budget is None: 

143 return True 

144 return (current_cost + estimated_cost) <= self.max_budget 

145