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

1from datetime import datetime, timedelta 

2 

3import aiohttp 

4 

5from . import ProviderUsage 

6 

7 

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

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

10 

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

12 now = datetime.now() 

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

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

15 

16 base_url_stripped = base_url.rstrip('/') 

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

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

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

20 else: 

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

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

23 

24 async with aiohttp.ClientSession() as session: 

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

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

27 raise Exception("Requires Org Admin Key") 

28 if resp.status != 200: 

29 text = await resp.text() 

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

31 completions_data = await resp.json() 

32 

33 cost_value = 0.0 

34 try: 

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

36 if resp_cost.status == 200: 

37 costs_data = await resp_cost.json() 

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

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

40 if isinstance(amt, dict): 

41 val = amt.get("value") 

42 if val is not None: 

43 cost_value += float(val) 

44 except Exception: 

45 pass # Ignore cost fetch errors as optional fallback 

46 

47 total_input = 0 

48 total_output = 0 

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

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

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

52 

53 total_tokens = total_input + total_output 

54 

55 return [ 

56 ProviderUsage( 

57 provider="openai", 

58 label="Tokens", 

59 used=float(total_tokens), 

60 limit=None, 

61 remaining=None, 

62 percent=None, 

63 reset_at=None, 

64 unit="tokens" 

65 ), 

66 ProviderUsage( 

67 provider="openai", 

68 label="Cost", 

69 used=cost_value, 

70 limit=None, 

71 remaining=None, 

72 percent=None, 

73 reset_at=None, 

74 unit="$" 

75 ) 

76 ]