Coverage for lintro / ai / cost.py: 95%

1"""Token counting and cost estimation for AI operations. 

2 

3Provides per-model pricing and utility functions for estimating 

4and formatting costs of AI API calls. 

5""" 

6 

7from __future__ import annotations 

8 

9from loguru import logger 

10 

11from lintro.ai.registry import DEFAULT_PRICING, PROVIDERS 

12 

13 

14def estimate_cost( 

15 model: str, 

16 input_tokens: int, 

17 output_tokens: int, 

18) -> float: 

19 """Estimate the cost of an AI API call. 

20 

21 Args: 

22 model: Model identifier (e.g., "claude-sonnet-4-20250514"). 

23 input_tokens: Number of input tokens. 

24 output_tokens: Number of output tokens. 

25 

26 Returns: 

27 float: Estimated cost in USD. 

28 """ 

29 pricing = PROVIDERS.model_pricing.get(model) 

30 if pricing is None: 

31 logger.debug(f"Unknown model {model!r}, using default pricing") 

32 pricing = DEFAULT_PRICING 

33 

34 input_cost = (input_tokens / 1_000_000) * pricing.input_per_million 

35 output_cost = (output_tokens / 1_000_000) * pricing.output_per_million 

36 

37 return input_cost + output_cost 

38 

39 

40def format_cost(cost: float) -> str: 

41 """Format a cost value for display. 

42 

43 Args: 

44 cost: Cost in USD. 

45 

46 Returns: 

47 str: Formatted cost string (e.g., "$0.003", "<$0.001"). 

48 """ 

49 if cost < 0: 

50 cost = 0 

51 if cost < 0.001: 

52 return "<$0.001" 

53 return f"${cost:.3f}" 

54 

55 

56def format_token_count(tokens: int) -> str: 

57 """Format a token count for display. 

58 

59 Args: 

60 tokens: Number of tokens. 

61 

62 Returns: 

63 str: Formatted token count (e.g., "~1,230"). 

64 """ 

65 return f"~{tokens:,}"