Coverage for src/kimi_code_usage/config.py: 97%

122 statements  

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

1import json 

2import os 

3from dataclasses import dataclass, field 

4from pathlib import Path 

5from typing import Dict, List, Optional 

6 

7DEFAULT_CONFIG_PATH = Path.home() / ".kimi-usage" / "config.json" 

8 

9@dataclass 

10class ProviderConfig: 

11 api_key: Optional[str] = None 

12 base_url: Optional[str] = None 

13 enabled: bool = True 

14 management_key: Optional[str] = None 

15 

16@dataclass 

17class AppConfig: 

18 providers: Dict[str, ProviderConfig] = field(default_factory=dict) 

19 refresh_interval_minutes: int = 5 

20 output_mode: str = "rich" 

21 provider_order: List[str] = field(default_factory=list) 

22 theme: str = "blue-dark" 

23 language: Optional[str] = None 

24 visible_providers: Optional[List[str]] = None 

25 or_metric: str = "requests" 

26 days_window: int = 30 

27 

28 @property 

29 def enabled_providers(self) -> List[str]: 

30 return [name for name, p_conf in self.providers.items() if p_conf.api_key and p_conf.enabled] 

31 

32class ConfigResolver: 

33 def __init__(self, config_path: Optional[str] = None): 

34 self.config_path = Path(config_path) if config_path else DEFAULT_CONFIG_PATH 

35 self.config = AppConfig() 

36 

37 def resolve(self) -> AppConfig: 

38 # 1. Load from JSON if exists 

39 json_data = {} 

40 if self.config_path.exists(): 

41 try: 

42 with open(self.config_path, "r", encoding="utf-8") as f: 

43 json_data = json.load(f) 

44 except Exception: 

45 pass # Ignore invalid JSON, fall back to default/env 

46 

47 # Parse JSON config 

48 providers_data = json_data.get("providers", {}) 

49 general_data = json_data.get("general", {}) 

50 

51 self.config.refresh_interval_minutes = general_data.get("refreshIntervalMinutes", 5) 

52 self.config.output_mode = general_data.get("outputMode", "rich") 

53 

54 # Get theme from JSON, else ENV, default blue-dark 

55 theme = general_data.get("theme") 

56 if not theme: 

57 theme = os.getenv("KIMI_USAGE_THEME", "blue-dark") 

58 self.config.theme = theme 

59 

60 # Get language and visible providers from JSON 

61 self.config.language = general_data.get("language") 

62 self.config.visible_providers = general_data.get("visibleProviders") 

63 

64 # Get OpenRouter metric preference, default requests 

65 or_metric = general_data.get("orMetric", "requests") 

66 if or_metric not in ("spend", "requests", "tokens"): 

67 or_metric = "requests" 

68 self.config.or_metric = or_metric 

69 

70 # Get days window for daily chart, default 30 

71 days_window = general_data.get("daysWindow", 30) 

72 try: 

73 days_window = int(days_window) 

74 except (TypeError, ValueError): 

75 days_window = 30 

76 if days_window not in (7, 14, 30, 60, 90): 

77 days_window = 30 

78 self.config.days_window = days_window 

79 

80 # 2. Setup providers 

81 all_providers = ["kimi", "openai", "anthropic", "openrouter"] 

82 

83 # Envs map 

84 env_keys = { 

85 "kimi": ("KIMI_API_KEY", "KIMI_CODING_API_KEY"), 

86 "openai": ("OPENAI_API_KEY",), 

87 "anthropic": ("ANTHROPIC_API_KEY",), 

88 "openrouter": ("OPENROUTER_API_KEY", "OPENROUTER_ADMIN_KEY"), 

89 } 

90 

91 env_urls = { 

92 "kimi": "KIMI_BASE_URL", 

93 "openai": "OPENAI_BASE_URL", 

94 "anthropic": "ANTHROPIC_BASE_URL", 

95 "openrouter": "OPENROUTER_BASE_URL", 

96 } 

97 

98 default_urls = { 

99 "kimi": "https://api.kimi.com/coding/v1", 

100 "openai": "https://api.openai.com", 

101 "anthropic": "https://api.anthropic.com", 

102 "openrouter": "https://openrouter.ai/api", 

103 } 

104 

105 for p in all_providers: 

106 p_data = providers_data.get(p, {}) 

107 # Get API Key from JSON, else ENV 

108 api_key = p_data.get("apiKey") 

109 if not api_key: 

110 for env_var in env_keys[p]: 

111 val = os.getenv(env_var) 

112 if val: 

113 api_key = val 

114 break 

115 

116 # Get Base URL from JSON, else ENV, else Default 

117 base_url = p_data.get("baseUrl") 

118 if not base_url: 

119 base_url = os.getenv(env_urls[p]) 

120 if not base_url: 

121 base_url = default_urls[p] 

122 

123 # Get optional management key for OpenRouter (activity endpoint) 

124 management_key = None 

125 if p == "openrouter": 

126 management_key = ( 

127 p_data.get("managementApiKey") 

128 or os.getenv("OPENROUTER_MANAGEMENT_KEY") 

129 or os.getenv("OPENROUTER_ADMIN_KEY") 

130 ) 

131 

132 # Get Enabled status from JSON, else ENV, default True 

133 enabled = p_data.get("enabled") 

134 if enabled is None: 

135 env_enabled = os.getenv(f"{p.upper()}_ENABLED") 

136 if env_enabled is not None: 

137 enabled = env_enabled.lower() not in ("false", "0", "no") 

138 else: 

139 enabled = True 

140 

141 self.config.providers[p] = ProviderConfig( 

142 api_key=api_key if api_key else None, 

143 base_url=base_url, 

144 enabled=enabled, 

145 management_key=management_key if management_key else None, 

146 ) 

147 

148 # 3. Determine provider ordering 

149 default_order = ["anthropic", "openai", "openrouter", "kimi"] 

150 json_providers = list(providers_data.keys()) 

151 

152 final_order = [] 

153 for p in json_providers: 

154 if p in all_providers and p not in final_order: 

155 final_order.append(p) 

156 for p in default_order: 

157 if p not in final_order: 

158 final_order.append(p) 

159 

160 self.config.provider_order = final_order 

161 

162 return self.config 

163 

164 

165def save_theme( 

166 theme_name: str, 

167 language: Optional[str] = None, 

168 visible_providers: Optional[List[str]] = None, 

169 or_metric: Optional[str] = None, 

170 days_window: Optional[int] = None, 

171 config_path: Optional[Path] = None 

172) -> None: 

173 """Persist settings into ``general`` in the JSON config file. 

174 

175 All other settings in the file are preserved. The config directory is 

176 created automatically if it does not exist yet. 

177 """ 

178 path = config_path or DEFAULT_CONFIG_PATH 

179 path.parent.mkdir(parents=True, exist_ok=True) 

180 

181 data: dict = {} 

182 if path.exists(): 

183 try: 

184 with open(path, "r", encoding="utf-8") as fh: 

185 data = json.load(fh) 

186 except Exception: 

187 data = {} 

188 

189 general = data.get("general", {}) 

190 general["theme"] = theme_name 

191 if language is not None: 

192 general["language"] = language 

193 if visible_providers is not None: 

194 general["visibleProviders"] = visible_providers 

195 if or_metric is not None: 

196 general["orMetric"] = or_metric 

197 if days_window is not None: 

198 general["daysWindow"] = days_window 

199 data["general"] = general 

200 

201 with open(path, "w", encoding="utf-8") as fh: 

202 json.dump(data, fh, indent=2, ensure_ascii=False)