Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/load_config_utils.py: 19%

1import asyncio 

2import os 

3import posixpath 

4from collections.abc import Awaitable, Mapping 

5from types import MappingProxyType 

6from typing import TYPE_CHECKING, Final, Protocol 

7 

8import yaml 

9from pydantic import TypeAdapter, ValidationError 

10 

11from litellm._logging import verbose_proxy_logger 

12from litellm.proxy.common_utils.config_includes import resolve_includes 

13 

14if TYPE_CHECKING: 14 ↛ 15line 14 didn't jump to line 15 because the condition on line 14 was never true

15 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase 

16 

17_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) 

18 

19 

20class BucketObjectFetcher(Protocol): 

21 def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... 21 ↛ exitline 21 didn't return from function '__call__' because

22 

23 

24class BucketObjectReader(Protocol): 

25 def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... 25 ↛ exitline 25 didn't return from function '__call__' because

26 

27 

28class SyncBucketObjectReader(Protocol): 

29 def __call__(self, object_key: str, /) -> object | None: ... 29 ↛ exitline 29 didn't return from function '__call__' because

30 

31 

32def _parsed_config(object_key: str, file_contents: str) -> object | None: 

33 try: 

34 parsed: Final = yaml.safe_load(file_contents) 

35 except yaml.YAMLError as e: 

36 verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) 

37 return None 

38 return MappingProxyType({}) if parsed is None else parsed 

39 

40 

41def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: 

42 """ 

43 Build one reader for a whole config, so an `include` tree costs one S3 client rather than one per object. 

44 """ 

45 try: 

46 # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc 

47 import boto3 

48 from botocore.credentials import Credentials 

49 

50 from litellm.main import bedrock_converse_chat_completion 

51 

52 credentials: Final[Credentials] = bedrock_converse_chat_completion.get_credentials() 

53 s3_client: Final = boto3.client( 

54 "s3", 

55 aws_access_key_id=credentials.access_key, 

56 aws_secret_access_key=credentials.secret_key, 

57 aws_session_token=credentials.token, # Optional, if using temporary credentials 

58 ) 

59 except ImportError as e: 

60 # this is most likely if a user is not using the litellm docker container 

61 verbose_proxy_logger.error("ImportError: %s", e) 

62 return lambda object_key: None 

63 except Exception as e: 

64 verbose_proxy_logger.error("Error creating the S3 client for bucket %s: %s", bucket_name, e) 

65 return lambda object_key: None 

66 

67 def read(object_key: str) -> object | None: 

68 try: 

69 verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) 

70 response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) 

71 file_contents: Final = response["Body"].read().decode("utf-8") 

72 except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object 

73 verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e) 

74 return None 

75 

76 return _parsed_config(object_key, file_contents) 

77 

78 return read 

79 

80 

81def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | None: 

82 return s3_object_reader(bucket_name)(object_key) 

83 

84 

85def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": 

86 """ 

87 Build a plain GCS client for reading config objects. 

88 

89 Reading a config out of a bucket is not GCS logging, so it neither needs the enterprise license 

90 that gate covers nor the batching task the logger starts and never stops. 

91 """ 

92 try: 

93 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase 

94 

95 return GCSBucketBase(bucket_name=bucket_name) 

96 except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket 

97 verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) 

98 return None 

99 

100 

101async def get_config_file_contents_from_gcs( 

102 bucket_name: str, 

103 object_key: str, 

104 gcs_bucket: "GCSBucketBase | None" = None, 

105) -> object | None: 

106 try: 

107 bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket 

108 if bucket is None: 

109 return None 

110 file_contents: Final = await bucket.download_gcs_object(object_key) 

111 if file_contents is None: 

112 raise Exception(f"File contents are None for {object_key}") 

113 decoded: Final = file_contents.decode("utf-8") 

114 

115 except Exception as e: 

116 verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e) 

117 return None 

118 

119 return _parsed_config(object_key, decoded) 

120 

121 

122def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: 

123 """ 

124 Resolve one `include` entry to the object key it names, relative to the config object's prefix. 

125 

126 A leading "/" means the bucket root, mirroring how an absolute path on disk ignores the 

127 directory the including config sits in. 

128 """ 

129 if include_entry.startswith("/"): 

130 return posixpath.normpath(include_entry).lstrip("/") 

131 return posixpath.normpath(posixpath.join(posixpath.dirname(config_object_key), include_entry)) 

132 

133 

134async def resolve_bucket_includes( 

135 *, 

136 config: Mapping[str, object], 

137 object_key: str, 

138 fetch: BucketObjectFetcher, 

139) -> dict[str, object]: 

140 async def read(include_key: str) -> Mapping[str, object]: 

141 included: Final = await fetch(include_key) 

142 if included is None: 

143 raise FileNotFoundError( 

144 f"Included config could not be read from bucket: {include_key}. " 

145 "The underlying bucket error is logged above." 

146 ) 

147 return included 

148 

149 def resolve(include_entry: str, declared_in: str) -> str: 

150 return resolve_include_object_key(declared_in, include_entry) 

151 

152 return await resolve_includes(config=config, location=object_key, resolve=resolve, read=read) 

153 

154 

155async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: 

156 """ 

157 Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object. 

158 """ 

159 if bucket_type != "gcs": 

160 read_object: Final = await asyncio.to_thread(s3_object_reader, bucket_name) 

161 

162 async def read_from_s3(object_key: str) -> object | None: 

163 return await asyncio.to_thread(read_object, object_key) 

164 

165 return read_from_s3 

166 

167 gcs_bucket: Final = gcs_config_bucket(bucket_name) 

168 

169 async def read_from_gcs(object_key: str) -> object | None: 

170 if gcs_bucket is None: 

171 return None 

172 return await get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket) 

173 

174 return read_from_gcs 

175 

176 

177async def get_config_from_bucket( 

178 *, 

179 bucket_type: str | None, 

180 bucket_name: str, 

181 object_key: str, 

182) -> dict[str, object] | None: 

183 read: Final = await bucket_object_reader(bucket_type, bucket_name) 

184 

185 async def fetch(key: str) -> Mapping[str, object] | None: 

186 raw: Final = await read(key) 

187 if raw is None: 

188 return None 

189 try: 

190 return _BUCKET_CONFIG_ADAPTER.validate_python(raw) 

191 except ValidationError as e: 

192 raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e 

193 

194 config: Final = await fetch(object_key) 

195 if not config: 

196 return None 

197 

198 return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) 

199 

200 

201def download_python_file_from_s3( 

202 bucket_name: str, 

203 object_key: str, 

204 local_file_path: str, 

205) -> bool: 

206 """ 

207 Download a Python file from S3 and save it to local filesystem. 

208 

209 Args: 

210 bucket_name (str): S3 bucket name 

211 object_key (str): S3 object key (file path in bucket) 

212 local_file_path (str): Local path where file should be saved 

213 

214 Returns: 

215 bool: True if successful, False otherwise 

216 """ 

217 try: 

218 import boto3 

219 from botocore.credentials import Credentials 

220 

221 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM 

222 

223 base_aws_llm: Final = BaseAWSLLM() 

224 

225 credentials: Final[Credentials] = base_aws_llm.get_credentials() 

226 s3_client: Final = boto3.client( 

227 "s3", 

228 aws_access_key_id=credentials.access_key, 

229 aws_secret_access_key=credentials.secret_key, 

230 aws_session_token=credentials.token, 

231 ) 

232 

233 verbose_proxy_logger.debug("Downloading Python file %s from S3 bucket: %s", object_key, bucket_name) 

234 response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) 

235 

236 # Read the file contents 

237 file_contents: Final = response["Body"].read().decode("utf-8") 

238 verbose_proxy_logger.debug("File contents: %s", file_contents) 

239 

240 # Ensure directory exists 

241 os.makedirs(os.path.dirname(local_file_path), exist_ok=True) 

242 

243 # Write to local file 

244 with open(local_file_path, "w") as f: 

245 f.write(file_contents) 

246 

247 verbose_proxy_logger.debug("Python file downloaded successfully to %s", local_file_path) 

248 return True 

249 

250 except ImportError as e: 

251 verbose_proxy_logger.error("ImportError: %s", e) 

252 return False 

253 except Exception as e: 

254 verbose_proxy_logger.exception("Error downloading Python file: %s", e) 

255 return False 

256 

257 

258async def download_python_file_from_gcs( 

259 bucket_name: str, 

260 object_key: str, 

261 local_file_path: str, 

262) -> bool: 

263 """ 

264 Download a Python file from GCS and save it to local filesystem. 

265 

266 Args: 

267 bucket_name (str): GCS bucket name 

268 object_key (str): GCS object key (file path in bucket) 

269 local_file_path (str): Local path where file should be saved 

270 

271 Returns: 

272 bool: True if successful, False otherwise 

273 """ 

274 try: 

275 from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase 

276 

277 gcs_bucket: Final = GCSBucketBase(bucket_name=bucket_name) 

278 file_contents = await gcs_bucket.download_gcs_object(object_key) 

279 if file_contents is None: 

280 raise Exception(f"File contents are None for {object_key}") 

281 

282 # file_contents is a bytes object, decode it 

283 file_contents = file_contents.decode("utf-8") 

284 

285 # Ensure directory exists 

286 os.makedirs(os.path.dirname(local_file_path), exist_ok=True) 

287 

288 # Write to local file 

289 with open(local_file_path, "w") as f: 

290 f.write(file_contents) 

291 

292 verbose_proxy_logger.debug("Python file downloaded successfully to %s", local_file_path) 

293 return True 

294 

295 except Exception as e: 

296 verbose_proxy_logger.exception("Error downloading Python file from GCS: %s", e) 

297 return False 

298 

299 

300# # Example usage 

301# bucket_name = 'litellm-proxy' 

302# object_key = 'litellm_proxy_config.yaml'