Coverage for tests/test_custom_prices.py: 98.85%

1from __future__ import annotations (empty)

2 

3from copy import deepcopy (empty)

4from dataclasses import dataclass (empty)

5from decimal import Decimal (empty)

6 

7import pytest (empty)

8from inline_snapshot import snapshot (empty)

9 

10from voice_prices import Usage, calc_price, data, types (empty)

11from voice_prices.data_snapshot import DataSnapshot, set_custom_snapshot (empty)

12from voice_prices.update_prices import UpdatePrices (empty)

13 

14 

15@dataclass (empty)

16class CustomModelPrice(types.ModelPrice): (empty)

17 sausage_price: Decimal | None = None (empty)

18 

19 def calc_price(self, usage: types.AbstractUsage) -> types.CalcPrice: (empty)

20 price = super().calc_price(usage) (empty)

21 if isinstance(usage, CustomUsage) and self.sausage_price is not None: 21 ↛ 23line 21 didn't jump to line 23 because the condition on line 21 was always true(empty)

22 price['total_price'] += self.sausage_price * usage.sausages (empty)

23 return price (empty)

24 

25 

26class AltUpdatePrices(UpdatePrices): (empty)

27 def fetch(self) -> DataSnapshot | None: (empty)

28 custom_providers = [ (empty)

29 types.Provider( 

30 id='testing', 

31 name='Testing', 

32 api_pattern=r'.*testing\.example\.com', 

33 models=[ 

34 types.ModelInfo( 

35 id='foobar', 

36 match=types.ClauseEquals('foobar'), 

37 name='Foobar', 

38 prices=types.ModelPrice(input_mtok=Decimal('1'), output_mtok=Decimal('2')), 

39 ), 

40 types.ModelInfo( 

41 id='sausage', 

42 match=types.ClauseStartsWith('sausage'), 

43 name='Sausage', 

44 prices=CustomModelPrice( 

45 input_mtok=Decimal('1'), 

46 output_mtok=Decimal('2'), 

47 sausage_price=Decimal('3'), 

48 ), 

49 ), 

50 ], 

51 ) 

52 ] 

53 return DataSnapshot(providers=custom_providers, from_auto_update=False) (empty)

54 

55 

56class ExtraUpdatePrices(UpdatePrices): (empty)

57 def fetch(self) -> DataSnapshot | None: (empty)

58 providers = deepcopy(data.providers) (empty)

59 openai = next(provider for provider in providers if provider.id == 'openai') (empty)

60 openai.models.insert( (empty)

61 0, 

62 types.ModelInfo( 

63 id='gpt-4o', 

64 match=types.ClauseStartsWith('gpt-4o'), 

65 name='gpt-4o Custom', 

66 prices=CustomModelPrice( 

67 input_mtok=Decimal('1'), 

68 output_mtok=Decimal('2'), 

69 sausage_price=Decimal('3'), 

70 ), 

71 ), 

72 ) 

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

74 

75 

76@dataclass (empty)

77class CustomUsage: (empty)

78 sausages: int (empty)

79 

80 requests: int | None = None (empty)

81 

82 input_tokens: int | None = None (empty)

83 

84 cache_write_tokens: int | None = None (empty)

85 cache_read_tokens: int | None = None (empty)

86 

87 output_tokens: int | None = None (empty)

88 

89 input_audio_tokens: int | None = None (empty)

90 cache_audio_read_tokens: int | None = None (empty)

91 output_audio_tokens: int | None = None (empty)

92 

93 # TTS + STT fields added in v0.x / v0.0.7 of the voice-prices fork; included 

94 # here so CustomUsage continues to satisfy the AbstractUsage protocol. 

95 # Duration fields migrated from int to Decimal in v0.0.7. 

96 characters: int | None = None (empty)

97 audio_output_seconds: Decimal | None = None (empty)

98 audio_input_seconds: Decimal | None = None (empty)

99 voice_class: str | None = None (empty)

100 

101 

102def test_alt_source(): (empty)

103 with AltUpdatePrices() as update_prices: (empty)

104 update_prices.wait() (empty)

105 price = calc_price( (empty)

106 Usage(input_tokens=1_000_000, output_tokens=1_000_000), 

107 model_ref='foobar', 

108 provider_id='testing', 

109 ) 

110 assert price.input_price == snapshot(Decimal('1')) (empty)

111 assert price.output_price == snapshot(Decimal('2')) (empty)

112 assert price.total_price == snapshot(Decimal('3')) (empty)

113 assert price.model.name == snapshot('Foobar') (empty)

114 assert price.provider.id == snapshot('testing') (empty)

115 assert price.auto_update_timestamp is None (empty)

116 

117 with pytest.raises(LookupError, match="Unable to find provider provider_id='testing'"): (empty)

118 calc_price( (empty)

119 Usage(input_tokens=1_000_000, output_tokens=1_000_000), 

120 model_ref='foobar', 

121 provider_id='testing', 

122 ) 

123 

124 

125def test_alt_source_sausage(): (empty)

126 set_custom_snapshot(AltUpdatePrices().fetch()) (empty)

127 try: (empty)

128 price = calc_price( (empty)

129 CustomUsage(sausages=3, input_tokens=1_000_000, output_tokens=1_000_000), 

130 model_ref='sausage', 

131 provider_id='testing', 

132 ) 

133 assert price.input_price == snapshot(Decimal('1')) (empty)

134 assert price.output_price == snapshot(Decimal('2')) (empty)

135 assert price.total_price == snapshot(Decimal('12')) (empty)

136 assert price.model.name == snapshot('Sausage') (empty)

137 assert price.provider.id == snapshot('testing') (empty)

138 assert price.auto_update_timestamp is None (empty)

139 finally: 

140 set_custom_snapshot(None) (empty)

141 

142 

143def test_extra_source_normal(): (empty)

144 with ExtraUpdatePrices() as update_prices: (empty)

145 update_prices.wait() (empty)

146 price = calc_price( (empty)

147 CustomUsage(sausages=3, input_tokens=1_000_000, output_tokens=1_000_000), 

148 model_ref='gpt-4', 

149 provider_id='openai', 

150 ) 

151 assert price.input_price == snapshot(Decimal('30')) (empty)

152 assert price.output_price == snapshot(Decimal('60')) (empty)

153 assert price.total_price == snapshot(Decimal('90')) (empty)

154 assert price.model.name == snapshot('gpt 4') (empty)

155 assert price.provider.id == snapshot('openai') (empty)

156 assert price.auto_update_timestamp is None (empty)

157 

158 

159def test_extra_source_sausage(): (empty)

160 with ExtraUpdatePrices() as update_prices: (empty)

161 update_prices.wait() (empty)

162 price = calc_price( (empty)

163 CustomUsage(sausages=3, input_tokens=1_000_000, output_tokens=1_000_000), 

164 model_ref='gpt-4o', 

165 provider_id='openai', 

166 ) 

167 assert price.input_price == snapshot(Decimal('1')) (empty)

168 assert price.output_price == snapshot(Decimal('2')) (empty)

169 assert price.total_price == snapshot(Decimal('12')) (empty)

170 assert price.model.name == snapshot('gpt-4o Custom') (empty)

171 assert price.provider.id == snapshot('openai') (empty)

172 assert price.auto_update_timestamp is None (empty)