Skip to content

Commit b99b935

Browse files
fix(check_batch_cost): persist managed file before mutating response and propagate team_id
- Move setattr after store_unified_file_id so the response only receives the managed ID once the DB record is successfully written. Avoids serializing an orphaned managed ID into file_object when the store call fails. - Populate team_id on the minimal UserAPIKeyAuth from job.team_id so the managed file record is created with the correct team ownership, allowing other team members to access the batch output file via /files/{id}/content. Co-authored-by: Yassin Kortam <yassin@berri.ai>
1 parent 13d34a1 commit b99b935

1 file changed

Lines changed: 3 additions & 2 deletions

File tree

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,8 @@ async def check_batch_cost(self):
306306
if managed_files_hook is not None:
307307
from litellm.proxy._types import UserAPIKeyAuth
308308
_minimal_auth = UserAPIKeyAuth(
309-
user_id=job.created_by or "default-user-id"
309+
user_id=job.created_by or "default-user-id",
310+
team_id=getattr(job, "team_id", None),
310311
)
311312
for _file_attr in ["output_file_id", "error_file_id"]:
312313
_raw_file_id = getattr(response, _file_attr, None)
@@ -317,14 +318,14 @@ async def check_batch_cost(self):
317318
model_id=model_id,
318319
model_name=str(model_name) if model_name else deployment_info.model_name or None,
319320
)
320-
setattr(response, _file_attr, _unified_file_id)
321321
await managed_files_hook.store_unified_file_id(
322322
file_id=_unified_file_id,
323323
file_object=None,
324324
litellm_parent_otel_span=None,
325325
model_mappings={model_id: _raw_file_id},
326326
user_api_key_dict=_minimal_auth,
327327
)
328+
setattr(response, _file_attr, _unified_file_id)
328329
verbose_proxy_logger.info(
329330
f"CheckBatchCost: converted {_file_attr} "
330331
f"{_raw_file_id!r} -> managed ID for batch {batch_id}"

0 commit comments

Comments
 (0)