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

1import builtins 

2from collections.abc import Mapping 

3from typing import Any, Final 

4 

5import requests 

6 

7from .exceptions import NotFoundError, UnauthorizedError 

8 

9 

10class ModelsManagementClient: 

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

12 """ 

13 Initialize the ModelsManagementClient. 

14 

15 Args: 

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

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

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

19 """ 

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

21 self._api_key = api_key 

22 self._timeout = timeout 

23 

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

25 """ 

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

27 

28 Returns: 

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

30 """ 

31 headers: Final = {} 

32 if self._api_key: 

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

34 return headers 

35 

36 def list(self, return_request: bool = False) -> builtins.list[dict[str, Any]] | requests.Request: 

37 """ 

38 Get the list of models supported by the server. 

39 

40 Args: 

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

42 Useful for inspection or modification before sending. 

43 

44 Returns: 

45 Union[List[Dict[str, Any]], requests.Request]: Either a list of model information dictionaries 

46 or a prepared request object if return_request is True. 

47 

48 Raises: 

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

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

51 """ 

52 url: Final = f"{self._base_url}/models" 

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

54 

55 if return_request: 

56 return request 

57 

58 # Prepare and send the request 

59 session: Final = requests.Session() 

60 try: 

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

62 response.raise_for_status() 

63 return response.json()["data"] 

64 except requests.exceptions.HTTPError as e: 

65 if e.response.status_code == 401: 

66 raise UnauthorizedError(e) 

67 raise 

68 

69 def new( 

70 self, 

71 model_name: str, 

72 model_params: Mapping[str, object], 

73 model_info: Mapping[str, object] | None = None, 

74 return_request: bool = False, 

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

76 """ 

77 Add a new model to the proxy. 

78 

79 Args: 

80 model_name (str): Name of the model to add 

81 model_params (Dict[str, Any]): Parameters for the model (e.g., model type, api_base, api_key) 

82 model_info (Optional[Dict[str, Any]]): Additional information about the model 

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

84 

85 Returns: 

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

87 a prepared request object if return_request is True 

88 

89 Raises: 

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

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

92 """ 

93 url: Final = f"{self._base_url}/model/new" 

94 

95 data: Final = { 

96 "model_name": model_name, 

97 "litellm_params": model_params, 

98 } 

99 if model_info: 

100 data["model_info"] = model_info 

101 

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

103 

104 if return_request: 

105 return request 

106 

107 # Prepare and send the request 

108 session: Final = requests.Session() 

109 try: 

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

111 response.raise_for_status() 

112 return response.json() 

113 except requests.exceptions.HTTPError as e: 

114 if e.response.status_code == 401: 

115 raise UnauthorizedError(e) 

116 raise 

117 

118 def delete(self, model_id: str, return_request: bool = False) -> dict[str, Any] | requests.Request: 

119 """ 

120 Delete a model from the proxy. 

121 

122 Args: 

123 model_id (str): ID of the model to delete (e.g., "2f23364f-4579-4d79-a43a-2d48dd551c2e") 

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

125 

126 Returns: 

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

128 a prepared request object if return_request is True 

129 

130 Raises: 

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

132 NotFoundError: If the request fails with a 404 status code or indicates the model was not found 

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

134 """ 

135 url: Final = f"{self._base_url}/model/delete" 

136 data: Final = {"id": model_id} 

137 

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

139 

140 if return_request: 

141 return request 

142 

143 # Prepare and send the request 

144 session: Final = requests.Session() 

145 try: 

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

147 response.raise_for_status() 

148 return response.json() 

149 except requests.exceptions.HTTPError as e: 

150 if e.response.status_code == 401: 

151 raise UnauthorizedError(e) 

152 if e.response.status_code == 404 or "not found" in e.response.text.lower(): 

153 raise NotFoundError(e) 

154 raise 

155 

156 def get( 

157 self, 

158 model_id: str | None = None, 

159 model_name: str | None = None, 

160 return_request: bool = False, 

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

162 """ 

163 Get information about a specific model by its ID or name. 

164 

165 Args: 

166 model_id (Optional[str]): ID of the model to retrieve 

167 model_name (Optional[str]): Name of the model to retrieve 

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

169 

170 Returns: 

171 Union[Dict[str, Any], requests.Request]: Either the model information from the server or 

172 a prepared request object if return_request is True 

173 

174 Raises: 

175 ValueError: If neither model_id nor model_name is provided, or if both are provided 

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

177 NotFoundError: If the model is not found 

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

179 """ 

180 if (model_id is None and model_name is None) or (model_id is not None and model_name is not None): 

181 raise ValueError("Exactly one of model_id or model_name must be provided") 

182 

183 # If return_request is True, delegate to info 

184 if return_request: 

185 result: Final = self.info(return_request=True) 

186 assert isinstance(result, requests.Request) 

187 return result 

188 

189 # Get all models and filter 

190 models: Final = self.info() 

191 assert isinstance(models, list) 

192 

193 # Find the matching model 

194 for model in models: 

195 if (model_id and model.get("model_info", {}).get("id") == model_id) or ( 

196 model_name and model.get("model_name") == model_name 

197 ): 

198 return model 

199 

200 # If we get here, no model was found 

201 if model_id: 

202 msg = f"Model with id={model_id} not found" 

203 elif model_name: 

204 msg = f"Model with model_name={model_name} not found" 

205 else: 

206 msg = "Unknown error trying to find model" 

207 raise NotFoundError( 

208 requests.exceptions.HTTPError( 

209 msg, 

210 response=requests.Response(), # Empty response since we didn't make a direct request 

211 ) 

212 ) 

213 

214 def info(self, return_request: bool = False) -> builtins.list[dict[str, Any]] | requests.Request: 

215 """ 

216 Get detailed information about all models from the server. 

217 

218 Args: 

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

220 

221 Returns: 

222 Union[List[Dict[str, Any]], requests.Request]: Either a list of model information dictionaries 

223 or a prepared request object if return_request is True 

224 

225 Raises: 

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

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

228 """ 

229 url: Final = f"{self._base_url}/v1/model/info" 

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

231 

232 if return_request: 

233 return request 

234 

235 # Prepare and send the request 

236 session: Final = requests.Session() 

237 try: 

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

239 response.raise_for_status() 

240 return response.json()["data"] 

241 except requests.exceptions.HTTPError as e: 

242 if e.response.status_code == 401: 

243 raise UnauthorizedError(e) 

244 raise 

245 

246 def update( 

247 self, 

248 model_id: str, 

249 model_params: Mapping[str, object], 

250 model_info: Mapping[str, object] | None = None, 

251 return_request: bool = False, 

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

253 """ 

254 Update an existing model's configuration. 

255 

256 Args: 

257 model_id (str): ID of the model to update 

258 model_params (Dict[str, Any]): New parameters for the model (e.g., model type, api_base, api_key) 

259 model_info (Optional[Dict[str, Any]]): Additional information about the model 

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

261 

262 Returns: 

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

264 a prepared request object if return_request is True 

265 

266 Raises: 

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

268 NotFoundError: If the model is not found 

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

270 """ 

271 url: Final = f"{self._base_url}/model/update" 

272 

273 data: Final = { 

274 "id": model_id, 

275 "litellm_params": model_params, 

276 } 

277 if model_info: 

278 data["model_info"] = model_info 

279 

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

281 

282 if return_request: 

283 return request 

284 

285 # Prepare and send the request 

286 session: Final = requests.Session() 

287 try: 

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

289 response.raise_for_status() 

290 return response.json() 

291 except requests.exceptions.HTTPError as e: 

292 if e.response.status_code == 401: 

293 raise UnauthorizedError(e) 

294 if e.response.status_code == 404 or "not found" in e.response.text.lower(): 

295 raise NotFoundError(e) 

296 raise