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

1"""

2 Budget control and enforcement for LLM costs.

3

4 Implements cost monitoring with threshold warnings and hard limits.

5"""

6

7 from decimal import Decimal

8 from typing import Optional

9

10 import structlog

11

12 logger = structlog.get_logger(__name__)

13

14

15 class BudgetExceededError(Exception):

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

17

18 pass

19

20

21 class 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