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

1import aiohttp 

2from datetime import datetime, timedelta 

3from typing import Any, List, Mapping, Optional, Sequence, Tuple, cast 

4from . import ProviderUsage 

5 

6def _to_int(v: Any) -> Optional[int]: 

7 try: 

8 return int(v) 

9 except (TypeError, ValueError): 

10 return None 

11 

12def _get_reset_info(data: Mapping[str, Any]) -> Optional[Tuple[str, str]]: 

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

14 if reset_at: 

15 try: 

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

17 dt = datetime.fromtimestamp(reset_at) 

18 else: 

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

20 

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

22 diff = dt - now 

23 if diff.total_seconds() <= 0: 

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

25 days = diff.days 

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

27 minutes, _ = divmod(rem, 60) 

28 parts = [] 

29 if days > 0: 

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

31 if hours > 0: 

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

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

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

35 except Exception: 

36 pass 

37 

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

39 if reset_in is not None: 

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

41 hours, rem = divmod(reset_in, 3600) 

42 minutes, _ = divmod(rem, 60) 

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

44 return None 

45 

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

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

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

49 if duration is not None: 

50 if "MINUTE" in time_unit: 

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

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

53 return f"{duration}m Limit" 

54 if "HOUR" in time_unit: 

55 return f"{duration}h Limit" 

56 if "DAY" in time_unit: 

57 return f"{duration}d Limit" 

58 if "MONTH" in time_unit: 

59 return f"{duration}mo Limit" 

60 return f"{duration}s Limit" 

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

62 

63class KimiRow: 

64 def __init__(self, label: str, used: int, limit: int, reset_at: Optional[str] = None, countdown: Optional[str] = None): 

65 self.label = label 

66 self.used = used 

67 self.limit = limit 

68 self.reset_at = reset_at 

69 self.countdown = countdown 

70 

71def _to_usage_row(data: Mapping[str, Any], default_label: str) -> Optional[KimiRow]: 

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

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

74 if used is None: 

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

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

77 used = limit - remaining 

78 if used is None and limit is None: 

79 return None 

80 reset_info = _get_reset_info(data) 

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

82 return KimiRow( 

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

84 used=used or 0, 

85 limit=limit or 0, 

86 reset_at=reset_at, 

87 countdown=countdown, 

88 ) 

89 

90def _parse_usage_payload(payload: Mapping[str, Any]) -> Tuple[Optional[KimiRow], List[KimiRow]]: 

91 summary = None 

92 limits = [] 

93 

94 data_list = payload.get("data") 

95 if isinstance(data_list, Sequence): 

96 for item in data_list: 

97 if not isinstance(item, Mapping): 

98 continue 

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

100 row = _to_usage_row(item, default_label=label) 

101 if row: 

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

103 summary = row 

104 else: 

105 limits.append(row) 

106 else: 

107 usage = payload.get("usage") 

108 if isinstance(usage, Mapping): 

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

110 raw_limits = payload.get("limits") 

111 if isinstance(raw_limits, Sequence): 

112 for idx, item in enumerate(raw_limits): 

113 if not isinstance(item, Mapping): 

114 continue 

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

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

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

118 if row: 

119 limits.append(row) 

120 

121 return summary, limits 

122 

123async def fetch_kimi_usage(api_key: str, base_url: str, management_key: Optional[str] = None) -> List[ProviderUsage]: 

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

125 async with aiohttp.ClientSession() as session: 

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

127 if resp.status != 200: 

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

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

130 if f_resp.status != 200: 

131 text = await f_resp.text() 

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

133 payload = await f_resp.json() 

134 else: 

135 payload = await resp.json() 

136 

137 summary, limits = _parse_usage_payload(payload) 

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

139 

140 res = [] 

141 for r in rows: 

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

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

144 res.append(ProviderUsage( 

145 provider="kimi", 

146 label=r.label, 

147 used=float(r.used), 

148 limit=float(r.limit), 

149 remaining=rem, 

150 percent=pct, 

151 reset_at=r.reset_at, 

152 countdown=r.countdown, 

153 unit="%" 

154 )) 

155 return res