Coverage for agentos/tests/test_token_counter.py: 0%

1"""Comprehensive tests for agentos/cost/token_counter.py.""" 

2 

3import pytest 

4 

5from agentos.cost.token_counter import ( 

6 PRICING_TABLE, 

7 CostEstimate, 

8 ModelFamily, 

9 TokenCount, 

10 TokenCounter, 

11) 

12 

13 

14class TestModelFamily: 

15 def test_all_families_exist(self): 

16 assert ModelFamily.GPT4.value == "gpt-4" 

17 assert ModelFamily.GPT4O.value == "gpt-4o" 

18 assert ModelFamily.GPT35.value == "gpt-3.5-turbo" 

19 assert ModelFamily.CLAUDE3.value == "claude-3" 

20 assert ModelFamily.CLAUDE35.value == "claude-3.5" 

21 assert ModelFamily.GEMINI.value == "gemini" 

22 assert ModelFamily.LLAMA.value == "llama" 

23 assert ModelFamily.MIXTRAL.value == "mixtral" 

24 assert ModelFamily.UNKNOWN.value == "unknown" 

25 

26 

27class TestTokenCount: 

28 def test_defaults(self): 

29 tc = TokenCount() 

30 assert tc.prompt_tokens == 0 

31 assert tc.completion_tokens == 0 

32 assert tc.total_tokens == 0 

33 assert tc.model == "" 

34 

35 def test_custom_values(self): 

36 tc = TokenCount(prompt_tokens=100, completion_tokens=50, total_tokens=150, model="gpt-4o") 

37 assert tc.prompt_tokens == 100 

38 assert tc.completion_tokens == 50 

39 assert tc.total_tokens == 150 

40 assert tc.model == "gpt-4o" 

41 

42 

43class TestCostEstimate: 

44 def test_defaults(self): 

45 ce = CostEstimate() 

46 assert ce.prompt_cost == 0.0 

47 assert ce.completion_cost == 0.0 

48 assert ce.total_cost == 0.0 

49 assert ce.currency == "USD" 

50 assert ce.token_count is None 

51 

52 def test_with_token_count(self): 

53 tc = TokenCount(prompt_tokens=1000, total_tokens=1000) 

54 ce = CostEstimate( 

55 prompt_cost=0.0025, 

56 completion_cost=0.0, 

57 total_cost=0.0025, 

58 token_count=tc, 

59 ) 

60 assert ce.prompt_cost == 0.0025 

61 assert ce.total_cost == 0.0025 

62 assert ce.token_count.prompt_tokens == 1000 

63 

64 

65class TestPricingTable: 

66 def test_major_models_have_pricing(self): 

67 assert "gpt-4o" in PRICING_TABLE 

68 assert "gpt-4o-mini" in PRICING_TABLE 

69 assert "gpt-4" in PRICING_TABLE 

70 assert "gpt-3.5-turbo" in PRICING_TABLE 

71 assert "claude-3.5-sonnet" in PRICING_TABLE 

72 assert "gemini-1.5-pro" in PRICING_TABLE 

73 assert "gemini-2.0-flash" in PRICING_TABLE 

74 

75 def test_pricing_is_tuple_of_two_floats(self): 

76 for model, pricing in PRICING_TABLE.items(): 

77 assert isinstance(pricing, tuple), f"{model} pricing not tuple" 

78 assert len(pricing) == 2, f"{model} pricing length != 2" 

79 assert isinstance(pricing[0], (int, float)), f"{model} input price not numeric" 

80 assert isinstance(pricing[1], (int, float)), f"{model} output price not numeric" 

81 

82 

83class TestTokenCounterInit: 

84 def test_initializes_without_error(self): 

85 counter = TokenCounter() 

86 assert isinstance(counter, TokenCounter) 

87 

88 def test_usage_log_starts_empty(self): 

89 counter = TokenCounter() 

90 assert counter._usage_log == [] 

91 

92 def test_reset_usage_clears_log(self): 

93 counter = TokenCounter() 

94 counter.count("hello") 

95 assert len(counter._usage_log) == 1 

96 counter.reset_usage() 

97 assert len(counter._usage_log) == 0 

98 

99 

100class TestTokenCounterCount: 

101 @pytest.fixture(autouse=True) 

102 def setup(self): 

103 self.counter = TokenCounter() 

104 

105 def test_count_short_text(self): 

106 tc = self.counter.count("hello", model="gpt-4o") 

107 assert tc.prompt_tokens > 0 

108 assert tc.total_tokens > 0 

109 assert tc.model == "gpt-4o" 

110 

111 def test_count_empty_text(self): 

112 tc = self.counter.count("", model="gpt-4o") 

113 assert tc.prompt_tokens >= 0 

114 assert tc.total_tokens >= 0 

115 

116 def test_count_long_text(self): 

117 long_text = "hello world " * 500 

118 tc = self.counter.count(long_text, model="gpt-4o") 

119 assert tc.total_tokens > 50 

120 

121 def test_count_different_models(self): 

122 for model in ["gpt-4o", "gpt-4", "gpt-3.5-turbo", "claude-3.5-sonnet", "gemini-1.5-pro"]: 

123 tc = self.counter.count("The quick brown fox jumps over the lazy dog.", model=model) 

124 assert tc.model == model 

125 assert tc.total_tokens > 0 

126 

127 def test_count_logs_usage(self): 

128 assert len(self.counter._usage_log) == 0 

129 self.counter.count("first") 

130 assert len(self.counter._usage_log) == 1 

131 self.counter.count("second") 

132 assert len(self.counter._usage_log) == 2 

133 

134 def test_count_unknown_model_falls_back(self): 

135 tc = self.counter.count("hello", model="nonexistent-model-xyz") 

136 assert tc.total_tokens > 0 

137 assert tc.model == "nonexistent-model-xyz" 

138 

139 

140class TestTokenCounterCountMessages: 

141 @pytest.fixture(autouse=True) 

142 def setup(self): 

143 self.counter = TokenCounter() 

144 

145 def test_single_message(self): 

146 msgs = [{"role": "user", "content": "Hello, how are you?"}] 

147 tc = self.counter.count_messages(msgs, model="gpt-4o") 

148 assert tc.prompt_tokens > 0 

149 assert tc.model == "gpt-4o" 

150 

151 def test_multiple_messages(self): 

152 msgs = [ 

153 {"role": "system", "content": "You are helpful."}, 

154 {"role": "user", "content": "What is Python?"}, 

155 {"role": "assistant", "content": "Python is a programming language."}, 

156 ] 

157 tc = self.counter.count_messages(msgs, model="gpt-4o") 

158 assert tc.prompt_tokens > 0 

159 

160 def test_empty_message_content(self): 

161 msgs = [{"role": "user", "content": ""}] 

162 tc = self.counter.count_messages(msgs) 

163 assert tc.prompt_tokens >= 0 

164 

165 def test_messages_log_to_usage(self): 

166 msgs = [{"role": "user", "content": "test"}] 

167 self.counter.count_messages(msgs) 

168 assert len(self.counter._usage_log) >= 1 

169 

170 

171class TestTokenCounterEstimateCost: 

172 @pytest.fixture(autouse=True) 

173 def setup(self): 

174 self.counter = TokenCounter() 

175 

176 def test_estimate_cost_gpt4o(self): 

177 tc = TokenCount(prompt_tokens=1000, total_tokens=1000, model="gpt-4o") 

178 cost = self.counter.estimate_cost(tc) 

179 assert cost.total_cost > 0 

180 assert cost.currency == "USD" 

181 assert cost.prompt_cost > 0 

182 

183 def test_estimate_cost_gpt4o_mini_is_cheaper(self): 

184 tc = TokenCount(prompt_tokens=1_000_000, total_tokens=1_000_000, model="gpt-4o-mini") 

185 cost = self.counter.estimate_cost(tc) 

186 # gpt-4o-mini input: $0.15/1M 

187 assert 0.10 < cost.prompt_cost < 0.20 

188 

189 def test_estimate_cost_with_model_override(self): 

190 tc = TokenCount(prompt_tokens=1_000_000, total_tokens=1_000_000, model="gpt-4o") 

191 cost = self.counter.estimate_cost(tc, model="gpt-4o-mini") 

192 assert 0.10 < cost.total_cost < 0.20 

193 

194 def test_estimate_cost_zero_tokens(self): 

195 tc = TokenCount() 

196 cost = self.counter.estimate_cost(tc) 

197 assert cost.total_cost == 0.0 

198 

199 def test_estimate_cost_unknown_model_default_pricing(self): 

200 tc = TokenCount(prompt_tokens=1_000_000, total_tokens=1_000_000, model="unknown-model") 

201 cost = self.counter.estimate_cost(tc) 

202 # Default: (1.00, 3.00) per 1M 

203 assert 0.50 < cost.total_cost < 5.00 

204 

205 def test_estimate_cost_with_completion_tokens(self): 

206 tc = TokenCount( 

207 prompt_tokens=500_000, 

208 completion_tokens=500_000, 

209 total_tokens=1_000_000, 

210 model="gpt-4o", 

211 ) 

212 cost = self.counter.estimate_cost(tc) 

213 assert cost.prompt_cost > 0 

214 assert cost.completion_cost > 0 

215 assert cost.completion_cost > cost.prompt_cost # output is more expensive 

216 

217 

218class TestTokenCounterClassifyModel: 

219 @pytest.fixture(autouse=True) 

220 def setup(self): 

221 self.counter = TokenCounter() 

222 

223 def test_classify_gpt4o(self): 

224 assert self.counter._classify_model("gpt-4o") == ModelFamily.GPT4O 

225 assert self.counter._classify_model("gpt-4o-mini") == ModelFamily.GPT4O 

226 

227 def test_classify_gpt4(self): 

228 assert self.counter._classify_model("gpt-4") == ModelFamily.GPT4 

229 assert self.counter._classify_model("gpt-4-turbo") == ModelFamily.GPT4 

230 

231 def test_classify_gpt35(self): 

232 assert self.counter._classify_model("gpt-3.5-turbo") == ModelFamily.GPT35 

233 

234 def test_classify_claude35(self): 

235 assert self.counter._classify_model("claude-3.5-sonnet") == ModelFamily.CLAUDE35 

236 

237 def test_classify_claude3(self): 

238 assert self.counter._classify_model("claude-3-opus") == ModelFamily.CLAUDE3 

239 assert self.counter._classify_model("claude-3-haiku") == ModelFamily.CLAUDE3 

240 

241 def test_classify_gemini(self): 

242 assert self.counter._classify_model("gemini-1.5-pro") == ModelFamily.GEMINI 

243 assert self.counter._classify_model("gemini-2.0-flash") == ModelFamily.GEMINI 

244 

245 def test_classify_llama(self): 

246 assert self.counter._classify_model("llama-3-70b") == ModelFamily.LLAMA 

247 

248 def test_classify_mixtral(self): 

249 assert self.counter._classify_model("mixtral-8x7b") == ModelFamily.MIXTRAL 

250 

251 def test_classify_unknown(self): 

252 assert self.counter._classify_model("random-model-123") == ModelFamily.UNKNOWN 

253 

254 

255class TestTokenCounterGetTotalUsage: 

256 def test_empty_log(self): 

257 counter = TokenCounter() 

258 total = counter.get_total_usage() 

259 assert total.prompt_tokens == 0 

260 assert total.completion_tokens == 0 

261 assert total.total_tokens == 0 

262 

263 def test_aggregates_all_entries(self): 

264 counter = TokenCounter() 

265 counter.count("first call") 

266 counter.count("second call") 

267 counter.count("third call") 

268 total = counter.get_total_usage() 

269 assert total.prompt_tokens > 0 

270 assert total.total_tokens > 0 

271 

272 def test_after_reset_returns_zero(self): 

273 counter = TokenCounter() 

274 counter.count("data") 

275 counter.reset_usage() 

276 total = counter.get_total_usage() 

277 assert total.total_tokens == 0 

278 

279 

280class TestTokenCounterGetTotalCost: 

281 def test_empty_log_zero_cost(self): 

282 counter = TokenCounter() 

283 cost = counter.get_total_cost() 

284 assert cost.total_cost == 0.0 

285 

286 def test_accumulates_cost(self): 

287 counter = TokenCounter() 

288 counter.count("hello " * 100, model="gpt-4o") 

289 counter.count("world " * 100, model="gpt-4o") 

290 cost = counter.get_total_cost() 

291 assert cost.total_cost > 0 

292 

293 

294class TestTokenCounterFormatCost: 

295 def test_tiny_cost(self): 

296 ce = CostEstimate(total_cost=0.000123) 

297 result = TokenCounter.format_cost(ce) 

298 assert "$" in result 

299 assert len(result.split(".")[1]) >= 6 

300 

301 def test_small_cost(self): 

302 ce = CostEstimate(total_cost=0.50) 

303 result = TokenCounter.format_cost(ce) 

304 assert "$" in result 

305 assert len(result.split(".")[1]) == 4 

306 

307 def test_large_cost(self): 

308 ce = CostEstimate(total_cost=42.0) 

309 result = TokenCounter.format_cost(ce) 

310 assert result == "$42.00" 

311 

312 

313class TestTokenCounterFormatTokens: 

314 def test_small_count(self): 

315 tc = TokenCount(total_tokens=500) 

316 result = TokenCounter.format_tokens(tc) 

317 assert result == "500" 

318 

319 def test_large_count(self): 

320 tc = TokenCount(total_tokens=2500) 

321 result = TokenCounter.format_tokens(tc) 

322 assert "2.5K" in result 

323 

324 

325class TestTokenCounterGetPricing: 

326 @pytest.fixture(autouse=True) 

327 def setup(self): 

328 self.counter = TokenCounter() 

329 

330 def test_exact_match(self): 

331 pricing = self.counter._get_pricing("gpt-4o") 

332 assert pricing == (2.50, 10.00) 

333 

334 def test_prefix_match(self): 

335 pricing = self.counter._get_pricing("gpt-4o-2024-08-06") 

336 assert pricing == (2.50, 10.00) 

337 

338 def test_unknown_model_default(self): 

339 pricing = self.counter._get_pricing("totally-unknown") 

340 assert pricing == (1.00, 3.00)