Coverage for src/kimi_code_usage/providers/openrouter.py: 95%
172 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-14 14:57 +0800
1import aiohttp
2import asyncio
3import os
4from typing import List, Optional
5from . import (
6 ProviderUsage,
7 ActivityTotals,
8 DailyUsage,
9 ModelUsage,
10)
13def _base_url_v1(base_url: str) -> str:
14 base = base_url.rstrip('/')
15 if base.endswith('/v1'):
16 return base
17 return base + '/v1'
20def _get_lang() -> bool:
21 return "zh" in os.getenv("LANG", "en").lower()
24def _build_activity(items: List[dict]) -> tuple[ActivityTotals, List[ModelUsage]]:
25 totals = ActivityTotals()
26 by_model: dict[str, ModelUsage] = {}
28 for item in items:
29 totals.spend += float(item.get("usage") or 0)
30 totals.requests += float(item.get("requests") or 0)
31 totals.prompt_tokens += float(item.get("prompt_tokens") or 0)
32 totals.completion_tokens += float(item.get("completion_tokens") or 0)
33 totals.reasoning_tokens += float(item.get("reasoning_tokens") or 0)
35 model_name = item.get("model", "unknown")
36 if model_name not in by_model:
37 by_model[model_name] = ModelUsage(model=model_name)
38 m = by_model[model_name]
39 m.spend += float(item.get("usage") or 0)
40 m.requests += float(item.get("requests") or 0)
41 m.prompt_tokens += float(item.get("prompt_tokens") or 0)
42 m.completion_tokens += float(item.get("completion_tokens") or 0)
43 m.reasoning_tokens += float(item.get("reasoning_tokens") or 0)
45 models = sorted(
46 by_model.values(),
47 key=lambda m: (m.spend, m.requests),
48 reverse=True,
49 )
50 return totals, models
53def _build_daily_activity(items: List[dict]) -> List[DailyUsage]:
54 by_date: dict[str, dict[str, ModelUsage]] = {}
56 for item in items:
57 date = item.get("date", "")
58 if not date:
59 continue
60 models = by_date.setdefault(date, {})
61 model_name = item.get("model", "unknown")
62 if model_name not in models:
63 models[model_name] = ModelUsage(model=model_name)
64 m = models[model_name]
65 m.spend += float(item.get("usage") or 0)
66 m.requests += float(item.get("requests") or 0)
67 m.prompt_tokens += float(item.get("prompt_tokens") or 0)
68 m.completion_tokens += float(item.get("completion_tokens") or 0)
69 m.reasoning_tokens += float(item.get("reasoning_tokens") or 0)
71 daily = []
72 for date in sorted(by_date.keys()):
73 models = sorted(
74 by_date[date].values(),
75 key=lambda m: m.spend,
76 reverse=True,
77 )
78 total = sum(m.spend for m in models)
79 daily.append(DailyUsage(date=date, models=models, total=total))
80 return daily
83def _model_short_name(model: str) -> str:
84 if "/" in model:
85 return model.rsplit("/", 1)[-1]
86 return model
89async def fetch_openrouter_usage(api_key: str, base_url: str, management_key: Optional[str] = None) -> List[ProviderUsage]:
90 base = _base_url_v1(base_url)
91 key_url = f"{base}/auth/key"
92 credits_url = f"{base}/credits"
93 activity_url = f"{base}/activity"
95 key_headers = {"Authorization": f"Bearer {api_key}"}
96 activity_headers = {"Authorization": f"Bearer {management_key or api_key}"}
98 async with aiohttp.ClientSession() as session:
99 async def fetch_key():
100 async with session.get(key_url, headers=key_headers) as resp:
101 if resp.status != 200:
102 text = await resp.text()
103 raise Exception(f"OpenRouter API Error {resp.status}: {text}")
104 return await resp.json()
106 async def fetch_credits():
107 async with session.get(credits_url, headers=key_headers) as resp:
108 if resp.status == 200:
109 return await resp.json()
110 return None
112 async def fetch_activity():
113 try:
114 async with session.get(activity_url, headers=activity_headers) as resp:
115 if resp.status == 200:
116 return await resp.json()
117 except Exception:
118 pass
119 return None
121 key_payload = await fetch_key()
122 if isinstance(key_payload, Exception):
123 raise key_payload
125 # Fetch activity for all keys; only management keys typically have access,
126 # so failures are silently ignored.
127 results = await asyncio.gather(fetch_credits(), fetch_activity(), return_exceptions=True)
129 credits_payload = results[0]
130 activity_payload = results[1]
132 if isinstance(credits_payload, Exception):
133 credits_payload = None
134 if isinstance(activity_payload, Exception):
135 activity_payload = None
137 key_data = key_payload.get("data", {})
138 is_management = bool(key_data.get("is_management_key"))
140 res: List[ProviderUsage] = []
141 is_zh = _get_lang()
143 key_usage = key_data.get("usage")
144 key_limit = key_data.get("limit")
145 key_used_val = float(key_usage) if key_usage is not None else 0.0
146 key_limit_val = float(key_limit) if key_limit is not None else None
148 credits_data = credits_payload.get("data", {}) if credits_payload else None
150 if credits_data:
151 total_credits = credits_data.get("total_credits")
152 total_usage = credits_data.get("total_usage")
153 if total_credits is not None and total_usage is not None:
154 t_credits = float(total_credits)
155 t_usage = float(total_usage)
156 remaining_val = t_credits - t_usage
157 percent_val = (t_usage / t_credits * 100) if t_credits > 0 else None
159 res.append(ProviderUsage(
160 provider="openrouter",
161 label="Credits",
162 used=t_usage,
163 limit=t_credits,
164 remaining=remaining_val,
165 percent=percent_val,
166 reset_at=None,
167 unit="$"
168 ))
170 k_remaining = (key_limit_val - key_used_val) if key_limit_val is not None else None
171 k_percent = (key_used_val / key_limit_val * 100) if (key_limit_val is not None and key_limit_val > 0) else None
173 res.append(ProviderUsage(
174 provider="openrouter",
175 label="Key Usage",
176 used=key_used_val,
177 limit=key_limit_val,
178 remaining=k_remaining,
179 percent=k_percent,
180 reset_at=None,
181 unit="$"
182 ))
184 if not res:
185 remaining_val = (key_limit_val - key_used_val) if key_limit_val is not None else None
186 percent_val = (key_used_val / key_limit_val * 100) if (key_limit_val is not None and key_limit_val > 0) else None
188 res.append(ProviderUsage(
189 provider="openrouter",
190 label="Credits",
191 used=key_used_val,
192 limit=key_limit_val,
193 remaining=remaining_val,
194 percent=percent_val,
195 reset_at=None,
196 unit="$"
197 ))
199 key_label = key_data.get("label")
200 if key_label:
201 if is_management:
202 display_label = "Management Key"
203 else:
204 display_label = "Key Name"
205 display_value = key_label.replace("...", "*********") if "..." in key_label else key_label
207 res.append(ProviderUsage(
208 provider="openrouter",
209 label=display_label,
210 used=0.0,
211 limit=None,
212 remaining=None,
213 percent=None,
214 reset_at=None,
215 unit="text",
216 text_value=display_value
217 ))
219 is_free_tier = key_data.get("is_free_tier")
220 if is_free_tier is not None:
221 res.append(ProviderUsage(
222 provider="openrouter",
223 label="Free Tier",
224 used=0.0,
225 limit=None,
226 remaining=None,
227 percent=None,
228 reset_at=None,
229 unit="text",
230 text_value="Yes" if is_free_tier else "No"
231 ))
233 limit_reset = key_data.get("limit_reset")
234 if limit_reset:
235 res.append(ProviderUsage(
236 provider="openrouter",
237 label="Limit Reset",
238 used=0.0,
239 limit=None,
240 remaining=None,
241 percent=None,
242 reset_at=None,
243 unit="text",
244 text_value=str(limit_reset)
245 ))
247 expires_at = key_data.get("expires_at")
248 if expires_at:
249 res.append(ProviderUsage(
250 provider="openrouter",
251 label="Expires At",
252 used=0.0,
253 limit=None,
254 remaining=None,
255 percent=None,
256 reset_at=None,
257 unit="text",
258 text_value=str(expires_at)
259 ))
261 is_provisioning = key_data.get("is_provisioning_key")
262 # NOTE: Is Provisioning is intentionally not displayed by default.
263 # It remains available in key_data for future configuration if needed.
264 _ = is_provisioning
266 rate_limit = key_data.get("rate_limit", {})
267 if isinstance(rate_limit, dict):
268 reqs = rate_limit.get("requests")
269 interval = rate_limit.get("interval")
270 if reqs is not None and interval is not None:
271 if reqs == -1:
272 if is_zh:
273 interval_zh = str(interval).replace("s", "秒").replace("m", "分钟").replace("h", "小时")
274 val_str = f"无限制/{interval_zh}"
275 else:
276 val_str = f"Unlimited/{interval}"
277 else:
278 if is_zh:
279 interval_zh = str(interval).replace("s", "秒").replace("m", "分钟").replace("h", "小时")
280 val_str = f"{reqs}次/{interval_zh}"
281 else:
282 val_str = f"{reqs} req/{interval}"
284 res.append(ProviderUsage(
285 provider="openrouter",
286 label="Rate Limit",
287 used=0.0,
288 limit=None,
289 remaining=None,
290 percent=None,
291 reset_at=None,
292 unit="text",
293 text_value=val_str
294 ))
296 u_daily = float(key_data.get("usage_daily") if key_data.get("usage_daily") is not None else 0.0)
297 u_weekly = float(key_data.get("usage_weekly") if key_data.get("usage_weekly") is not None else 0.0)
298 u_monthly = float(key_data.get("usage_monthly") if key_data.get("usage_monthly") is not None else 0.0)
300 label_str = "周期已用" if is_zh else "Usage"
301 if is_zh:
302 text_val = f"今日: ${u_daily:.4f} | 本周: ${u_weekly:.4f} | 本月: ${u_monthly:.4f}"
303 else:
304 text_val = f"Daily: ${u_daily:.4f} | Weekly: ${u_weekly:.4f} | Monthly: ${u_monthly:.4f}"
306 res.append(ProviderUsage(
307 provider="openrouter",
308 label=label_str,
309 used=0.0,
310 limit=None,
311 remaining=None,
312 percent=None,
313 reset_at=None,
314 unit="text",
315 text_value=text_val
316 ))
318 if activity_payload and isinstance(activity_payload, dict):
319 items = activity_payload.get("data", [])
320 if isinstance(items, list) and items:
321 totals, top_models = _build_activity(items)
322 daily_activity = _build_daily_activity(items)
324 res.append(ProviderUsage(
325 provider="openrouter",
326 label="Activity" if not is_zh else "活动",
327 used=0.0,
328 limit=None,
329 remaining=None,
330 percent=None,
331 reset_at=None,
332 unit="text",
333 text_value=None,
334 activity_totals=totals,
335 ))
337 if daily_activity:
338 res.append(ProviderUsage(
339 provider="openrouter",
340 label="Daily Spend" if not is_zh else "每日支出",
341 used=0.0,
342 limit=None,
343 remaining=None,
344 percent=None,
345 reset_at=None,
346 unit="text",
347 text_value=None,
348 daily_activity=daily_activity,
349 ))
351 if top_models:
352 res.append(ProviderUsage(
353 provider="openrouter",
354 label="Top Models" if not is_zh else "Top 模型",
355 used=0.0,
356 limit=None,
357 remaining=None,
358 percent=None,
359 reset_at=None,
360 unit="text",
361 text_value=None,
362 top_models=top_models,
363 ))
365 return res