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
4import aiohttp
6from . import (
7 ActivityTotals,
8 DailyUsage,
9 ModelUsage,
10 ProviderUsage,
11)
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'
21def _get_lang() -> bool:
22 return "zh" in os.getenv("LANG", "en").lower()
25def _build_activity(items: list[dict]) -> tuple[ActivityTotals, list[ModelUsage]]:
26 totals = ActivityTotals()
27 by_model: dict[str, ModelUsage] = {}
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)
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)
46 models = sorted(
47 by_model.values(),
48 key=lambda m: (m.spend, m.requests),
49 reverse=True,
50 )
51 return totals, models
54def _build_daily_activity(items: list[dict]) -> list[DailyUsage]:
55 by_date: dict[str, dict[str, ModelUsage]] = {}
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)
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
84def _model_short_name(model: str) -> str:
85 if "/" in model:
86 return model.rsplit("/", 1)[-1]
87 return model
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"
96 key_headers = {"Authorization": f"Bearer {api_key}"}
97 activity_headers = {"Authorization": f"Bearer {management_key or api_key}"}
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()
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
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
122 key_payload = await fetch_key()
123 if isinstance(key_payload, Exception): # pragma: no cover
124 raise key_payload
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)
130 credits_payload = results[0]
131 activity_payload = results[1]
133 if isinstance(credits_payload, Exception):
134 credits_payload = None
135 if isinstance(activity_payload, Exception): # pragma: no cover
136 activity_payload = None
138 key_data = key_payload.get("data", {})
139 is_management = bool(key_data.get("is_management_key"))
141 res: list[ProviderUsage] = []
142 is_zh = _get_lang()
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
149 credits_data = credits_payload.get("data", {}) if credits_payload else None
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
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 ))
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
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 ))
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
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 ))
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
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 ))
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 ))
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 ))
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 ))
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
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}"
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 ))
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)
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}"
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 ))
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)
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 ))
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 ))
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 ))
366 return res