Coverage for src/lilbee/providers/key_check.py: 100%

1"""Hosted provider API-key status: one authenticated call per key, cached.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6import os 

7import threading 

8from collections.abc import Callable 

9from concurrent.futures import ThreadPoolExecutor 

10from dataclasses import dataclass 

11 

12import httpx 

13from cachetools import TTLCache, cached 

14 

15from lilbee.catalog.types import KeyStatus 

16from lilbee.providers.sdk_backend import PROVIDER_API_KEY_ENV, get_provider_api_key 

17 

18log = logging.getLogger(__name__) 

19 

20KEY_CHECK_TTL_S = 300 

21KEY_CHECK_TIMEOUT_S = 5.0 

22_KEYS_CACHED_PER_PROVIDER = 4 

23_AUTH_REJECTED_STATUS = frozenset({httpx.codes.UNAUTHORIZED, httpx.codes.FORBIDDEN}) 

24_ANTHROPIC_API_VERSION = "2023-06-01" 

25_GEMINI_INVALID_KEY_REASON = "API_KEY_INVALID" 

26 

27 

28def _bearer(api_key: str) -> dict[str, str]: 

29 return {"Authorization": f"Bearer {api_key}"} 

30 

31 

32def _anthropic_headers(api_key: str) -> dict[str, str]: 

33 return {"x-api-key": api_key, "anthropic-version": _ANTHROPIC_API_VERSION} 

34 

35 

36def _gemini_headers(api_key: str) -> dict[str, str]: 

37 return {"x-goog-api-key": api_key} 

38 

39 

40def _auth_rejected(resp: httpx.Response) -> bool: 

41 return resp.status_code in _AUTH_REJECTED_STATUS 

42 

43 

44def _gemini_rejected(resp: httpx.Response) -> bool: 

45 """Gemini answers a bad key with 400 and the structured reason ``API_KEY_INVALID``.""" 

46 if _auth_rejected(resp): 

47 return True 

48 if resp.status_code != httpx.codes.BAD_REQUEST: 

49 return False 

50 try: 

51 details = resp.json()["error"]["details"] 

52 except (ValueError, KeyError, TypeError): 

53 return False 

54 # Untyped JSON error body: details must be a list, and only dict entries carry a reason. 

55 if not isinstance(details, list): 

56 return False 

57 return any( 

58 isinstance(item, dict) and item.get("reason") == _GEMINI_INVALID_KEY_REASON 

59 for item in details 

60 ) 

61 

62 

63@dataclass(frozen=True) 

64class _KeyProbe: 

65 """An authenticated GET that fails only when the key is bad.""" 

66 

67 url: str 

68 headers: Callable[[str], dict[str, str]] 

69 rejected: Callable[[httpx.Response], bool] = _auth_rejected 

70 

71 

72# Not litellm: get_models raises untyped errors; check_valid_key and get_valid_models hide them. 

73# OpenRouter's model list is public, so its probe reads the key's own record instead. 

74_PROBES: dict[str, _KeyProbe] = { 

75 "openrouter": _KeyProbe("https://openrouter.ai/api/v1/key", _bearer), 

76 "gemini": _KeyProbe( 

77 "https://generativelanguage.googleapis.com/v1beta/models", _gemini_headers, _gemini_rejected 

78 ), 

79 "anthropic": _KeyProbe("https://api.anthropic.com/v1/models", _anthropic_headers), 

80 "openai": _KeyProbe("https://api.openai.com/v1/models", _bearer), 

81 "mistral": _KeyProbe("https://api.mistral.ai/v1/models", _bearer), 

82 "deepseek": _KeyProbe("https://api.deepseek.com/models", _bearer), 

83} 

84 

85_probe_condition = threading.Condition() 

86 

87 

88def _http_get(url: str, *, headers: dict[str, str]) -> httpx.Response: 

89 """GET one probe URL (module seam; tests stub here).""" 

90 return httpx.get(url, headers=headers, timeout=KEY_CHECK_TIMEOUT_S) 

91 

92 

93def _sendable(api_key: str) -> bool: 

94 """True when the key can travel in an HTTP header: printable ASCII only.""" 

95 return api_key.isascii() and api_key.isprintable() 

96 

97 

98def _probe_status(provider: str, api_key: str) -> KeyStatus: 

99 """Send one probe and read the verdict from its status.""" 

100 probe = _PROBES[provider] 

101 resp = _http_get(probe.url, headers=probe.headers(api_key)) 

102 if probe.rejected(resp): 

103 return KeyStatus.INVALID_KEY 

104 if resp.is_error: 

105 log.warning("Could not verify the %s API key: HTTP %d", provider, resp.status_code) 

106 return KeyStatus.READY 

107 

108 

109@cached( 

110 TTLCache(maxsize=len(_PROBES) * _KEYS_CACHED_PER_PROVIDER, ttl=KEY_CHECK_TTL_S), 

111 lock=_probe_condition, 

112 condition=_probe_condition, 

113) 

114def _checked_key_status(provider: str, api_key: str) -> KeyStatus: 

115 """INVALID_KEY for a key that cannot be sent or is refused; an unanswered check stays READY.""" 

116 if not _sendable(api_key): 

117 return KeyStatus.INVALID_KEY 

118 try: 

119 return _probe_status(provider, api_key) 

120 except Exception as exc: # a failed check must never fail the catalog or chat routing 

121 log.warning("Could not verify the %s API key: %s", provider, type(exc).__name__) 

122 return KeyStatus.READY 

123 

124 

125def provider_api_key_in_use(provider: str) -> str | None: 

126 """The key a hosted request carries: its env var when set, else the lilbee config field.""" 

127 return os.environ.get(PROVIDER_API_KEY_ENV[provider]) or get_provider_api_key(provider) 

128 

129 

130def provider_key_set(provider: str) -> bool: 

131 """True when *provider* has an API key in its env var or the lilbee config.""" 

132 return provider_api_key_in_use(provider) is not None 

133 

134 

135def provider_key_status(provider: str) -> KeyStatus: 

136 """Whether *provider*'s API key is missing, rejected, or usable.""" 

137 api_key = provider_api_key_in_use(provider) 

138 if api_key is None: 

139 return KeyStatus.MISSING_KEY 

140 return _checked_key_status(provider, api_key) 

141 

142 

143def provider_key_statuses(providers: list[str]) -> dict[str, KeyStatus]: 

144 """Key status per provider, the uncached checks running concurrently.""" 

145 if not providers: 

146 return {} 

147 with ThreadPoolExecutor(max_workers=len(providers)) as pool: 

148 return dict(zip(providers, pool.map(provider_key_status, providers), strict=True))