Fix Bedrock batch datetime logging

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

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}},
             )
 

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():