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

1"""Token counting and cost estimation utilities. 

2 

3Example: 

4 >>> from lexigram.ai.llm import TiktokenCounter 

5 >>> 

6 >>> counter = TiktokenCounter(model="gpt-4") 

7 >>> tokens = counter.count("Hello, world!") 

8 >>> print(f"Tokens: {tokens}") 

9 

10""" 

11 

12from __future__ import annotations 

13 

14from dataclasses import dataclass 

15from datetime import UTC, datetime 

16from typing import Any 

17 

18from lexigram.ai.llm.types import ChatMessage 

19from lexigram.domain import DomainModel 

20from lexigram.logging import ( 

21 get_logger, 

22) 

23from lexigram.validation import Field 

24 

25logger = get_logger(__name__) 

26 

27 

28@dataclass(init=False) 

29class TokenCount(DomainModel): 

30 """Token count result with metadata. 

31 

32 Attributes: 

33 total: Total number of tokens. 

34 prompt_tokens: Number of tokens in the prompt. 

35 completion_tokens: Number of tokens in the completion (if applicable). 

36 model: Model name used for counting. 

37 timestamp: When the count was performed. 

38 

39 """ 

40 

41 total: int = Field(..., description="Total token count") 

42 prompt_tokens: int = Field(default=0, description="Prompt token count") 

43 completion_tokens: int = Field(default=0, description="Completion token count") 

44 model: str = Field(..., description="Model used for counting") 

45 timestamp: datetime = Field( 

46 default_factory=lambda: datetime.now(UTC), 

47 description="Count timestamp", 

48 ) 

49 

50 

51@dataclass(init=False) 

52class CostEstimate(DomainModel): 

53 """Cost estimation result. 

54 

55 Attributes: 

56 prompt_cost: Cost for prompt tokens. 

57 completion_cost: Cost for completion tokens. 

58 total_cost: Total estimated cost. 

59 currency: Currency code (default: USD). 

60 model: Model name. 

61 rate_per_1k_prompt: Rate per 1000 prompt tokens. 

62 rate_per_1k_completion: Rate per 1000 completion tokens. 

63 

64 """ 

65 

66 prompt_cost: float = Field(..., description="Prompt token cost") 

67 completion_cost: float = Field(..., description="Completion token cost") 

68 total_cost: float = Field(..., description="Total cost") 

69 currency: str = Field(default="USD", description="Currency code") 

70 model: str = Field(..., description="Model name") 

71 rate_per_1k_prompt: float = Field(..., description="Rate per 1k prompt tokens") 

72 rate_per_1k_completion: float = Field( 

73 ..., 

74 description="Rate per 1k completion tokens", 

75 ) 

76 

77 

78class TiktokenCounter: 

79 """Token counter using tiktoken (OpenAI/compatible models). 

80 

81 Implements TokenCounterProtocol using tiktoken for precise counting. 

82 tiktoken is a required dependency for this counter. 

83 

84 Args: 

85 model: Model name (e.g. 'gpt-4', 'gpt-3.5-turbo'). 

86 encoding_name: Optional tiktoken encoding name override. 

87 """ 

88 

89 def __init__( 

90 self, 

91 model: str = "gpt-3.5-turbo", 

92 encoding_name: str | None = None, 

93 ) -> None: 

94 """Initialize TiktokenCounter. 

95 

96 Args: 

97 model: Model name for token counting. 

98 encoding_name: Optional tiktoken encoding name override. 

99 

100 Raises: 

101 ImportError: If tiktoken is not installed. 

102 """ 

103 import tiktoken 

104 

105 self._model = model 

106 self._encoding_name = encoding_name 

107 self._encoder: Any | None = None 

108 self._tiktoken = tiktoken 

109 

110 @property 

111 def model(self) -> str: 

112 """The model this counter is calibrated for.""" 

113 return self._model 

114 

115 def _get_encoder(self) -> Any: 

116 """Get tiktoken encoder lazily.""" 

117 if self._encoder is None: 

118 try: 

119 if self._encoding_name: 

120 self._encoder = self._tiktoken.get_encoding(self._encoding_name) 

121 else: 

122 self._encoder = self._tiktoken.encoding_for_model(self._model) 

123 except (KeyError, ValueError): 

124 self._encoder = self._tiktoken.get_encoding("cl100k_base") 

125 return self._encoder 

126 

127 def count(self, text: str) -> int: 

128 """Count tokens in a text string.""" 

129 encoder = self._get_encoder() 

130 return len(encoder.encode(text)) 

131 

132 def count_messages(self, messages: list[ChatMessage]) -> int: 

133 """Count tokens in a list of chat messages, including overhead.""" 

134 encoder = self._get_encoder() 

135 tokens_per_message = 3 

136 tokens_per_name = 1 

137 num_tokens = 0 

138 for msg in messages: 

139 num_tokens += tokens_per_message 

140 num_tokens += len(encoder.encode(str(msg.content or ""))) 

141 if hasattr(msg, "role"): 

142 num_tokens += len(encoder.encode(str(msg.role))) 

143 if hasattr(msg, "name") and msg.name: 

144 num_tokens += tokens_per_name 

145 num_tokens += 3 

146 return num_tokens 

147 

148 

149class HuggingFaceCounter: 

150 """Token counter using HuggingFace AutoTokenizer (lazy-loaded). 

151 

152 When constructed without a model, uses character estimation (~4 chars/token). 

153 When constructed with a model name, lazy-loads that model's tokenizer on first use. 

154 

155 Args: 

156 model: Optional HuggingFace model name. If None, uses char estimation fallback. 

157 """ 

158 

159 def __init__(self, model: str | None = None) -> None: 

160 """Initialize HuggingFaceCounter. 

161 

162 Args: 

163 model: Optional HuggingFace model name for tokenizer loading. 

164 """ 

165 self._model = model 

166 self._tokenizer: Any | None = None 

167 self._loaded: bool = False 

168 

169 @property 

170 def model(self) -> str: 

171 """Backend identifier.""" 

172 return self._model or "huggingface" 

173 

174 def _get_tokenizer(self) -> Any | None: 

175 """Lazily load tokenizer on first use.""" 

176 if self._model is None: 

177 return None # Use char estimation 

178 if not self._loaded: 

179 self._loaded = True 

180 try: 

181 from transformers import AutoTokenizer 

182 

183 self._tokenizer = AutoTokenizer.from_pretrained( 

184 self._model, 

185 use_fast=True, 

186 ) 

187 except ImportError as e: 

188 logger.warning("huggingface_counter_import_failed", error=str(e)) 

189 except OSError as e: 

190 logger.warning( 

191 "huggingface_counter_load_failed", 

192 model=self._model, 

193 error=str(e), 

194 ) 

195 return self._tokenizer 

196 

197 def count(self, text: str) -> int: 

198 """Count tokens in a text string.""" 

199 tok = self._get_tokenizer() 

200 if tok is None: 

201 return max(1, len(text) // 4) 

202 return len(tok.encode(text)) 

203 

204 def count_messages(self, messages: list[ChatMessage]) -> int: 

205 """Count tokens in a list of chat messages.""" 

206 return ( 

207 sum(self.count(str(m.content or "")) for m in messages) + len(messages) * 4 

208 ) 

209 

210 

211class MistralCounter: 

212 """Token counter using mistral-common tokenizer (lazy-loaded). 

213 

214 Tokenizer is loaded on first use, not at construction time. 

215 """ 

216 

217 def __init__(self) -> None: 

218 """Initialize MistralCounter.""" 

219 self._tokenizer: Any | None = None 

220 self._loaded: bool = False 

221 

222 @property 

223 def model(self) -> str: 

224 """Backend identifier.""" 

225 return "mistral" 

226 

227 def _get_tokenizer(self) -> Any | None: 

228 """Lazily load tokenizer on first use.""" 

229 if not self._loaded: 

230 self._loaded = True 

231 try: 

232 from mistral_common.tokens.tokenizers.mistral import ( 

233 MistralTokenizer, 

234 ) 

235 

236 self._tokenizer = MistralTokenizer.v3() 

237 except ImportError as e: 

238 logger.warning("mistral_counter_import_failed", error=str(e)) 

239 except Exception as e: # noqa: BLE001 — tokenizer init raises varied errors (OS, network, parse) 

240 logger.warning("mistral_counter_load_failed", error=str(e)) 

241 return self._tokenizer 

242 

243 def count(self, text: str) -> int: 

244 """Count tokens in a text string.""" 

245 tok = self._get_tokenizer() 

246 if tok is None: 

247 return max(1, len(text) // 4) 

248 try: 

249 # Try the instruct_tokenizer path (mistral-common v1+) 

250 encoded = tok.instruct_tokenizer.tokenizer.encode( 

251 text, bos=False, eos=False 

252 ) 

253 return len(encoded) 

254 except AttributeError: 

255 # Fallback: try direct encode (older API or alternate path) 

256 try: 

257 return len(tok.encode(text)) 

258 except AttributeError: 

259 return max(1, len(text) // 4) 

260 

261 def count_messages(self, messages: list[ChatMessage]) -> int: 

262 """Count tokens in a list of chat messages.""" 

263 return ( 

264 sum(self.count(str(m.content or "")) for m in messages) + len(messages) * 4 

265 ) 

266 

267 

268class CharEstimateCounter: 

269 """Character-based token count estimator (~4 chars per token). 

270 

271 Always available without any optional dependencies. 

272 Suitable as a safe fallback counter. 

273 

274 Args: 

275 model: Model name (used for identification only). 

276 """ 

277 

278 def __init__(self, model: str = "unknown") -> None: 

279 """Initialize CharEstimateCounter. 

280 

281 Args: 

282 model: Model name for identification. 

283 """ 

284 self._model = model 

285 

286 @property 

287 def model(self) -> str: 

288 """The model this counter is calibrated for.""" 

289 return self._model 

290 

291 def count(self, text: str) -> int: 

292 """Count tokens using character estimation.""" 

293 return max(1, len(text) // 4) 

294 

295 def count_messages(self, messages: list[ChatMessage]) -> int: 

296 """Count tokens in a list of chat messages.""" 

297 return ( 

298 sum(self.count(str(m.content or "")) for m in messages) + len(messages) * 4 

299 ) 

300 

301 

302def create_token_counter( 

303 model: str = "gpt-3.5-turbo", 

304 encoding_name: str | None = None, 

305) -> TiktokenCounter: 

306 """Factory function for creating token counters. 

307 

308 Args: 

309 model: Model name. 

310 encoding_name: Optional encoding name override. 

311 

312 Returns: 

313 TiktokenCounter instance. 

314 

315 Example: 

316 >>> from lexigram.ai.llm import create_token_counter 

317 >>> 

318 >>> counter = create_token_counter("gpt-4") 

319 >>> count = counter.count("Hello!") 

320 >>> print(count) 

321 

322 """ 

323 return TiktokenCounter( 

324 model=model, 

325 encoding_name=encoding_name, 

326 )