From 4f587087ffbd90d79d15b7239517ac2e4e25c2d6 Mon Sep 17 00:00:00 2001 From: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com> Date: Mon, 11 May 2026 21:28:37 +0000 Subject: [PATCH] Fix Bedrock batch datetime logging Co-authored-by: ishaan-berri --- litellm/llms/bedrock/batches/handler.py | 15 +++++++- .../llms/bedrock/batches/test_handler.py | 37 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index c071f331337..554fa3c9974 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -75,6 +75,19 @@ def _to_epoch(value: Any) -> Optional[int]: return None +def _sanitize_response_for_logging(value: Any) -> Any: + if isinstance(value, datetime): + return value.isoformat() + if isinstance(value, dict): + return { + key: _sanitize_response_for_logging(response_value) + for key, response_value in value.items() + } + if isinstance(value, (list, tuple)): + return [_sanitize_response_for_logging(item) for item in value] + return value + + class BedrockBatchesHandler: """ Handler for Bedrock Batches. @@ -270,7 +283,7 @@ class BedrockBatchesHandler: logging_obj.post_call( input=batch_id, api_key="", - original_response=response, + original_response=_sanitize_response_for_logging(response), additional_args={"complete_input_dict": {"jobIdentifier": batch_id}}, ) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 18780ccce0f..8478be2302f 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,6 +8,7 @@ the tests don't hit AWS. from __future__ import annotations +import json import os import sys from datetime import datetime, timezone @@ -22,6 +23,7 @@ from litellm.llms.bedrock.batches.handler import ( # noqa: E402 _extract_job_id_from_arn, _extract_region_from_bedrock_arn, _predict_output_file_uri, + _sanitize_response_for_logging, _to_epoch, ) @@ -104,6 +106,27 @@ def test_to_epoch_handles_supported_types(value, expected): assert _to_epoch(value) == expected +def test_sanitize_response_for_logging_converts_nested_datetimes(): + response = { + "submitTime": SUBMIT_TIME, + "nested": { + "endTimes": [END_TIME], + "tupleTimes": (SUBMIT_TIME,), + }, + } + + sanitized_response = _sanitize_response_for_logging(response) + + assert sanitized_response == { + "submitTime": SUBMIT_TIME.isoformat(), + "nested": { + "endTimes": [END_TIME.isoformat()], + "tupleTimes": [SUBMIT_TIME.isoformat()], + }, + } + assert response["submitTime"] == SUBMIT_TIME + + def test_extract_job_id_from_arn(): assert _extract_job_id_from_arn(JOB_ARN) == JOB_ID assert ( @@ -301,6 +324,20 @@ def test_logging_obj_pre_and_post_call_invoked(patched_boto3): post_kwargs = logging_obj.post_call.call_args.kwargs assert post_kwargs["input"] == JOB_ARN assert post_kwargs["original_response"]["jobArn"] == JOB_ARN + assert post_kwargs["original_response"]["submitTime"] == SUBMIT_TIME.isoformat() + + +def test_logging_obj_post_call_gets_json_serializable_response(patched_boto3): + class JsonSerializingLogger: + def pre_call(self, **kwargs): + pass + + def post_call(self, **kwargs): + json.dumps(kwargs["original_response"]) + + BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, logging_obj=JsonSerializingLogger() + ) def test_missing_boto3_raises_helpful_import_error():