Address Bedrock logging review feedback

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin 2026-05-11 21:41:47 +00:00
parent 14ea4f57f4
commit fc35360606
No known key found for this signature in database
2 changed files with 14 additions and 5 deletions

View file

@ -80,9 +80,7 @@ def _sanitize_response_for_logging(value: Any) -> Any:
def _default(obj: Any) -> Any:
if isinstance(obj, datetime):
return obj.isoformat()
raise TypeError(
f"Object of type {obj.__class__.__name__} is not JSON serializable"
)
return str(obj)
return json.loads(json.dumps(value, default=_default))

View file

@ -11,6 +11,7 @@ from __future__ import annotations
import json
import os
import sys
from decimal import Decimal
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
@ -112,6 +113,7 @@ def test_sanitize_response_for_logging_converts_nested_datetimes():
"nested": {
"endTimes": [END_TIME],
"tupleTimes": (SUBMIT_TIME,),
"cost": Decimal("1.25"),
},
}
@ -122,9 +124,11 @@ def test_sanitize_response_for_logging_converts_nested_datetimes():
"nested": {
"endTimes": [END_TIME.isoformat()],
"tupleTimes": [SUBMIT_TIME.isoformat()],
"cost": "1.25",
},
}
assert response["submitTime"] == SUBMIT_TIME
json.dumps(sanitized_response)
def test_extract_job_id_from_arn():
@ -329,14 +333,21 @@ def test_logging_obj_pre_and_post_call_invoked(patched_boto3):
def test_logging_obj_post_call_gets_json_serializable_response(patched_boto3):
class JsonSerializingLogger:
serialized_response: str
def pre_call(self, **kwargs):
pass
def post_call(self, **kwargs):
json.dumps(kwargs["original_response"])
self.serialized_response = json.dumps(kwargs["original_response"])
logging_obj = JsonSerializingLogger()
BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=JOB_ARN, logging_obj=JsonSerializingLogger()
batch_id=JOB_ARN, logging_obj=logging_obj
)
assert (
f'"submitTime": "{SUBMIT_TIME.isoformat()}"' in logging_obj.serialized_response
)