Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/batches_endpoints/endpoints.py: 38%

1###################################################################### 

2 

3# /v1/batches Endpoints 

4 

5 

6###################################################################### 

7import asyncio 

8import os 

9from collections.abc import Mapping, MutableMapping 

10from datetime import datetime 

11from types import MappingProxyType 

12from typing import TYPE_CHECKING, Any, Final, Literal, cast 

13 

14from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response 

15from pydantic import TypeAdapter 

16 

17import litellm 

18from litellm._logging import verbose_proxy_logger 

19from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest 

20from litellm.proxy._types import * 

21from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

22from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit 

23from litellm.proxy.batches_endpoints.litellm_executed_batches import ( 

24 LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE, 

25 LiteLLMExecutedBatchRunner, 

26 ManagedBatchStore, 

27 batch_error, 

28 executed_batch_runner_lost, 

29 litellm_executed_provider_for, 

30 resolve_litellm_executed_provider, 

31) 

32from litellm.proxy.common_request_processing import ( 

33 ProxyBaseLLMRequestProcessing, 

34 log_llm_api_exception, 

35 request_litellm_call_id, 

36) 

37from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata 

38from litellm.proxy.common_utils.http_parsing_utils import _read_request_body 

39from litellm.proxy.common_utils.openai_endpoint_utils import ( 

40 get_custom_llm_provider_from_request_headers, 

41 get_custom_llm_provider_from_request_query, 

42) 

43from litellm.proxy.openai_files_endpoints.common_utils import ( 

44 BATCH_CREATE_HIDDEN_PARAM, 

45 _is_base64_encoded_unified_file_id, 

46 add_deployment_model_info, 

47 add_internal_model_credentials, 

48 apply_team_provider_credentials, 

49 authorize_model_for_key, 

50 batch_cost_poller_is_active, 

51 decode_model_from_file_id, 

52 encode_batch_response_ids, 

53 encode_file_id_with_model, 

54 ensure_batch_response_managed_file_ids, 

55 get_authorized_credentials_for_model, 

56 get_batch_from_database, 

57 get_batch_id_from_unified_batch_id, 

58 get_model_id_from_unified_batch_id, 

59 get_models_from_unified_file_id, 

60 get_original_file_id, 

61 is_litellm_executed_batch, 

62 prepare_data_with_credentials, 

63 update_batch_in_database, 

64 validate_managed_id_requirement, 

65) 

66from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata 

67from litellm.proxy.route_llm_request import raise_if_required_body_param_missing 

68from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model 

69from litellm.repositories.managed_batch_repository import ManagedBatchRepository 

70from litellm.repositories.table_repositories import ManagedFileRepository 

71from litellm.router import Router 

72from litellm.types.llms.openai import LiteLLMBatchCreateRequest 

73from litellm.types.utils import LiteLLMBatch 

74 

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

76 from prisma.models import LiteLLM_ManagedObjectTable 

77 

78router: Final = APIRouter() 

79_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) 

80 

81 

82def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None: 

83 metadata: Final = data.get("litellm_metadata") 

84 if metadata is None: 

85 return None 

86 return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata)) 

87 

88 

89def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner: 

90 from litellm.proxy.proxy_server import general_settings, prisma_client 

91 

92 managed_files: Final = proxy_logging_obj.get_proxy_hook("managed_files") 

93 if prisma_client is None or not isinstance(managed_files, ManagedBatchStore): 

94 raise batch_error( 

95 400, 

96 "LiteLLM-executed batches need a database: set DATABASE_URL so LiteLLM can keep the batch and its files", 

97 ) 

98 return LiteLLMExecutedBatchRunner( 

99 llm_router=llm_router, 

100 prisma_client=prisma_client, 

101 managed_files=managed_files, 

102 batches=ManagedBatchRepository(prisma_client), 

103 proxy_logging_obj=proxy_logging_obj, 

104 general_settings=general_settings, 

105 ) 

106 

107 

108async def _batch_from_database( 

109 batch_id: str, 

110 unified_batch_id: str | Literal[False], 

111 executed_batch: bool, 

112 managed_files_obj: object, 

113 prisma_client: PrismaClient | None, 

114 llm_router: Router | None, 

115 proxy_logging_obj: ProxyLogging, 

116 user_api_key_dict: UserAPIKeyAuth, 

117) -> tuple["LiteLLM_ManagedObjectTable | None", LiteLLMBatch | None]: 

118 row, batch = await get_batch_from_database( 

119 batch_id=batch_id, 

120 unified_batch_id=unified_batch_id, 

121 managed_files_obj=managed_files_obj, 

122 prisma_client=prisma_client, 

123 verbose_proxy_logger=verbose_proxy_logger, 

124 ) 

125 updated_at: Final[object] = getattr(row, "updated_at", None) 

126 if not executed_batch or batch is None or llm_router is None or not isinstance(updated_at, datetime): 126 ↛ 128line 126 didn't jump to line 128 because the condition on line 126 was always true

127 return row, batch 

128 if not executed_batch_runner_lost(batch.status, updated_at): 

129 return row, batch 

130 runner: Final = _litellm_executed_batch_runner(llm_router, proxy_logging_obj) 

131 return row, await runner.fail_abandoned(batch, user_api_key_dict) 

132 

133 

134async def _raise_when_input_file_must_be_managed(model: str, credentials: Mapping[str, object]) -> None: 

135 if await litellm_executed_provider_for(credentials) is None: 

136 return 

137 raise batch_error( 

138 400, 

139 f"Batches for {model} run inside LiteLLM, so the input file must be a LiteLLM managed file: " 

140 f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}", 

141 ) 

142 

143 

144def _litellm_metadata_of(data: MutableMapping[str, object]) -> MutableMapping[str, object]: 

145 """The request's litellm_metadata mapping, created on the request when it carries none. 

146 

147 The success handler reads this mapping, so a flag or a model group set here has to live 

148 inside it rather than beside it. 

149 """ 

150 existing: Final = data.get("litellm_metadata") 

151 if isinstance(existing, MutableMapping): 

152 return existing 

153 created: Final[dict[str, object]] = {} # mutable-ok: the logging layer copies and extends this mapping 

154 data["litellm_metadata"] = created # rebind-ok: the success handler reads the request's own mapping 

155 return created 

156 

157 

158def _raise_not_found_when_openai_fallback_unservable( 

159 requested_provider: "str | None", 

160 data: Mapping[str, object], 

161 not_found_message: str, 

162) -> None: 

163 if requested_provider is not None: 

164 return 

165 if data.get("api_key") or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY"): 165 ↛ 167line 165 didn't jump to line 167 because the condition on line 165 was always true

166 return 

167 raise ProxyException( 

168 message=not_found_message, 

169 type="invalid_request_error", 

170 param=None, 

171 code=404, 

172 ) 

173 

174 

175async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str | None": 

176 """Resolve a managed (unified) input_file_id to its backend storage_url. 

177 

178 Provider batch handlers (e.g. Vertex AI, which parses a `publishers/` 

179 segment out of the file URI) need a real storage location; the opaque 

180 unified token crashes them. Returns None whenever a storage_url cannot be 

181 produced (no database, lookup error, no managed-file row, or a row without 

182 a storage_url yet) so callers fall back to dispatching the original id, 

183 which the managed-files deployment hook still maps. This adds resolution 

184 without changing behavior on any path that did not resolve before. 

185 """ 

186 from litellm.proxy.proxy_server import prisma_client 

187 

188 if prisma_client is None: 

189 return None 

190 try: 

191 db_file = await ManagedFileRepository(prisma_client).table.find_first(where={"unified_file_id": input_file_id}) 

192 except Exception as e: 

193 verbose_proxy_logger.warning("create_batch: managed file lookup failed for %s: %s", input_file_id, e) 

194 return None 

195 if db_file is None: 

196 return None 

197 return db_file.storage_url or None 

198 

199 

200async def _create_provider_batch_for_managed_file( 

201 llm_router: Router, 

202 create_batch_data: LiteLLMBatchCreateRequest, 

203 input_file_id: str, 

204 unified_file_id: str, 

205) -> LiteLLMBatch: 

206 resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id) 

207 request: Final[LiteLLMBatchCreateRequest] = { 

208 **create_batch_data, 

209 "input_file_id": resolved_storage_url or input_file_id, 

210 "disable_fallbacks": True, 

211 } 

212 response: Final = await llm_router.acreate_batch(**request) 

213 response.input_file_id = input_file_id 

214 response._hidden_params["unified_file_id"] = unified_file_id 

215 return response 

216 

217 

218@router.post( 

219 "/{provider}/v1/batches", 

220 dependencies=[Depends(user_api_key_auth)], 

221 tags=["batch"], 

222) 

223@router.post( 

224 "/v1/batches", 

225 dependencies=[Depends(user_api_key_auth)], 

226 tags=["batch"], 

227) 

228@router.post( 

229 "/batches", 

230 dependencies=[Depends(user_api_key_auth)], 

231 tags=["batch"], 

232) 

233async def create_batch( 

234 request: Request, 

235 fastapi_response: Response, 

236 provider: str | None = None, 

237 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

238): 

239 """ 

240 Create large batches of API requests for asynchronous processing. 

241 This is the equivalent of POST https://api.openai.com/v1/batch 

242 Supports Identical Params as: https://platform.openai.com/docs/api-reference/batch 

243 

244 Example Curl 

245 ``` 

246 curl http://localhost:4000/v1/batches \ 

247 -H "Authorization: Bearer sk-1234" \ 

248 -H "Content-Type: application/json" \ 

249 -d '{ 

250 "input_file_id": "file-abc123", 

251 "endpoint": "/v1/chat/completions", 

252 "completion_window": "24h" 

253 }' 

254 ``` 

255 """ 

256 from litellm.proxy.proxy_server import ( 

257 general_settings, 

258 llm_router, 

259 proxy_config, 

260 proxy_logging_obj, 

261 version, 

262 ) 

263 

264 data: dict = {} 

265 try: 

266 data = await _read_request_body(request=request) 

267 verbose_proxy_logger.debug( 

268 "Request received by LiteLLM:\n%s", 

269 json.dumps(data, indent=4), 

270 ) 

271 base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) 

272 ( 

273 data, 

274 litellm_logging_obj, 

275 ) = await base_llm_response_processor.common_processing_pre_call_logic( 

276 request=request, 

277 general_settings=general_settings, 

278 user_api_key_dict=user_api_key_dict, 

279 version=version, 

280 proxy_logging_obj=proxy_logging_obj, 

281 proxy_config=proxy_config, 

282 route_type="acreate_batch", 

283 ) 

284 data["metadata"] = sanitize_openai_provider_metadata(data.get("metadata")) 

285 

286 raise_if_required_body_param_missing(route_type="acreate_batch", data=data) 

287 

288 ## check if model is a loadbalanced model 

289 router_model: str | None = None 

290 is_router_model = False 

291 if litellm.enable_loadbalancing_on_batch_endpoints is True: 

292 router_model = data.get("model", None) 

293 is_router_model = is_known_model(model=router_model, llm_router=llm_router) 

294 

295 requested_provider: Final = ( 

296 provider 

297 or data.pop("custom_llm_provider", None) 

298 or get_custom_llm_provider_from_request_headers(request=request) 

299 ) 

300 custom_llm_provider: Final = requested_provider or "openai" 

301 _create_batch_data: Final = LiteLLMBatchCreateRequest(**data) 

302 

303 # Apply team-level batch output expiry enforcement 

304 team_metadata: Final = user_api_key_dict.team_metadata or {} 

305 enforced_batch_expiry: Final = team_metadata.get("enforced_batch_output_expires_after") 

306 if enforced_batch_expiry is not None: 

307 if "anchor" not in enforced_batch_expiry or "seconds" not in enforced_batch_expiry: 

308 raise HTTPException( 

309 status_code=500, 

310 detail={ 

311 "error": "Server configuration error: team metadata field 'enforced_batch_output_expires_after' is malformed - must contain 'anchor' and 'seconds' keys. Contact your team or proxy admin to fix this setting.", 

312 }, 

313 ) 

314 if enforced_batch_expiry["anchor"] != "created_at": 

315 raise HTTPException( 

316 status_code=500, 

317 detail={ 

318 "error": f"Server configuration error: team metadata field 'enforced_batch_output_expires_after' has invalid anchor '{enforced_batch_expiry['anchor']}' - must be 'created_at'. Contact your team or proxy admin to fix this setting.", 

319 }, 

320 ) 

321 _create_batch_data["output_expires_after"] = { 

322 "anchor": "created_at", 

323 "seconds": int(enforced_batch_expiry["seconds"]), 

324 } 

325 

326 input_file_id: Final = _create_batch_data.get("input_file_id", None) 

327 await validate_managed_id_requirement( 

328 resource_id=input_file_id, 

329 resource_kind="file", 

330 user_api_key_dict=user_api_key_dict, 

331 managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), 

332 ) 

333 unified_file_id: str | Literal[False] = False 

334 

335 model_from_file_id = None 

336 if input_file_id: 

337 model_from_file_id = decode_model_from_file_id(input_file_id) 

338 unified_file_id = _is_base64_encoded_unified_file_id(input_file_id) 

339 

340 # SCENARIO 1: File ID is encoded with model info 

341 if model_from_file_id is not None and input_file_id: 

342 credentials = await get_authorized_credentials_for_model( 

343 llm_router=llm_router, 

344 model_id=model_from_file_id, 

345 user_api_key_dict=user_api_key_dict, 

346 operation_context="batch creation (file created with model)", 

347 ) 

348 

349 original_file_id: Final = get_original_file_id(input_file_id) 

350 _create_batch_data["input_file_id"] = original_file_id 

351 prepare_data_with_credentials( 

352 data=_create_batch_data, 

353 credentials=credentials, 

354 ) 

355 

356 # Create batch using model credentials 

357 response = await litellm.acreate_batch( 

358 custom_llm_provider=credentials["custom_llm_provider"], 

359 **_create_batch_data, 

360 ) 

361 

362 # Encode the batch ID and related file IDs with model information 

363 if response and hasattr(response, "id") and response.id: 

364 original_batch_id: Final = response.id 

365 encoded_batch_id: Final = encode_file_id_with_model( 

366 file_id=original_batch_id, 

367 model=model_from_file_id, 

368 id_type="batch", 

369 ) 

370 response.id = encoded_batch_id 

371 

372 if hasattr(response, "output_file_id") and response.output_file_id: 

373 response.output_file_id = encode_file_id_with_model( 

374 file_id=response.output_file_id, model=model_from_file_id 

375 ) 

376 

377 if hasattr(response, "error_file_id") and response.error_file_id: 

378 response.error_file_id = encode_file_id_with_model( 

379 file_id=response.error_file_id, model=model_from_file_id 

380 ) 

381 

382 verbose_proxy_logger.debug( 

383 "Created batch using model: %s, original_batch_id: %s, encoded: %s", 

384 model_from_file_id, 

385 original_batch_id, 

386 encoded_batch_id, 

387 ) 

388 

389 response.input_file_id = input_file_id 

390 

391 elif litellm.enable_loadbalancing_on_batch_endpoints is True and is_router_model and router_model is not None: 

392 if llm_router is None: 

393 raise HTTPException( 

394 status_code=500, 

395 detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, 

396 ) 

397 

398 response = await llm_router.acreate_batch(**_create_batch_data) 

399 elif ( 

400 unified_file_id and input_file_id 

401 ): # litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;target_model_names,gpt-4o-mini 

402 target_model_names: Final = get_models_from_unified_file_id(unified_file_id) 

403 ## EXPECTS 1 MODEL 

404 if len(target_model_names) != 1: 

405 raise HTTPException( 

406 status_code=400, 

407 detail={"error": f"Expected 1 model, got {len(target_model_names)}"}, 

408 ) 

409 model: Final = target_model_names[0] 

410 await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict) 

411 _create_batch_data["model"] = model 

412 

413 if llm_router is None: 

414 raise HTTPException( 

415 status_code=500, 

416 detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, 

417 ) 

418 

419 executed_provider: Final = await resolve_litellm_executed_provider( 

420 llm_router, model, user_api_key_dict.team_id 

421 ) 

422 response = ( 

423 await _litellm_executed_batch_runner(llm_router, proxy_logging_obj).create( 

424 create_request=_create_batch_data, 

425 unified_input_file_id=input_file_id, 

426 model=model, 

427 provider=executed_provider, 

428 user_api_key_dict=user_api_key_dict, 

429 request_tags=_request_tags(_create_batch_data), 

430 ) 

431 if executed_provider is not None 

432 else await _create_provider_batch_for_managed_file( 

433 llm_router, _create_batch_data, input_file_id, unified_file_id 

434 ) 

435 ) 

436 else: 

437 # Check if model specified via header/query/body param 

438 model_param: Final = ( 

439 _create_batch_data.get("model") 

440 or request.query_params.get("model") 

441 or request.headers.get("x-litellm-model") 

442 ) 

443 

444 # SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback 

445 if model_param: 

446 # SCENARIO 2: Use model-based routing from header/query/body 

447 credentials = await get_authorized_credentials_for_model( 

448 llm_router=llm_router, 

449 model_id=model_param, 

450 user_api_key_dict=user_api_key_dict, 

451 operation_context="batch creation", 

452 ) 

453 await _raise_when_input_file_must_be_managed(model_param, credentials) 

454 

455 prepare_data_with_credentials( 

456 data=_create_batch_data, 

457 credentials=credentials, 

458 ) 

459 

460 # Create batch using model credentials 

461 response = await litellm.acreate_batch( 

462 custom_llm_provider=credentials["custom_llm_provider"], 

463 **_create_batch_data, 

464 ) 

465 

466 encode_batch_response_ids(response, model=model_param) 

467 

468 verbose_proxy_logger.debug("Created batch using model: %s", model_param) 

469 else: 

470 # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) 

471 apply_team_provider_credentials( 

472 data=cast(dict, _create_batch_data), # cast-ok: TypedDict is a dict at runtime 

473 llm_router=llm_router, 

474 user_api_key_dict=user_api_key_dict, 

475 custom_llm_provider=custom_llm_provider, 

476 ) 

477 _raise_not_found_when_openai_fallback_unservable( 

478 requested_provider=requested_provider, 

479 data=cast(dict, _create_batch_data), # cast-ok: TypedDict is a dict at runtime 

480 not_found_message=f"No such File object: {input_file_id}", 

481 ) 

482 response = await litellm.acreate_batch( 

483 custom_llm_provider=custom_llm_provider, 

484 **_create_batch_data, 

485 ) 

486 

487 response._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True 

488 

489 ### CALL HOOKS ### - modify outgoing data 

490 response = await proxy_logging_obj.post_call_success_hook( 

491 data=data, user_api_key_dict=user_api_key_dict, response=response 

492 ) 

493 

494 ### ALERTING ### 

495 asyncio.create_task( 

496 proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") 

497 ) 

498 

499 ### RESPONSE HEADERS ### 

500 hidden_params: Final = getattr(response, "_hidden_params", {}) or {} 

501 model_id: Final = hidden_params.get("model_id", None) or "" 

502 cache_key: Final = hidden_params.get("cache_key", None) or "" 

503 api_base: Final = hidden_params.get("api_base", None) or "" 

504 

505 fastapi_response.headers.update( 

506 ProxyBaseLLMRequestProcessing.get_custom_headers( 

507 user_api_key_dict=user_api_key_dict, 

508 model_id=model_id, 

509 cache_key=cache_key, 

510 api_base=api_base, 

511 version=version, 

512 model_region=getattr(user_api_key_dict, "allowed_model_region", ""), 

513 request_data=data, 

514 ) 

515 ) 

516 

517 return response 

518 except Exception as e: 

519 await proxy_logging_obj.post_call_failure_hook( 

520 user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data 

521 ) 

522 litellm_call_id: Final = request_litellm_call_id(data) 

523 log_llm_api_exception(e, litellm_call_id) 

524 raise handle_exception_on_proxy(e, litellm_call_id) 

525 

526 

527@router.get( 

528 "/{provider}/v1/batches/{batch_id:path}", 

529 dependencies=[Depends(user_api_key_auth)], 

530 tags=["batch"], 

531) 

532@router.get( 

533 "/v1/batches/{batch_id:path}", 

534 dependencies=[Depends(user_api_key_auth)], 

535 tags=["batch"], 

536) 

537@router.get( 

538 "/batches/{batch_id:path}", 

539 dependencies=[Depends(user_api_key_auth)], 

540 tags=["batch"], 

541) 

542async def retrieve_batch( 

543 request: Request, 

544 fastapi_response: Response, 

545 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

546 provider: str | None = None, 

547 batch_id: str = Path(title="Batch ID to retrieve", description="The ID of the batch to retrieve"), 

548): 

549 """ 

550 Retrieves a batch. 

551 This is the equivalent of GET https://api.openai.com/v1/batches/{batch_id} 

552 Supports Identical Params as: https://platform.openai.com/docs/api-reference/batch/retrieve 

553 

554 Example Curl 

555 ``` 

556 curl http://localhost:4000/v1/batches/batch_abc123 \ 

557 -H "Authorization: Bearer sk-1234" \ 

558 -H "Content-Type: application/json" \ 

559 

560 ``` 

561 """ 

562 from litellm.proxy.proxy_server import ( 

563 general_settings, 

564 llm_router, 

565 proxy_config, 

566 proxy_logging_obj, 

567 version, 

568 ) 

569 

570 data: dict = {} 

571 try: 

572 await validate_managed_id_requirement( 

573 resource_id=batch_id, 

574 resource_kind="batch", 

575 user_api_key_dict=user_api_key_dict, 

576 managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), 

577 ) 

578 model_from_id: Final = decode_model_from_file_id(batch_id) 

579 _retrieve_batch_request: Final = RetrieveBatchRequest( 

580 batch_id=batch_id, 

581 ) 

582 

583 data = cast(dict, _retrieve_batch_request) 

584 unified_batch_id: Final = _is_base64_encoded_unified_file_id(batch_id) 

585 

586 base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) 

587 ( 

588 data, 

589 litellm_logging_obj, 

590 ) = await base_llm_response_processor.common_processing_pre_call_logic( 

591 request=request, 

592 general_settings=general_settings, 

593 user_api_key_dict=user_api_key_dict, 

594 version=version, 

595 proxy_logging_obj=proxy_logging_obj, 

596 proxy_config=proxy_config, 

597 route_type="aretrieve_batch", 

598 ) 

599 

600 unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) if unified_batch_id else None 

601 if unified_model_id is not None: 601 ↛ 602line 601 didn't jump to line 602 because the condition on line 601 was never true

602 resolved_unified_model: Final = ( 

603 llm_router.resolve_model_name_from_model_id(unified_model_id) if llm_router is not None else None 

604 ) 

605 await authorize_model_for_key( 

606 model_id=resolved_unified_model or unified_model_id, 

607 llm_router=llm_router, 

608 user_api_key_dict=user_api_key_dict, 

609 ) 

610 

611 # FIX: First, try to read from ManagedObjectTable for consistent state 

612 managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") 

613 from litellm.proxy.proxy_server import prisma_client 

614 

615 executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id) 

616 db_batch_object, response = await _batch_from_database( 

617 batch_id=batch_id, 

618 unified_batch_id=unified_batch_id, 

619 executed_batch=executed_batch, 

620 managed_files_obj=managed_files_obj, 

621 prisma_client=prisma_client, 

622 llm_router=llm_router, 

623 proxy_logging_obj=proxy_logging_obj, 

624 user_api_key_dict=user_api_key_dict, 

625 ) 

626 

627 if executed_batch and response is None: 627 ↛ 628line 627 didn't jump to line 628 because the condition on line 627 was never true

628 raise batch_error(404, f"No batch found with id '{batch_id}'.") 

629 

630 # If batch is in a terminal state, return immediately. 

631 # Include "complete" (DB-normalized form of "completed"). 

632 if response is not None and ( 632 ↛ 636line 632 didn't jump to line 636 because the condition on line 632 was never true

633 response.status in ("completed", "complete", "failed", "cancelled", "expired") or executed_batch 

634 ): 

635 # Call hooks and return 

636 response = await proxy_logging_obj.post_call_success_hook( 

637 data=data, user_api_key_dict=user_api_key_dict, response=response 

638 ) 

639 

640 # The DB may store raw provider file IDs (before hooks translate them). 

641 # Register any missing managed-file rows and return unified IDs. 

642 if unified_batch_id: 

643 await ensure_batch_response_managed_file_ids( 

644 response=response, 

645 managed_files_obj=managed_files_obj, 

646 prisma_client=prisma_client, 

647 verbose_proxy_logger=verbose_proxy_logger, 

648 db_batch_object=db_batch_object, 

649 unified_batch_id=unified_batch_id, 

650 ) 

651 

652 asyncio.create_task( 

653 proxy_logging_obj.update_request_status( 

654 litellm_call_id=data.get("litellm_call_id", ""), status="success" 

655 ) 

656 ) 

657 

658 hidden_params = getattr(response, "_hidden_params", {}) or {} 

659 model_id = hidden_params.get("model_id", None) or "" 

660 cache_key = hidden_params.get("cache_key", None) or "" 

661 api_base = hidden_params.get("api_base", None) or "" 

662 

663 fastapi_response.headers.update( 

664 ProxyBaseLLMRequestProcessing.get_custom_headers( 

665 user_api_key_dict=user_api_key_dict, 

666 model_id=model_id, 

667 cache_key=cache_key, 

668 api_base=api_base, 

669 version=version, 

670 model_region=getattr(user_api_key_dict, "allowed_model_region", ""), 

671 request_data=data, 

672 ) 

673 ) 

674 

675 return response 

676 

677 # If batch is still processing, sync with provider to get latest state 

678 if response is not None: 678 ↛ 679line 678 didn't jump to line 679 because the condition on line 678 was never true

679 verbose_proxy_logger.debug( 

680 "Batch %s is in non-terminal state %s, syncing with provider", batch_id, response.status 

681 ) 

682 

683 poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active() 

684 if poller_owns_accounting: 684 ↛ 685line 684 didn't jump to line 685 because the condition on line 684 was never true

685 _litellm_metadata_of(data)["batch_ignore_default_logging"] = True 

686 

687 # Retrieve from provider (for non-terminal states or if DB lookup failed) 

688 # SCENARIO 1: Batch ID is encoded with model info 

689 if model_from_id is not None: 689 ↛ 690line 689 didn't jump to line 690 because the condition on line 689 was never true

690 credentials: Final = await get_authorized_credentials_for_model( 

691 llm_router=llm_router, 

692 model_id=model_from_id, 

693 user_api_key_dict=user_api_key_dict, 

694 operation_context="batch retrieval (batch created with model)", 

695 ) 

696 

697 original_batch_id: Final = get_original_file_id(batch_id) 

698 prepare_data_with_credentials( 

699 data=data, 

700 credentials=credentials, 

701 file_id=original_batch_id, # Sets data["batch_id"] = original_batch_id 

702 ) 

703 # Fix: The helper sets "file_id" but we need "batch_id" 

704 data["batch_id"] = data.pop("file_id", original_batch_id) 

705 # Provider-config providers (e.g. bedrock) require `model` in kwargs 

706 # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without 

707 # it the call falls into the legacy provider switch and 400s. 

708 data["model"] = model_from_id 

709 add_deployment_model_info(data=data, llm_router=llm_router, model_id=model_from_id) 

710 _litellm_metadata_of(data).setdefault("model_group", model_from_id) 

711 

712 # Retrieve batch using model credentials 

713 response = await litellm.aretrieve_batch( 

714 custom_llm_provider=credentials["custom_llm_provider"], 

715 **data, 

716 ) 

717 

718 encode_batch_response_ids(response, model=model_from_id) 

719 

720 verbose_proxy_logger.debug( 

721 "Retrieved batch using model: %s, original_id: %s", model_from_id, original_batch_id 

722 ) 

723 

724 elif litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id: 724 ↛ 725line 724 didn't jump to line 725 because the condition on line 724 was never true

725 if llm_router is None: 

726 raise HTTPException( 

727 status_code=500, 

728 detail={"error": "LLM Router not initialized. Ensure models added to proxy."}, 

729 ) 

730 

731 if unified_batch_id: 

732 add_internal_model_credentials( 

733 data=data, 

734 llm_router=llm_router, 

735 model_id=unified_model_id, 

736 ) 

737 

738 response = await llm_router.aretrieve_batch(**data) 

739 response._hidden_params["unified_batch_id"] = unified_batch_id 

740 if unified_batch_id: 

741 model_id_from_batch: Final = get_model_id_from_unified_batch_id(unified_batch_id) 

742 if model_id_from_batch: 

743 response._hidden_params["model_id"] = model_id_from_batch 

744 

745 # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) 

746 else: 

747 requested_provider: Final = ( 

748 provider 

749 or get_custom_llm_provider_from_request_headers(request=request) 

750 or get_custom_llm_provider_from_request_query(request=request) 

751 ) 

752 custom_llm_provider: Final = requested_provider or "openai" 

753 apply_team_provider_credentials( 

754 data=data, 

755 llm_router=llm_router, 

756 user_api_key_dict=user_api_key_dict, 

757 custom_llm_provider=custom_llm_provider, 

758 ) 

759 _raise_not_found_when_openai_fallback_unservable( 

760 requested_provider=requested_provider, 

761 data=data, 

762 not_found_message=f"No batch found with id '{batch_id}'.", 

763 ) 

764 response = await litellm.aretrieve_batch( 

765 custom_llm_provider=custom_llm_provider, 

766 **data, 

767 ) 

768 

769 # FIX: Update the database with the latest state from provider 

770 await update_batch_in_database( 

771 batch_id=batch_id, 

772 unified_batch_id=unified_batch_id, 

773 response=response, 

774 managed_files_obj=managed_files_obj, 

775 prisma_client=prisma_client, 

776 verbose_proxy_logger=verbose_proxy_logger, 

777 db_batch_object=db_batch_object, 

778 operation="retrieve", 

779 poller_owns_accounting=poller_owns_accounting, 

780 ) 

781 

782 ### CALL HOOKS ### - modify outgoing data 

783 response = await proxy_logging_obj.post_call_success_hook( 

784 data=data, user_api_key_dict=user_api_key_dict, response=response 

785 ) 

786 

787 # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id 

788 # Register any missing managed-file rows and return unified IDs. 

789 if unified_batch_id: 

790 await ensure_batch_response_managed_file_ids( 

791 response=response, 

792 managed_files_obj=managed_files_obj, 

793 prisma_client=prisma_client, 

794 verbose_proxy_logger=verbose_proxy_logger, 

795 db_batch_object=db_batch_object, 

796 unified_batch_id=unified_batch_id, 

797 ) 

798 

799 ### ALERTING ### 

800 asyncio.create_task( 

801 proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success") 

802 ) 

803 

804 ### RESPONSE HEADERS ### 

805 hidden_params = getattr(response, "_hidden_params", {}) or {} 

806 model_id = hidden_params.get("model_id", None) or "" 

807 cache_key = hidden_params.get("cache_key", None) or "" 

808 api_base = hidden_params.get("api_base", None) or "" 

809 

810 fastapi_response.headers.update( 

811 ProxyBaseLLMRequestProcessing.get_custom_headers( 

812 user_api_key_dict=user_api_key_dict, 

813 model_id=model_id, 

814 cache_key=cache_key, 

815 api_base=api_base, 

816 version=version, 

817 model_region=getattr(user_api_key_dict, "allowed_model_region", ""), 

818 request_data=data, 

819 ) 

820 ) 

821 

822 return response 

823 except Exception as e: 

824 await proxy_logging_obj.post_call_failure_hook( 

825 user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data 

826 ) 

827 litellm_call_id: Final = request_litellm_call_id(data) 

828 log_llm_api_exception(e, litellm_call_id) 

829 raise handle_exception_on_proxy(e, litellm_call_id) 

830 

831 

832@router.get( 

833 "/{provider}/v1/batches", 

834 dependencies=[Depends(user_api_key_auth)], 

835 tags=["batch"], 

836) 

837@router.get( 

838 "/v1/batches", 

839 dependencies=[Depends(user_api_key_auth)], 

840 tags=["batch"], 

841) 

842@router.get( 

843 "/batches", 

844 dependencies=[Depends(user_api_key_auth)], 

845 tags=["batch"], 

846) 

847async def list_batches( 

848 request: Request, 

849 fastapi_response: Response, 

850 provider: str | None = None, 

851 limit: int | None = None, 

852 after: str | None = None, 

853 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 

854 target_model_names: str | None = None, 

855): 

856 """ 

857 Lists  

858 This is the equivalent of GET https://api.openai.com/v1/batches/ 

859 Supports Identical Params as: https://platform.openai.com/docs/api-reference/batch/list 

860 

861 Example Curl 

862 ``` 

863 curl http://localhost:4000/v1/batches?limit=2 \ 

864 -H "Authorization: Bearer sk-1234" \