Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/keys.py: 0%

1import builtins 

2from collections.abc import Mapping 

3from typing import Any, Final 

4 

5import requests 

6 

7from litellm.litellm_core_utils.secret_redaction import redact_string 

8 

9from .exceptions import UnauthorizedError 

10 

11 

12class KeysManagementClient: 

13 def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): 

14 """ 

15 Initialize the KeysManagementClient. 

16 

17 Args: 

18 base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") 

19 api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. 

20 timeout (int): Request timeout in seconds (default: 30) 

21 """ 

22 self._base_url = base_url.rstrip("/") # Remove trailing slash if present 

23 self._api_key = api_key 

24 self._timeout = timeout 

25 

26 def _get_headers(self) -> dict[str, str]: 

27 """ 

28 Get the headers for API requests, including authorization if api_key is set. 

29 

30 Returns: 

31 Dict[str, str]: Headers to use for API requests 

32 """ 

33 headers: Final = {"Content-Type": "application/json"} 

34 if self._api_key: 

35 headers["Authorization"] = f"Bearer {self._api_key}" 

36 return headers 

37 

38 def list( 

39 self, 

40 page: int | None = None, 

41 size: int | None = None, 

42 user_id: str | None = None, 

43 team_id: str | None = None, 

44 organization_id: str | None = None, 

45 key_hash: str | None = None, 

46 key_alias: str | None = None, 

47 return_full_object: bool | None = None, 

48 include_team_keys: bool | None = None, 

49 return_request: bool = False, 

50 ) -> dict[str, Any] | requests.Request: 

51 """ 

52 List all API keys with optional filtering and pagination. 

53 

54 Args: 

55 page (Optional[int]): Page number for pagination 

56 size (Optional[int]): Number of items per page 

57 user_id (Optional[str]): Filter keys by user ID 

58 team_id (Optional[str]): Filter keys by team ID 

59 organization_id (Optional[str]): Filter keys by organization ID 

60 key_hash (Optional[str]): Filter by specific key hash 

61 key_alias (Optional[str]): Filter by key alias 

62 return_full_object (Optional[bool]): Whether to return the full key object 

63 include_team_keys (Optional[bool]): Whether to include team keys in the response 

64 return_request (bool): If True, returns the prepared request object instead of executing it 

65 

66 Returns: 

67 Union[Dict[str, Any], requests.Request]: Either the response from the server or 

68 a prepared request object if return_request is True. The response contains a list 

69 of API keys with their configurations. 

70 

71 Raises: 

72 UnauthorizedError: If the request fails with a 401 status code 

73 requests.exceptions.RequestException: If the request fails with any other error 

74 """ 

75 url: Final = f"{self._base_url}/key/list" 

76 params: Final[dict[str, int | str]] = {} 

77 

78 # Add optional query parameters 

79 if page is not None: 

80 params["page"] = page 

81 if size is not None: 

82 params["size"] = size 

83 if user_id is not None: 

84 params["user_id"] = user_id 

85 if team_id is not None: 

86 params["team_id"] = team_id 

87 if organization_id is not None: 

88 params["organization_id"] = organization_id 

89 if key_hash is not None: 

90 params["key_hash"] = key_hash 

91 if key_alias is not None: 

92 params["key_alias"] = key_alias 

93 if return_full_object is not None: 

94 params["return_full_object"] = str(return_full_object).lower() 

95 if include_team_keys is not None: 

96 params["include_team_keys"] = str(include_team_keys).lower() 

97 

98 request: Final = requests.Request("GET", url, headers=self._get_headers(), params=params) 

99 

100 if return_request: 

101 return request 

102 

103 session: Final = requests.Session() 

104 try: 

105 response: Final = session.send(request.prepare(), timeout=self._timeout) 

106 response.raise_for_status() 

107 return response.json() 

108 except requests.exceptions.HTTPError as e: 

109 if e.response.status_code == 401: 

110 raise UnauthorizedError(e) 

111 raise 

112 

113 def generate( 

114 self, 

115 models: builtins.list[str] | None = None, 

116 aliases: dict[str, str] | None = None, 

117 spend: float | None = None, 

118 duration: str | None = None, 

119 key_alias: str | None = None, 

120 team_id: str | None = None, 

121 user_id: str | None = None, 

122 budget_id: str | None = None, 

123 config: Mapping[str, object] | None = None, 

124 return_request: bool = False, 

125 ) -> dict[str, object] | requests.Request: 

126 """ 

127 Generate an API key based on the provided data. 

128 

129 Docs: https://docs.litellm.ai/docs/proxy/virtual_keys 

130 

131 Args: 

132 models (Optional[List[str]]): List of allowed models for this key 

133 aliases (Optional[Dict[str, str]]): Model alias mappings 

134 spend (Optional[float]): Maximum spend limit for this key 

135 duration (Optional[str]): Duration for which the key is valid (e.g. "24h", "7d") 

136 key_alias (Optional[str]): Alias/name for the key for easier identification 

137 team_id (Optional[str]): Team ID to associate the key with 

138 user_id (Optional[str]): User ID to associate the key with 

139 budget_id (Optional[str]): Budget ID to associate the key with 

140 config (Optional[Dict[str, Any]]): Additional configuration parameters 

141 return_request (bool): If True, returns the prepared request object instead of executing it 

142 

143 Returns: 

144 Union[Dict[str, Any], requests.Request]: Either the response from the server or 

145 a prepared request object if return_request is True 

146 

147 Raises: 

148 UnauthorizedError: If the request fails with a 401 status code 

149 requests.exceptions.RequestException: If the request fails with any other error 

150 """ 

151 url: Final = f"{self._base_url}/key/generate" 

152 

153 data: Final[dict[str, object]] = {} 

154 if models is not None: 

155 data["models"] = models 

156 if aliases is not None: 

157 data["aliases"] = aliases 

158 if spend is not None: 

159 data["spend"] = spend 

160 if duration is not None: 

161 data["duration"] = duration 

162 if key_alias is not None: 

163 data["key_alias"] = key_alias 

164 if team_id is not None: 

165 data["team_id"] = team_id 

166 if user_id is not None: 

167 data["user_id"] = user_id 

168 if budget_id is not None: 

169 data["budget_id"] = budget_id 

170 if config is not None: 

171 data["config"] = config 

172 

173 request: Final = requests.Request("POST", url, headers=self._get_headers(), json=data) 

174 

175 if return_request: 

176 return request 

177 

178 session: Final = requests.Session() 

179 try: 

180 response: Final = session.send(request.prepare(), timeout=self._timeout) 

181 response.raise_for_status() 

182 return response.json() 

183 except requests.exceptions.HTTPError as e: 

184 if e.response.status_code == 401: 

185 raise UnauthorizedError(e) 

186 raise 

187 

188 def delete( 

189 self, 

190 keys: builtins.list[str] | None = None, 

191 key_aliases: builtins.list[str] | None = None, 

192 return_request: bool = False, 

193 ) -> dict[str, object] | requests.Request: 

194 """ 

195 Delete existing keys 

196 

197 Args: 

198 keys (List[str]): List of API keys to delete 

199 key_aliases (List[str]): List of key aliases to delete 

200 return_request (bool): If True, returns the prepared request object instead of executing it 

201 

202 Returns: 

203 Union[Dict[str, Any], requests.Request]: Either the response from the server or 

204 a prepared request object if return_request is True 

205 

206 Raises: 

207 UnauthorizedError: If the request fails with a 401 status code 

208 requests.exceptions.RequestException: If the request fails with any other error 

209 """ 

210 url: Final = f"{self._base_url}/key/delete" 

211 

212 data: Final = { 

213 "keys": keys, 

214 "key_aliases": key_aliases, 

215 } 

216 

217 request: Final = requests.Request("POST", url, headers=self._get_headers(), json=data) 

218 

219 if return_request: 

220 return request 

221 

222 session: Final = requests.Session() 

223 try: 

224 response: Final = session.send(request.prepare(), timeout=self._timeout) 

225 response.raise_for_status() 

226 return response.json() 

227 except requests.exceptions.HTTPError as e: 

228 if e.response.status_code == 401: 

229 raise UnauthorizedError(e) 

230 raise 

231 

232 def update( 

233 self, 

234 key: str, 

235 models: builtins.list[str] | None = None, 

236 aliases: dict[str, str] | None = None, 

237 spend: float | None = None, 

238 duration: str | None = None, 

239 key_alias: str | None = None, 

240 team_id: str | None = None, 

241 user_id: str | None = None, 

242 ) -> dict[str, object] | requests.Request: 

243 """ 

244 Update an existing API key's parameters. 

245 

246 Args: 

247 models: Optional[List[str]] = None, 

248 aliases: Optional[Dict[str, str]] = None, 

249 spend: Optional[float] = None, 

250 duration: Optional[str] = None, 

251 key_alias: Optional[str] = None, 

252 team_id: Optional[str] = None, 

253 user_id: Optional[str] = None, 

254 

255 Returns: 

256 Union[Dict[str, Any], requests.Request]: Either the response from the server or 

257 a prepared request object if return_request is True 

258 

259 Raises: 

260 UnauthorizedError: If the request fails with a 401 status code 

261 requests.exceptions.RequestException: If the request fails with any other error 

262 """ 

263 url: Final = f"{self._base_url}/key/update" 

264 

265 data: Final[dict[str, object]] = {"key": key} 

266 

267 if key_alias is not None: 

268 data["key_alias"] = key_alias 

269 if user_id is not None: 

270 data["user_id"] = user_id 

271 if team_id is not None: 

272 data["team_id"] = team_id 

273 if models is not None: 

274 data["models"] = models 

275 if spend is not None: 

276 data["spend"] = spend 

277 if duration is not None: 

278 data["duration"] = duration 

279 if aliases is not None: 

280 data["aliases"] = aliases 

281 request: Final = requests.Request("POST", url, headers=self._get_headers(), json=data) 

282 session: Final = requests.Session() 

283 response_text: str | None = None 

284 try: 

285 response: Final = session.send(request.prepare(), timeout=self._timeout) 

286 response_text = response.text 

287 response.raise_for_status() 

288 return response.json() 

289 except Exception: 

290 raise Exception(f"Error updating key: {response_text}") 

291 

292 def info(self, key: str, return_request: bool = False) -> dict[str, object] | requests.Request: 

293 """ 

294 Get information about API keys. 

295 

296 Args: 

297 key (str): The key hash to get information about 

298 return_request (bool): If True, returns the prepared request object instead of executing it 

299 

300 Returns: 

301 Union[Dict[str, Any], requests.Request]: Either the response from the server or a prepared request object if return_request is True 

302 

303 Raises: 

304 UnauthorizedError: If the request fails with a 401 status code 

305 requests.exceptions.RequestException: If the request fails with any other error 

306 """ 

307 url: Final = f"{self._base_url}/key/info?key={key}" 

308 request: Final = requests.Request("GET", url, headers=self._get_headers()) 

309 

310 if return_request: 

311 return request 

312 

313 session: Final = requests.Session() 

314 try: 

315 response: Final = session.send(request.prepare(), timeout=self._timeout) 

316 response.raise_for_status() 

317 return response.json() 

318 except requests.exceptions.HTTPError as e: 

319 redacted_message: Final = redact_string(str(e)) 

320 if e.response.status_code == 401: 

321 raise UnauthorizedError(e) from None 

322 raise requests.exceptions.HTTPError(redacted_message, response=e.response) from None