From d4b0a77698cef9eee9c0bb6243048e5d38773553 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:45:24 +0000 Subject: [PATCH] fix(proxy): register bedrock batch output files with a file object so they list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 21 ++++++++++- .../proxy_unit_tests/test_check_batch_cost.py | 36 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 13e9e5093a8..d12d2cc3dcb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,6 +2,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ +import time from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType @@ -13,6 +14,7 @@ from litellm.constants import ( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) +from litellm.types.llms.openai import OpenAIFileObject if TYPE_CHECKING: from prisma import models as prisma_models @@ -41,6 +43,19 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( ) +def _batch_output_file_object(unified_file_id: str, raw_file_id: str, size_bytes: int) -> OpenAIFileObject: + filename: Final = raw_file_id.rsplit("/", 1)[-1] or raw_file_id + return OpenAIFileObject( + id=unified_file_id, + object="file", + purpose="batch_output", + filename=filename, + created_at=int(time.time()), + bytes=size_bytes, + status="processed", + ) + + class CheckBatchCost: def __init__( self, @@ -787,7 +802,11 @@ class CheckBatchCost: ) await managed_files_hook.store_unified_file_id( file_id=_unified_file_id, - file_object=None, + file_object=_batch_output_file_object( + unified_file_id=_unified_file_id, + raw_file_id=_raw_file_id, + size_bytes=len(content_bytes) if _file_attr == "output_file_id" else 0, + ), litellm_parent_otel_span=None, model_mappings={model_id: _raw_file_id}, user_api_key_dict=_minimal_auth, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6417c7c8aa6..ec716e05925 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -23,6 +23,30 @@ _CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA==" _CLAIM_OUTPUT_FILE_ID = "file-output-123" +def test_batch_output_file_object_derives_metadata(): + from litellm_enterprise.proxy.common_utils.check_batch_cost import _batch_output_file_object + + output_file = _batch_output_file_object( + unified_file_id="unified-output", + raw_file_id="s3://bucket/path/to/output.jsonl.out", + size_bytes=4321, + ) + provider_file = _batch_output_file_object( + unified_file_id="unified-provider", + raw_file_id="file-abc", + size_bytes=0, + ) + + assert output_file.id == "unified-output" + assert output_file.object == "file" + assert output_file.purpose == "batch_output" + assert output_file.filename == "output.jsonl.out" + assert output_file.bytes == 4321 + assert output_file.status == "processed" + assert provider_file.filename == "file-abc" + assert provider_file.bytes == 0 + + def _batch_cost_result( cost: float, usage: dict, @@ -1777,6 +1801,18 @@ class TestCheckBatchCost: raw_output_file_id: fake_managed_output_id, raw_error_file_id: fake_managed_error_id, } + for store_call in mock_hook.store_unified_file_id.call_args_list: + file_object = store_call.kwargs["file_object"] + assert file_object is not None + assert file_object.object == "file" + assert file_object.purpose == "batch_output" + assert file_object.id == store_call.kwargs["file_id"] + stored_file_objects = { + next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_object"] + for c in mock_hook.store_unified_file_id.call_args_list + } + assert stored_file_objects[raw_output_file_id].bytes == len(mock_file_content.content) + assert stored_file_objects[raw_error_file_id].bytes == 0 assert mock_response.output_file_id == fake_managed_output_id assert mock_response.error_file_id == fake_managed_error_id