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

1import aiohttp 

2from datetime import datetime, timedelta 

3from typing import List, Optional 

4from . import ProviderUsage 

5 

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

7 headers = {"Authorization": f"Bearer {api_key}"} 

8 

9 # Calculate start_time (1st day of current month) and end_time (tomorrow) 

10 now = datetime.now() 

11 start_str = now.strftime("%Y-%m-01") 

12 end_str = (now + timedelta(days=1)).strftime("%Y-%m-%d") 

13 

14 base_url_stripped = base_url.rstrip('/') 

15 if base_url_stripped.endswith('/v1'): 

16 completions_url = f"{base_url_stripped}/organization/usage/completions?start_time={start_str}&end_time={end_str}" 

17 costs_url = f"{base_url_stripped}/organization/usage/costs?start_time={start_str}&end_time={end_str}" 

18 else: 

19 completions_url = f"{base_url_stripped}/v1/organization/usage/completions?start_time={start_str}&end_time={end_str}" 

20 costs_url = f"{base_url_stripped}/v1/organization/usage/costs?start_time={start_str}&end_time={end_str}" 

21 

22 async with aiohttp.ClientSession() as session: 

23 async with session.get(completions_url, headers=headers) as resp: 

24 if resp.status in (401, 403): 

25 raise Exception("Requires Org Admin Key") 

26 if resp.status != 200: 

27 text = await resp.text() 

28 raise Exception(f"OpenAI API Error {resp.status}: {text}") 

29 completions_data = await resp.json() 

30 

31 cost_value = 0.0 

32 try: 

33 async with session.get(costs_url, headers=headers) as resp_cost: 

34 if resp_cost.status == 200: 

35 costs_data = await resp_cost.json() 

36 for item in costs_data.get("data", []): 

37 amt = item.get("amount", {}) 

38 if isinstance(amt, dict): 

39 val = amt.get("value") 

40 if val is not None: 

41 cost_value += float(val) 

42 except Exception: 

43 pass # Ignore cost fetch errors as optional fallback 

44 

45 total_input = 0 

46 total_output = 0 

47 for bucket in completions_data.get("data", []): 

48 total_input += bucket.get("input_tokens", 0) 

49 total_output += bucket.get("output_tokens", 0) 

50 

51 total_tokens = total_input + total_output 

52 

53 return [ 

54 ProviderUsage( 

55 provider="openai", 

56 label="Tokens", 

57 used=float(total_tokens), 

58 limit=None, 

59 remaining=None, 

60 percent=None, 

61 reset_at=None, 

62 unit="tokens" 

63 ), 

64 ProviderUsage( 

65 provider="openai", 

66 label="Cost", 

67 used=cost_value, 

68 limit=None, 

69 remaining=None, 

70 percent=None, 

71 reset_at=None, 

72 unit="$" 

73 ) 

74 ]