Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/openai_files_endpoints/batch_file_validation.py: 34%

1import json 

2from collections.abc import Iterator 

3from dataclasses import dataclass 

4from itertools import chain 

5from typing import BinaryIO, Final, NoReturn 

6 

7from typing_extensions import assert_never 

8 

9from litellm.proxy._types import ProxyException 

10 

11_MB: Final = 1024 * 1024 

12 

13 

14@dataclass(frozen=True, slots=True) 

15class BatchLineShape: 

16 required_keys: tuple[str, ...] 

17 hint: str 

18 

19 

20BATCH_LINE_SHAPE: Final = BatchLineShape( 

21 required_keys=("custom_id", "method", "url", "body"), 

22 hint="Each line must be a JSON object with keys custom_id, method, url, body", 

23) 

24PASSTHROUGH_BATCH_LINE_SHAPE: Final = BatchLineShape( 

25 required_keys=("request",), 

26 hint="A passthrough upload takes native Vertex batch rows, so each line must be a JSON object with a request key", 

27) 

28 

29 

30@dataclass(frozen=True, slots=True) 

31class BatchFileTooLarge: 

32 size_bytes: int 

33 limit_mb: int 

34 

35 

36@dataclass(frozen=True, slots=True) 

37class BatchFileWrongExtension: 

38 filename: str 

39 

40 

41@dataclass(frozen=True, slots=True) 

42class BatchFileEmpty: 

43 pass 

44 

45 

46@dataclass(frozen=True, slots=True) 

47class BatchFileInvalidJsonLine: 

48 line_number: int 

49 

50 

51@dataclass(frozen=True, slots=True) 

52class BatchFileLineNotObject: 

53 line_number: int 

54 

55 

56@dataclass(frozen=True, slots=True) 

57class BatchFileMissingLineKey: 

58 line_number: int 

59 key: str 

60 line_shape: BatchLineShape = BATCH_LINE_SHAPE 

61 

62 

63BatchFileValidationFailure = ( 

64 BatchFileTooLarge 

65 | BatchFileWrongExtension 

66 | BatchFileEmpty 

67 | BatchFileInvalidJsonLine 

68 | BatchFileLineNotObject 

69 | BatchFileMissingLineKey 

70) 

71 

72 

73def _file_size_bytes(file_source: bytes | BinaryIO) -> int: 

74 if isinstance(file_source, bytes): 

75 return len(file_source) 

76 file_source.seek(0, 2) 

77 size: Final = file_source.tell() 

78 file_source.seek(0) 

79 return size 

80 

81 

82def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: 

83 if isinstance(file_source, bytes): 

84 return iter(file_source.splitlines()) 

85 file_source.seek(0) 

86 return iter(file_source) 

87 

88 

89def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: 

90 try: 

91 parsed: Final = json.loads(raw_line) 

92 except (json.JSONDecodeError, UnicodeDecodeError): 

93 return BatchFileInvalidJsonLine(line_number=line_number) 

94 if not isinstance(parsed, dict): 

95 return BatchFileLineNotObject(line_number=line_number) 

96 missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) 

97 if missing is None: 

98 return None 

99 return BatchFileMissingLineKey(line_number=line_number, key=missing, line_shape=line_shape) 

100 

101 

102def _scan_lines(file_source: bytes | BinaryIO, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: 

103 content_lines: Final = ( 

104 (line_number, raw_line) 

105 for line_number, raw_line in enumerate(_iter_lines(file_source), start=1) 

106 if raw_line.strip() 

107 ) 

108 first_line: Final = next(content_lines, None) 

109 if first_line is None: 

110 return BatchFileEmpty() 

111 return next( 

112 ( 

113 failure 

114 for line_number, raw_line in chain((first_line,), content_lines) 

115 for failure in (_check_line(line_number, raw_line, line_shape),) 

116 if failure is not None 

117 ), 

118 None, 

119 ) 

120 

121 

122def check_batch_file_upload( 

123 filename: str | None, 

124 file_source: bytes | BinaryIO, 

125 max_batch_file_size_mb: int | None, 

126 line_shape: BatchLineShape = BATCH_LINE_SHAPE, 

127) -> BatchFileValidationFailure | None: 

128 if filename is None or not filename.lower().endswith(".jsonl"): 

129 return BatchFileWrongExtension(filename=filename or "") 

130 if max_batch_file_size_mb is not None and max_batch_file_size_mb > 0: 

131 size_bytes: Final = _file_size_bytes(file_source) 

132 if size_bytes > max_batch_file_size_mb * _MB: 

133 return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) 

134 scan_failure: Final = _scan_lines(file_source, line_shape) 

135 if not isinstance(file_source, bytes): 

136 file_source.seek(0) 

137 return scan_failure 

138 

139 

140def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> NoReturn: 

141 match failure: 

142 case BatchFileTooLarge(size_bytes=size_bytes, limit_mb=limit_mb): 

143 raise ProxyException( 

144 message=( 

145 f"Batch input file is {size_bytes / _MB:.1f} MB, which exceeds the configured " 

146 f"max_batch_file_size_mb of {limit_mb} MB. The file was not forwarded to the provider." 

147 ), 

148 type="invalid_request_error", 

149 param="file", 

150 code=413, 

151 ) 

152 case BatchFileWrongExtension(filename=filename): 

153 raise ProxyException( 

154 message=( 

155 f"Invalid file format for Batch API: '{filename}'. " 

156 "Batch input files must be .jsonl files. The file was not forwarded to the provider." 

157 ), 

158 type="invalid_request_error", 

159 param="file", 

160 code=400, 

161 ) 

162 case BatchFileEmpty(): 

163 raise ProxyException( 

164 message="Batch input file has no request lines. The file was not forwarded to the provider.", 

165 type="invalid_request_error", 

166 param="file", 

167 code=400, 

168 ) 

169 case BatchFileInvalidJsonLine(line_number=line_number): 

170 raise ProxyException( 

171 message=( 

172 f"Batch input file line {line_number} is not valid JSON. " 

173 "The file was not forwarded to the provider." 

174 ), 

175 type="invalid_request_error", 

176 param="file", 

177 code=400, 

178 ) 

179 case BatchFileLineNotObject(line_number=line_number): 

180 raise ProxyException( 

181 message=( 

182 f"Batch input file line {line_number} must be a JSON object. " 

183 "The file was not forwarded to the provider." 

184 ), 

185 type="invalid_request_error", 

186 param="file", 

187 code=400, 

188 ) 

189 case BatchFileMissingLineKey(line_number=line_number, key=key, line_shape=line_shape): 

190 raise ProxyException( 

191 message=( 

192 f"Missing required parameter: '{key}' (batch input file line {line_number}). " 

193 f"{line_shape.hint}. " 

194 "The file was not forwarded to the provider." 

195 ), 

196 type="invalid_request_error", 

197 param=key, 

198 code=400, 

199 ) 

200 case _: 

201 assert_never(failure)