Coverage for src/kimi_code_usage/providers/openrouter.py: 100%

167 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-20 20:05 +0800

1import asyncio 

2import os 

3 

4import aiohttp 

5 

6from . import ( 

7 ActivityTotals, 

8 DailyUsage, 

9 ModelUsage, 

10 ProviderUsage, 

11) 

12 

13 

14def _base_url_v1(base_url: str) -> str: 

15 base = base_url.rstrip('/') 

16 if base.endswith('/v1'): 

17 return base 

18 return base + '/v1' 

19 

20 

21def _get_lang() -> bool: 

22 return "zh" in os.getenv("LANG", "en").lower() 

23 

24 

25def _build_activity(items: list[dict]) -> tuple[ActivityTotals, list[ModelUsage]]: 

26 totals = ActivityTotals() 

27 by_model: dict[str, ModelUsage] = {} 

28 

29 for item in items: 

30 totals.spend += float(item.get("usage") or 0) 

31 totals.requests += float(item.get("requests") or 0) 

32 totals.prompt_tokens += float(item.get("prompt_tokens") or 0) 

33 totals.completion_tokens += float(item.get("completion_tokens") or 0) 

34 totals.reasoning_tokens += float(item.get("reasoning_tokens") or 0) 

35 

36 model_name = item.get("model", "unknown") 

37 if model_name not in by_model: 

38 by_model[model_name] = ModelUsage(model=model_name) 

39 m = by_model[model_name] 

40 m.spend += float(item.get("usage") or 0) 

41 m.requests += float(item.get("requests") or 0) 

42 m.prompt_tokens += float(item.get("prompt_tokens") or 0) 

43 m.completion_tokens += float(item.get("completion_tokens") or 0) 

44 m.reasoning_tokens += float(item.get("reasoning_tokens") or 0) 

45 

46 models = sorted( 

47 by_model.values(), 

48 key=lambda m: (m.spend, m.requests), 

49 reverse=True, 

50 ) 

51 return totals, models 

52 

53 

54def _build_daily_activity(items: list[dict]) -> list[DailyUsage]: 

55 by_date: dict[str, dict[str, ModelUsage]] = {} 

56 

57 for item in items: 

58 date = item.get("date", "") 

59 if not date: 

60 continue 

61 models = by_date.setdefault(date, {}) 

62 model_name = item.get("model", "unknown") 

63 if model_name not in models: 

64 models[model_name] = ModelUsage(model=model_name) 

65 m = models[model_name] 

66 m.spend += float(item.get("usage") or 0) 

67 m.requests += float(item.get("requests") or 0) 

68 m.prompt_tokens += float(item.get("prompt_tokens") or 0) 

69 m.completion_tokens += float(item.get("completion_tokens") or 0) 

70 m.reasoning_tokens += float(item.get("reasoning_tokens") or 0) 

71 

72 daily = [] 

73 for date in sorted(by_date.keys()): 

74 models = sorted( 

75 by_date[date].values(), 

76 key=lambda m: m.spend, 

77 reverse=True, 

78 ) 

79 total = sum(m.spend for m in models) 

80 daily.append(DailyUsage(date=date, models=models, total=total)) 

81 return daily 

82 

83 

84def _model_short_name(model: str) -> str: 

85 if "/" in model: 

86 return model.rsplit("/", 1)[-1] 

87 return model 

88 

89 

90async def fetch_openrouter_usage(api_key: str, base_url: str, management_key: str | None = None) -> list[ProviderUsage]: 

91 base = _base_url_v1(base_url) 

92 key_url = f"{base}/auth/key" 

93 credits_url = f"{base}/credits" 

94 activity_url = f"{base}/activity" 

95 

96 key_headers = {"Authorization": f"Bearer {api_key}"} 

97 activity_headers = {"Authorization": f"Bearer {management_key or api_key}"} 

98 

99 async with aiohttp.ClientSession() as session: 

100 async def fetch_key(): 

101 async with session.get(key_url, headers=key_headers) as resp: 

102 if resp.status != 200: 

103 text = await resp.text() 

104 raise Exception(f"OpenRouter API Error {resp.status}: {text}") 

105 return await resp.json() 

106 

107 async def fetch_credits(): 

108 async with session.get(credits_url, headers=key_headers) as resp: 

109 if resp.status == 200: 

110 return await resp.json() 

111 return None 

112 

113 async def fetch_activity(): 

114 try: 

115 async with session.get(activity_url, headers=activity_headers) as resp: 

116 if resp.status == 200: 

117 return await resp.json() 

118 except Exception: 

119 pass 

120 return None 

121 

122 key_payload = await fetch_key() 

123 if isinstance(key_payload, Exception): # pragma: no cover 

124 raise key_payload 

125 

126 # Fetch activity for all keys; only management keys typically have access, 

127 # so failures are silently ignored. 

128 results = await asyncio.gather(fetch_credits(), fetch_activity(), return_exceptions=True) 

129 

130 credits_payload = results[0] 

131 activity_payload = results[1] 

132 

133 if isinstance(credits_payload, Exception): 

134 credits_payload = None 

135 if isinstance(activity_payload, Exception): # pragma: no cover 

136 activity_payload = None 

137 

138 key_data = key_payload.get("data", {}) 

139 is_management = bool(key_data.get("is_management_key")) 

140 

141 res: list[ProviderUsage] = [] 

142 is_zh = _get_lang() 

143 

144 key_usage = key_data.get("usage") 

145 key_limit = key_data.get("limit") 

146 key_used_val = float(key_usage) if key_usage is not None else 0.0 

147 key_limit_val = float(key_limit) if key_limit is not None else None 

148 

149 credits_data = credits_payload.get("data", {}) if credits_payload else None 

150 

151 if credits_data: 

152 total_credits = credits_data.get("total_credits") 

153 total_usage = credits_data.get("total_usage") 

154 if total_credits is not None and total_usage is not None: 

155 t_credits = float(total_credits) 

156 t_usage = float(total_usage) 

157 remaining_val = t_credits - t_usage 

158 percent_val = (t_usage / t_credits * 100) if t_credits > 0 else None 

159 

160 res.append(ProviderUsage( 

161 provider="openrouter", 

162 label="Credits", 

163 used=t_usage, 

164 limit=t_credits, 

165 remaining=remaining_val, 

166 percent=percent_val, 

167 reset_at=None, 

168 unit="$" 

169 )) 

170 

171 k_remaining = (key_limit_val - key_used_val) if key_limit_val is not None else None 

172 k_percent = (key_used_val / key_limit_val * 100) if (key_limit_val is not None and key_limit_val > 0) else None 

173 

174 res.append(ProviderUsage( 

175 provider="openrouter", 

176 label="Key Usage", 

177 used=key_used_val, 

178 limit=key_limit_val, 

179 remaining=k_remaining, 

180 percent=k_percent, 

181 reset_at=None, 

182 unit="$" 

183 )) 

184 

185 if not res: 

186 remaining_val = (key_limit_val - key_used_val) if key_limit_val is not None else None 

187 percent_val = (key_used_val / key_limit_val * 100) if (key_limit_val is not None and key_limit_val > 0) else None 

188 

189 res.append(ProviderUsage( 

190 provider="openrouter", 

191 label="Credits", 

192 used=key_used_val, 

193 limit=key_limit_val, 

194 remaining=remaining_val, 

195 percent=percent_val, 

196 reset_at=None, 

197 unit="$" 

198 )) 

199 

200 key_label = key_data.get("label") 

201 if key_label: 

202 if is_management: 

203 display_label = "Management Key" 

204 else: 

205 display_label = "Key Name" 

206 display_value = key_label.replace("...", "*********") if "..." in key_label else key_label 

207 

208 res.append(ProviderUsage( 

209 provider="openrouter", 

210 label=display_label, 

211 used=0.0, 

212 limit=None, 

213 remaining=None, 

214 percent=None, 

215 reset_at=None, 

216 unit="text", 

217 text_value=display_value 

218 )) 

219 

220 is_free_tier = key_data.get("is_free_tier") 

221 if is_free_tier is not None: 

222 res.append(ProviderUsage( 

223 provider="openrouter", 

224 label="Free Tier", 

225 used=0.0, 

226 limit=None, 

227 remaining=None, 

228 percent=None, 

229 reset_at=None, 

230 unit="text", 

231 text_value="Yes" if is_free_tier else "No" 

232 )) 

233 

234 limit_reset = key_data.get("limit_reset") 

235 if limit_reset: 

236 res.append(ProviderUsage( 

237 provider="openrouter", 

238 label="Limit Reset", 

239 used=0.0, 

240 limit=None, 

241 remaining=None, 

242 percent=None, 

243 reset_at=None, 

244 unit="text", 

245 text_value=str(limit_reset) 

246 )) 

247 

248 expires_at = key_data.get("expires_at") 

249 if expires_at: 

250 res.append(ProviderUsage( 

251 provider="openrouter", 

252 label="Expires At", 

253 used=0.0, 

254 limit=None, 

255 remaining=None, 

256 percent=None, 

257 reset_at=None, 

258 unit="text", 

259 text_value=str(expires_at) 

260 )) 

261 

262 is_provisioning = key_data.get("is_provisioning_key") 

263 # NOTE: Is Provisioning is intentionally not displayed by default. 

264 # It remains available in key_data for future configuration if needed. 

265 _ = is_provisioning 

266 

267 rate_limit = key_data.get("rate_limit", {}) 

268 if isinstance(rate_limit, dict): 

269 reqs = rate_limit.get("requests") 

270 interval = rate_limit.get("interval") 

271 if reqs is not None and interval is not None: 

272 if reqs == -1: 

273 if is_zh: 

274 interval_zh = str(interval).replace("s", "秒").replace("m", "分钟").replace("h", "小时") 

275 val_str = f"无限制/{interval_zh}" 

276 else: 

277 val_str = f"Unlimited/{interval}" 

278 else: 

279 if is_zh: 

280 interval_zh = str(interval).replace("s", "秒").replace("m", "分钟").replace("h", "小时") 

281 val_str = f"{reqs}次/{interval_zh}" 

282 else: 

283 val_str = f"{reqs} req/{interval}" 

284 

285 res.append(ProviderUsage( 

286 provider="openrouter", 

287 label="Rate Limit", 

288 used=0.0, 

289 limit=None, 

290 remaining=None, 

291 percent=None, 

292 reset_at=None, 

293 unit="text", 

294 text_value=val_str 

295 )) 

296 

297 u_daily = float(key_data.get("usage_daily") if key_data.get("usage_daily") is not None else 0.0) 

298 u_weekly = float(key_data.get("usage_weekly") if key_data.get("usage_weekly") is not None else 0.0) 

299 u_monthly = float(key_data.get("usage_monthly") if key_data.get("usage_monthly") is not None else 0.0) 

300 

301 label_str = "周期已用" if is_zh else "Usage" 

302 if is_zh: 

303 text_val = f"今日: ${u_daily:.4f} | 本周: ${u_weekly:.4f} | 本月: ${u_monthly:.4f}" 

304 else: 

305 text_val = f"Daily: ${u_daily:.4f} | Weekly: ${u_weekly:.4f} | Monthly: ${u_monthly:.4f}" 

306 

307 res.append(ProviderUsage( 

308 provider="openrouter", 

309 label=label_str, 

310 used=0.0, 

311 limit=None, 

312 remaining=None, 

313 percent=None, 

314 reset_at=None, 

315 unit="text", 

316 text_value=text_val 

317 )) 

318 

319 if activity_payload and isinstance(activity_payload, dict): 

320 items = activity_payload.get("data", []) 

321 if isinstance(items, list) and items: 

322 totals, top_models = _build_activity(items) 

323 daily_activity = _build_daily_activity(items) 

324 

325 res.append(ProviderUsage( 

326 provider="openrouter", 

327 label="Activity" if not is_zh else "活动", 

328 used=0.0, 

329 limit=None, 

330 remaining=None, 

331 percent=None, 

332 reset_at=None, 

333 unit="text", 

334 text_value=None, 

335 activity_totals=totals, 

336 )) 

337 

338 if daily_activity: 

339 res.append(ProviderUsage( 

340 provider="openrouter", 

341 label="Daily Spend" if not is_zh else "每日支出", 

342 used=0.0, 

343 limit=None, 

344 remaining=None, 

345 percent=None, 

346 reset_at=None, 

347 unit="text", 

348 text_value=None, 

349 daily_activity=daily_activity, 

350 )) 

351 

352 if top_models: 

353 res.append(ProviderUsage( 

354 provider="openrouter", 

355 label="Top Models" if not is_zh else "Top 模型", 

356 used=0.0, 

357 limit=None, 

358 remaining=None, 

359 percent=None, 

360 reset_at=None, 

361 unit="text", 

362 text_value=None, 

363 top_models=top_models, 

364 )) 

365 

366 return res