Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/pricing/estimator.py: 32%

1"""Synchronous cost estimator over preloaded pricing data. 

2 

3Provides :class:`PricingCostEstimator`, a concrete implementation of 

4:class:`~lexigram.contracts.ai.llm.CostEstimatorProtocol` that prices 

5token usage against a snapshot of :class:`ModelPricing` entries. 

6 

7The snapshot is preloaded at container boot (see ``LLMProvider``); the 

8estimator itself is fully synchronous so it can be called on the hot 

9path without blocking the event loop. 

10""" 

11 

12from __future__ import annotations 

13 

14from typing import TYPE_CHECKING 

15 

16from lexigram.ai.llm.pricing.types import ModelPricing 

17from lexigram.logging import ( 

18 get_logger, 

19) 

20 

21if TYPE_CHECKING: 

22 from lexigram.ai.llm.pricing.manager import PricingManager 

23 

24logger = get_logger(__name__) 

25 

26 

27class PricingCostEstimator: 

28 """Estimate USD cost of LLM usage from a pricing snapshot. 

29 

30 Implements ``CostEstimatorProtocol.estimate_cost``. Model lookup is 

31 exact (case-insensitive), then a substring match in both directions 

32 (mirrors ``PricingManager`` fuzzy matching) when enabled. Unknown 

33 models price at ``0.0`` — callers must skip cost tracking rather than 

34 fabricate estimates. 

35 

36 Attributes: 

37 pricing: Snapshot of model name to pricing data. 

38 enable_fuzzy_match: Allow substring model name matching. 

39 

40 Example: 

41 >>> estimator = PricingCostEstimator( 

42 ... {"gpt-4o": ModelPricing(model="gpt-4o", prompt_per_1m=2.5, 

43 ... completion_per_1m=10.0, provider="openai")} 

44 ... ) 

45 >>> estimator.estimate_cost("gpt-4o", 1500, prompt_tokens=1000, 

46 ... completion_tokens=500) 

47 0.0075 

48 

49 """ 

50 

51 def __init__( 

52 self, 

53 pricing: dict[str, ModelPricing], 

54 *, 

55 enable_fuzzy_match: bool = True, 

56 ) -> None: 

57 """Initialize the estimator. 

58 

59 Args: 

60 pricing: Preloaded snapshot of model name to pricing data. 

61 enable_fuzzy_match: Allow substring model name matching 

62 (default: True). 

63 """ 

64 self.pricing = {k.lower(): v for k, v in pricing.items()} 

65 self.enable_fuzzy_match = enable_fuzzy_match 

66 

67 async def warm(self, manager: PricingManager) -> None: 

68 """Reload the pricing snapshot from a manager. 

69 

70 Aggregates all sources via :meth:`PricingManager.preload`; source 

71 failures degrade to empty entries, which price at ``0.0``. 

72 

73 Args: 

74 manager: Pricing manager to load from. 

75 """ 

76 merged = await manager.preload() 

77 self.pricing = {k.lower(): v for k, v in merged.items()} 

78 logger.info( 

79 "Cost estimator warmed with %d models", 

80 len(self.pricing), 

81 ) 

82 

83 def estimate_cost( 

84 self, 

85 model: str, 

86 total_tokens: int, 

87 provider: str | None = None, 

88 prompt_tokens: int = 0, 

89 completion_tokens: int = 0, 

90 ) -> float: 

91 """Estimate cost in USD for the given token usage. 

92 

93 When *prompt_tokens* and *completion_tokens* are known they are 

94 priced at the input and output rates. When both are ``0`` the 

95 split is unknown and *total_tokens* is priced at the input rate — 

96 prompt tokens dominate agent turns, so this is the closest 

97 approximation available. 

98 

99 Args: 

100 model: Model identifier (e.g. ``gpt-4o``). 

101 total_tokens: Total tokens consumed (prompt + completion). 

102 provider: Provider name (e.g. ``openai``) — currently unused, 

103 pricing is keyed by model. 

104 prompt_tokens: Input token count. ``0`` means unknown. 

105 completion_tokens: Output token count. ``0`` means unknown. 

106 

107 Returns: 

108 Estimated cost in USD. ``0.0`` when the model has no pricing. 

109 

110 Example: 

111 >>> estimator.estimate_cost("gpt-4o", 1500, prompt_tokens=1000, 

112 ... completion_tokens=500) 

113 0.0075 

114 """ 

115 entry = self._lookup(model) 

116 if entry is None: 

117 return 0.0 

118 

119 if prompt_tokens > 0 or completion_tokens > 0: 

120 prompt_cost = (prompt_tokens / 1_000_000) * entry.prompt_per_1m 

121 completion_cost = (completion_tokens / 1_000_000) * entry.completion_per_1m 

122 return prompt_cost + completion_cost 

123 

124 # Unknown split: price the total at the input rate (prompt tokens 

125 # dominate agent turns). 

126 return (total_tokens / 1_000_000) * entry.prompt_per_1m 

127 

128 def _lookup(self, model: str) -> ModelPricing | None: 

129 """Resolve pricing for a model name. 

130 

131 Args: 

132 model: Model identifier. 

133 

134 Returns: 

135 Matching pricing entry or None. 

136 """ 

137 normalized = model.lower().strip() 

138 entry = self.pricing.get(normalized) 

139 if entry is not None: 

140 return entry 

141 

142 if not self.enable_fuzzy_match: 

143 return None 

144 

145 for known_model, candidate in self.pricing.items(): 

146 if normalized in known_model or known_model in normalized: 

147 return candidate 

148 

149 return None 

150 

151 

152__all__ = ["PricingCostEstimator"]