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

1from collections.abc import Mapping, Sequence 

2from datetime import datetime, timedelta 

3from typing import Any, cast 

4 

5import aiohttp 

6 

7from . import ProviderUsage 

8 

9 

10def _to_int(v: Any) -> int | None: 

11 try: 

12 return int(v) 

13 except (TypeError, ValueError): 

14 return None 

15 

16def _get_reset_info(data: Mapping[str, Any]) -> tuple[str, str] | None: 

17 reset_at = data.get("resetTime") or data.get("reset_at") or data.get("reset_time") 

18 if reset_at: 

19 try: 

20 if isinstance(reset_at, (int, float)): 

21 dt = datetime.fromtimestamp(reset_at) 

22 else: 

23 dt = datetime.fromisoformat(reset_at.replace("Z", "+00:00")).astimezone() 

24 

25 now = datetime.now(dt.tzinfo) if dt.tzinfo else datetime.now() 

26 diff = dt - now 

27 if diff.total_seconds() <= 0: 

28 return dt.strftime("%m-%d %H:%M"), "0m" 

29 days = diff.days 

30 hours, rem = divmod(diff.seconds, 3600) 

31 minutes, _ = divmod(rem, 60) 

32 parts = [] 

33 if days > 0: 

34 parts.append(f"{days}d") 

35 if hours > 0: 

36 parts.append(f"{hours}h") 

37 parts.append(f"{minutes}m") 

38 return dt.strftime("%m-%d %H:%M"), " ".join(parts) 

39 except Exception: 

40 pass 

41 

42 reset_in = _to_int(data.get("reset_in")) 

43 if reset_in is not None: 

44 dt = datetime.now() + timedelta(seconds=reset_in) 

45 hours, rem = divmod(reset_in, 3600) 

46 minutes, _ = divmod(rem, 60) 

47 return dt.strftime("%m-%d %H:%M"), f"{hours}h {minutes}m" 

48 return None 

49 

50def _limit_label(window: Mapping[str, Any], idx: int) -> str: 

51 duration = _to_int(window.get("duration")) 

52 time_unit = str(window.get("timeUnit") or window.get("time_unit") or "").upper() 

53 if duration is not None: 

54 if "MINUTE" in time_unit: 

55 if duration >= 60 and duration % 60 == 0: 

56 return f"{duration // 60}h Limit" 

57 return f"{duration}m Limit" 

58 if "HOUR" in time_unit: 

59 return f"{duration}h Limit" 

60 if "DAY" in time_unit: 

61 return f"{duration}d Limit" 

62 if "MONTH" in time_unit: 

63 return f"{duration}mo Limit" 

64 return f"{duration}s Limit" 

65 return f"Limit #{idx + 1}" 

66 

67class KimiRow: 

68 def __init__(self, label: str, used: int, limit: int, reset_at: str | None = None, countdown: str | None = None): 

69 self.label = label 

70 self.used = used 

71 self.limit = limit 

72 self.reset_at = reset_at 

73 self.countdown = countdown 

74 

75def _to_usage_row(data: Mapping[str, Any], default_label: str) -> KimiRow | None: 

76 limit = _to_int(data.get("limit") or data.get("limit_amount")) 

77 used = _to_int(data.get("used") or data.get("used_amount")) 

78 if used is None: 

79 remaining = _to_int(data.get("remaining")) 

80 if remaining is not None and limit is not None: 

81 used = limit - remaining 

82 if used is None and limit is None: 

83 return None 

84 reset_info = _get_reset_info(data) 

85 reset_at, countdown = reset_info if reset_info else (None, None) 

86 return KimiRow( 

87 label=str(data.get("name") or data.get("title") or data.get("model_name") or default_label), 

88 used=used or 0, 

89 limit=limit or 0, 

90 reset_at=reset_at, 

91 countdown=countdown, 

92 ) 

93 

94def _parse_usage_payload(payload: Mapping[str, Any]) -> tuple[KimiRow | None, list[KimiRow]]: 

95 summary = None 

96 limits = [] 

97 

98 data_list = payload.get("data") 

99 if isinstance(data_list, Sequence): 

100 for item in data_list: 

101 if not isinstance(item, Mapping): 

102 continue 

103 label = "Weekly Usage" if item.get("model_name") == "all" else "Limit" 

104 row = _to_usage_row(item, default_label=label) 

105 if row: 

106 if item.get("model_name") == "all": 

107 summary = row 

108 else: 

109 limits.append(row) 

110 else: 

111 usage = payload.get("usage") 

112 if isinstance(usage, Mapping): 

113 summary = _to_usage_row(cast(Mapping, usage), default_label="Weekly Usage") 

114 raw_limits = payload.get("limits") 

115 if isinstance(raw_limits, Sequence): 

116 for idx, item in enumerate(raw_limits): 

117 if not isinstance(item, Mapping): 

118 continue 

119 detail = item.get("detail") if isinstance(item.get("detail"), Mapping) else item 

120 window = item.get("window") if isinstance(item.get("window"), Mapping) else {} 

121 row = _to_usage_row(detail, default_label=_limit_label(window, idx)) 

122 if row: 

123 limits.append(row) 

124 

125 return summary, limits 

126 

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

128 url = base_url.rstrip("/") + "/usages" 

129 async with aiohttp.ClientSession() as session: 

130 async with session.get(url, headers={"Authorization": f"Bearer {api_key}"}) as resp: 

131 if resp.status != 200: 

132 fallback_url = base_url.rstrip("/") + "/usage" 

133 async with session.get(fallback_url, headers={"Authorization": f"Bearer {api_key}"}) as f_resp: 

134 if f_resp.status != 200: 

135 text = await f_resp.text() 

136 raise Exception(f"API Error {f_resp.status}: {text}") 

137 payload = await f_resp.json() 

138 else: 

139 payload = await resp.json() 

140 

141 summary, limits = _parse_usage_payload(payload) 

142 rows = ([summary] if summary else []) + limits 

143 

144 res = [] 

145 for r in rows: 

146 pct = (r.used / r.limit * 100) if r.limit > 0 else 0.0 

147 rem = float(r.limit - r.used) 

148 res.append(ProviderUsage( 

149 provider="kimi", 

150 label=r.label, 

151 used=float(r.used), 

152 limit=float(r.limit), 

153 remaining=rem, 

154 percent=pct, 

155 reset_at=r.reset_at, 

156 countdown=r.countdown, 

157 unit="%" 

158 )) 

159 return res