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

1"""Token counter registry for managing model-to-counter mappings.""" 

2 

3from __future__ import annotations 

4 

5import re 

6 

7from lexigram.contracts.ai.llm import TokenCounterProtocol 

8from lexigram.logging import ( 

9 get_logger, 

10) 

11 

12logger = get_logger(__name__) 

13 

14 

15def _tiktoken_available() -> bool: 

16 """Check if tiktoken is installed.""" 

17 try: 

18 import tiktoken # noqa: F401 

19 

20 return True 

21 except ImportError: 

22 return False 

23 

24 

25def _transformers_available() -> bool: 

26 """Check if HuggingFace transformers is installed.""" 

27 try: 

28 import transformers # noqa: F401 

29 

30 return True 

31 except ImportError: 

32 return False 

33 

34 

35def _mistral_available() -> bool: 

36 """Check if mistral-common is installed.""" 

37 try: 

38 import mistral_common # noqa: F401 

39 

40 return True 

41 except ImportError: 

42 return False 

43 

44 

45class TokenCounterRegistry: 

46 """Registry mapping model-name patterns to TokenCounterProtocol backends. 

47 

48 Uses named backend keys and regex patterns for flexible model mapping. 

49 

50 Usage:: 

51 

52 registry = TokenCounterRegistry.with_defaults() 

53 counter = registry.for_model("gpt-4o") 

54 tokens = counter.count("Hello!") 

55 """ 

56 

57 def __init__(self) -> None: 

58 """Create an empty registry.""" 

59 self._backends: dict[str, TokenCounterProtocol] = {} 

60 self._patterns: list[tuple[re.Pattern[str], str]] = [] 

61 

62 @classmethod 

63 def with_defaults(cls) -> TokenCounterRegistry: 

64 """Create registry with all available tokenizer backends. 

65 

66 Registers: 

67 - char_estimate (always available, fallback) 

68 - tiktoken (if installed, for OpenAI/Anthropic models) 

69 - huggingface (if installed, for HuggingFace models) 

70 - mistral (if installed, for Mistral models) 

71 

72 Returns: 

73 TokenCounterRegistry pre-populated with default backends. 

74 """ 

75 from lexigram.ai.llm.pricing.tokens import CharEstimateCounter 

76 

77 registry = cls() 

78 registry.register("char_estimate", CharEstimateCounter()) # type: ignore[arg-type] 

79 

80 if _tiktoken_available(): 

81 from lexigram.ai.llm.pricing.tokens import TiktokenCounter 

82 

83 registry.register("tiktoken", TiktokenCounter()) # type: ignore[arg-type] 

84 registry.map_models(r"gpt-.*|o[0-9].*|text-embedding-.*", "tiktoken") 

85 logger.debug("token_counter_registry_tiktoken_registered") 

86 else: 

87 logger.warning( 

88 "token_counter_registry_tiktoken_unavailable", 

89 fallback="char_estimate", 

90 ) 

91 

92 if _transformers_available(): 

93 from lexigram.ai.llm.pricing.tokens import HuggingFaceCounter 

94 

95 registry.register("huggingface", HuggingFaceCounter()) # type: ignore[arg-type] 

96 registry.map_models( 

97 r"llama-.*|qwen-.*|deepseek-.*|gemma-.*", 

98 "huggingface", 

99 ) 

100 logger.debug("token_counter_registry_huggingface_registered") 

101 

102 if _mistral_available(): 

103 from lexigram.ai.llm.pricing.tokens import MistralCounter 

104 

105 registry.register("mistral", MistralCounter()) # type: ignore[arg-type] 

106 registry.map_models(r"mistral-.*|codestral-.*", "mistral") 

107 logger.debug("token_counter_registry_mistral_registered") 

108 

109 return registry 

110 

111 def register(self, key: str, counter: TokenCounterProtocol) -> None: 

112 """Register a counter backend under a named key. 

113 

114 Args: 

115 key: Backend name (e.g., 'tiktoken', 'huggingface', 'char_estimate'). 

116 counter: Counter implementing TokenCounterProtocol. 

117 """ 

118 self._backends[key] = counter 

119 

120 def map_models(self, pattern: str, counter_key: str) -> None: 

121 """Map a regex pattern of model names to a backend key. 

122 

123 Args: 

124 pattern: Regex pattern matching model names (case-insensitive). 

125 counter_key: Backend key (must be registered). 

126 """ 

127 if counter_key not in self._backends: 

128 logger.warning("token_counter_map_key_not_found", key=counter_key) 

129 return 

130 self._patterns.append((re.compile(pattern, re.IGNORECASE), counter_key)) 

131 

132 def for_model(self, model: str) -> TokenCounterProtocol: 

133 """Get the best counter for the given model name. 

134 

135 Tries exact regex match in _patterns first, falls back to 'char_estimate'. 

136 

137 Args: 

138 model: Model name. 

139 

140 Returns: 

141 TokenCounterProtocol implementation. 

142 """ 

143 for compiled_pattern, key in self._patterns: 

144 if compiled_pattern.match(model): 

145 return self._backends[key] 

146 if "char_estimate" in self._backends: 

147 return self._backends["char_estimate"] 

148 from lexigram.ai.llm.pricing.tokens import CharEstimateCounter 

149 

150 return CharEstimateCounter() # type: ignore[return-value]