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

1import asyncio 

2from dataclasses import dataclass, field 

3from typing import Dict, List, Optional, Tuple 

4from kimi_code_usage.config import AppConfig 

5 

6 

7@dataclass 

8class ModelUsage: 

9 model: str 

10 spend: float = 0.0 

11 requests: float = 0.0 

12 prompt_tokens: float = 0.0 

13 completion_tokens: float = 0.0 

14 reasoning_tokens: float = 0.0 

15 

16 

17@dataclass 

18class DailyUsage: 

19 date: str 

20 models: List[ModelUsage] = field(default_factory=list) 

21 total: float = 0.0 

22 

23 

24@dataclass 

25class ActivityTotals: 

26 spend: float = 0.0 

27 requests: float = 0.0 

28 prompt_tokens: float = 0.0 

29 completion_tokens: float = 0.0 

30 reasoning_tokens: float = 0.0 

31 

32 

33@dataclass 

34class ProviderUsage: 

35 provider: str # "kimi" | "openai" | "anthropic" | "openrouter" 

36 label: str # e.g., "Weekly Usage", "5 Hours", "Cost" 

37 used: float 

38 limit: Optional[float] # None if unlimited 

39 remaining: Optional[float] 

40 percent: Optional[float] # 0-100, None if no limit 

41 reset_at: Optional[str] # ISO timestamp or formatted reset date 

42 unit: str # "%" | "tokens" | "$" | "credits" 

43 countdown: Optional[str] = None # Reset countdown, e.g. "5d 12h" 

44 text_value: Optional[str] = None 

45 activity_totals: Optional[ActivityTotals] = None 

46 top_models: Optional[List[ModelUsage]] = None 

47 daily_activity: Optional[List[DailyUsage]] = None 

48 

49# We will dynamically import fetchers to avoid circular dependencies 

50# and make it easier to load/mock them. 

51 

52async def dispatch_all(config: AppConfig) -> Tuple[Dict[str, List[ProviderUsage]], Dict[str, str]]: 

53 """ 

54 Fetch usage from all enabled providers in parallel. 

55 Returns a tuple of (results, errors) dicts. 

56 """ 

57 from .kimi import fetch_kimi_usage 

58 from .openai import fetch_openai_usage 

59 from .anthropic import fetch_anthropic_usage 

60 from .openrouter import fetch_openrouter_usage 

61 

62 fetchers = { 

63 "kimi": fetch_kimi_usage, 

64 "openai": fetch_openai_usage, 

65 "anthropic": fetch_anthropic_usage, 

66 "openrouter": fetch_openrouter_usage, 

67 } 

68 

69 results: Dict[str, List[ProviderUsage]] = {} 

70 errors: Dict[str, str] = {} 

71 

72 enabled = config.enabled_providers 

73 if not enabled: 

74 return results, errors 

75 

76 async def fetch_one(provider_name: str) -> Tuple[str, Optional[List[ProviderUsage]], Optional[str]]: 

77 fetch_func = fetchers.get(provider_name) 

78 if not fetch_func: 

79 return provider_name, None, f"Unknown provider {provider_name}" 

80 

81 p_conf = config.providers.get(provider_name) 

82 if not p_conf or not p_conf.api_key or not p_conf.base_url: 

83 return provider_name, None, "Not configured" 

84 

85 try: 

86 # 10s timeout 

87 res = await asyncio.wait_for( 

88 fetch_func(p_conf.api_key, p_conf.base_url, management_key=p_conf.management_key), 

89 timeout=10.0 

90 ) 

91 return provider_name, res, None 

92 except asyncio.TimeoutError: 

93 return provider_name, None, "Request timed out" 

94 except Exception as e: 

95 return provider_name, None, str(e) 

96 

97 tasks = [fetch_one(p) for p in enabled] 

98 responses = await asyncio.gather(*tasks) 

99 

100 for provider, res, err in responses: 

101 if err: 

102 errors[provider] = err 

103 elif res is not None: 

104 results[provider] = res 

105 

106 return results, errors