Coverage for packages/python/voice_prices/data_snapshot.py: 97.54%

1from __future__ import annotations as _annotations (empty)

2 

3import re (empty)

4from dataclasses import dataclass, field (empty)

5from datetime import datetime, timedelta, timezone (empty)

6from functools import cache (empty)

7from typing import Any (empty)

8 

9from . import types (empty)

10 

11__all__ = 'DataSnapshot', 'set_custom_snapshot' (empty)

12 

13# snapshot set by UpdatePrices, or manually by the user 

14_custom_snapshot: DataSnapshot | None = None (empty)

15 

16 

17def get_snapshot() -> DataSnapshot: (empty)

18 if _custom_snapshot is not None: (empty)

19 return _custom_snapshot (empty)

20 return _bundled_snapshot() (empty)

21 

22 

23@cache (empty)

24def _bundled_snapshot() -> DataSnapshot: (empty)

25 from .data import providers (empty)

26 

27 return DataSnapshot(providers=providers, from_auto_update=False) (empty)

28 

29 

30def set_custom_snapshot(snapshot: DataSnapshot | None): (empty)

31 """Install a snapshot for every later price lookup in this process, or reset with None. 

32 

33 Process-wide by design: `get_snapshot` reads this module global, so one call changes what 

34 every subsequent `calc_price` sees. `UpdatePrices` calls it from a background daemon thread, 

35 which means a process that has started an updater can return a different price for two 

36 identical calls, and a test that has not reset can see downloaded data instead of the 

37 snapshot bundled with the installed version. 

38 

39 `set_custom_snapshot(None)` is the supported reset and restores the bundled data. Prefer 

40 `with UpdatePrices():`, whose `stop()` resets for you after joining the thread. 

41 """ 

42 global _custom_snapshot 

43 _custom_snapshot = snapshot (empty)

44 

45 

46@dataclass (empty)

47class DataSnapshot: (empty)

48 providers: list[types.Provider] (empty)

49 from_auto_update: bool (empty)

50 _lookup_cache: dict[tuple[str | None, str | None, str], tuple[types.Provider, types.ModelInfo]] = field( (empty)

51 default_factory=lambda: {} 

52 ) 

53 timestamp: datetime = field(default_factory=datetime.now) (empty)

54 

55 def active(self, ttl: timedelta) -> bool: (empty)

56 """Check if the snapshot is "active" (e.g. hasn't expired) based on a time to live.""" 

57 return self.timestamp + ttl > datetime.now() 

58 

59 def calc( (empty)

60 self, 

61 usage: types.AbstractUsage, 

62 model_ref: str, 

63 provider_id: str | None, 

64 provider_api_url: str | None, 

65 genai_request_timestamp: datetime | None, 

66 ) -> types.PriceCalculation: 

67 """Calculate the price for the given usage.""" 

68 genai_request_timestamp = genai_request_timestamp or datetime.now(tz=timezone.utc) (empty)

69 

70 provider, model = self.find_provider_model(model_ref, None, provider_id, provider_api_url) (empty)

71 return model.calc_price( (empty)

72 usage, 

73 provider, 

74 genai_request_timestamp=genai_request_timestamp, 

75 auto_update_timestamp=self.timestamp if self.from_auto_update else None, 

76 ) 

77 

78 def extract_usage( (empty)

79 self, 

80 response_data: Any, 

81 provider_id: types.ProviderID | str | None = None, 

82 provider_api_url: str | None = None, 

83 api_flavor: str = 'default', 

84 ) -> types.ExtractedUsage: 

85 provider = self.find_provider(None, provider_id, provider_api_url) (empty)

86 model_ref, usage = provider.extract_usage(response_data, api_flavor=api_flavor) (empty)

87 if model_ref is not None: (empty)

88 _, model = self.find_provider_model(model_ref, provider, None, None) (empty)

89 else: 

90 model = None (empty)

91 return types.ExtractedUsage(usage, model, provider, self.timestamp if self.from_auto_update else None) (empty)

92 

93 def find_provider_model( (empty)

94 self, 

95 model_ref: str, 

96 provider: types.Provider | None, 

97 provider_id: str | None, 

98 provider_api_url: str | None, 

99 ) -> tuple[types.Provider, types.ModelInfo]: 

100 """Find the provider and model for the given model reference and optional provider identifier.""" 

101 model_ref = model_ref.lower() (empty)

102 

103 # Handle litellm provider_id by extracting actual provider from model name prefix 

104 if provider_id and provider_id.lower() == 'litellm' and '/' in model_ref: (empty)

105 actual_provider_id, actual_model_ref = model_ref.split('/', 1) (empty)

106 # Only use the extracted provider if it exists 

107 if actual_provider_id and find_provider_by_id(self.providers, actual_provider_id): 107 ↛ 111line 107 didn't jump to line 111 because the condition on line 107 was always true(empty)

108 provider_id = actual_provider_id (empty)

109 model_ref = actual_model_ref (empty)

110 

111 if provider: (empty)

112 if provider_model := self._lookup_cache.get((provider.id, None, model_ref)): (empty)

113 return provider_model (empty)

114 else: 

115 if provider_model := self._lookup_cache.get((provider_id, provider_api_url, model_ref)): (empty)

116 return provider_model (empty)

117 

118 provider = self.find_provider(model_ref, provider_id, provider_api_url) (empty)

119 

120 if model := provider.find_model(model_ref, all_providers=self.providers): (empty)

121 self._lookup_cache[(provider_id, provider_api_url, model_ref)] = ret = provider, model (empty)

122 return ret (empty)

123 else: 

124 raise LookupError(f'Unable to find model with {model_ref=!r} in {provider.id}') (empty)

125 

126 def find_provider( (empty)

127 self, 

128 model_ref: str | None, 

129 provider_id: str | None, 

130 provider_api_url: str | None, 

131 ) -> types.Provider: 

132 if provider_id is not None: (empty)

133 if provider := find_provider_by_id(self.providers, provider_id): (empty)

134 return provider (empty)

135 # Special case for litellm: fall back to model matching if provider not found 

136 if provider_id.lower() != 'litellm': (empty)

137 raise LookupError(f'Unable to find provider {provider_id=!r}') (empty)

138 

139 if provider_api_url is not None: (empty)

140 for provider in self.providers: (empty)

141 if re.match(provider.api_pattern, provider_api_url): (empty)

142 return provider (empty)

143 raise LookupError(f'Unable to find provider {provider_api_url=!r}') (empty)

144 

145 if model_ref: 145 ↛ 150line 145 didn't jump to line 150 because the condition on line 145 was always true(empty)

146 for provider in self.providers: (empty)

147 if provider.model_match is not None and provider.model_match.is_match(model_ref): (empty)

148 return provider (empty)

149 

150 raise LookupError(f'Unable to find provider with model matching {model_ref!r}') (empty)

151 

152 

153def find_provider_by_id(providers: list[types.Provider], provider_id: str) -> types.Provider | None: (empty)

154 """Find a provider by matching against provider_match logic. 

155 

156 Args: 

157 providers: List of available providers 

158 provider_id: The provider ID to match 

159 

160 Returns: 

161 The matching provider or None 

162 """ 

163 normalized_provider_id = provider_id.lower().strip() (empty)

164 

165 for provider in providers: (empty)

166 if provider.id == normalized_provider_id: (empty)

167 return provider (empty)

168 

169 for provider in providers: (empty)

170 if provider.provider_match and provider.provider_match.is_match(normalized_provider_id): (empty)

171 return provider (empty)

172 

173 return None (empty)