Coverage for lintro / ai / budget.py: 100%

1"""Session cost budget tracker.""" 

2 

3from __future__ import annotations 

4 

5import threading 

6from dataclasses import dataclass, field 

7 

8 

9@dataclass 

10class CostBudget: 

11 """Track cumulative AI session cost against an optional ceiling. 

12 

13 Attributes: 

14 max_cost_usd: Maximum total cost in USD per AI session. 

15 None disables the limit. 

16 """ 

17 

18 max_cost_usd: float | None = None 

19 _spent: float = field(default=0.0, init=False) 

20 _lock: threading.Lock = field( 

21 default_factory=threading.Lock, 

22 init=False, 

23 repr=False, 

24 ) 

25 

26 def record(self, cost: float) -> None: 

27 """Record a cost increment.""" 

28 with self._lock: 

29 self._spent += cost 

30 

31 @property 

32 def spent(self) -> float: 

33 """Return total cost spent so far.""" 

34 with self._lock: 

35 return self._spent 

36 

37 @property 

38 def remaining(self) -> float | None: 

39 """Return remaining budget, or None if unlimited.""" 

40 if self.max_cost_usd is None: 

41 return None 

42 return max(0.0, self.max_cost_usd - self.spent) 

43 

44 def check(self) -> None: 

45 """Raise AIError if the budget has been exceeded.""" 

46 if self.max_cost_usd is not None and self.spent >= self.max_cost_usd: 

47 from lintro.ai.exceptions import AIError 

48 

49 raise AIError( 

50 f"AI cost budget exceeded: ${self.spent:.4f} spent, " 

51 f"limit is ${self.max_cost_usd:.2f}", 

52 )