mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Address Bedrock logging review feedback
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
parent
14ea4f57f4
commit
fc35360606
2 changed files with 14 additions and 5 deletions
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue