fix(bedrock): treat partial record counts as unknown on batch retrieve

litellm/llms/bedrock/batches/handler.py

@@ -71,13 +71,13 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N
 
 def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequestCounts | None:
     total_records: Final = response.get("totalRecordCount")
-    if not isinstance(total_records, int):
-        return None
     success_records: Final = response.get("successRecordCount")
+    if not isinstance(total_records, int) or not isinstance(success_records, int):
+        return None
     error_records: Final = response.get("errorRecordCount")
     return BatchRequestCounts(
         total=total_records,
-        completed=success_records if isinstance(success_records, int) else 0,
+        completed=success_records,
         failed=error_records if isinstance(error_records, int) else 0,
     )
 

tests/test_litellm/llms/bedrock/batches/test_handler.py

@@ -159,9 +159,12 @@ def test_handle_model_invocation_job_status_completed(patched_boto3):
 @pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)])
 def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count):
     fake_client, _ = patched_boto3
-    counted_response = _fake_boto3_response()
-    counted_response.update(totalRecordCount=100, successRecordCount=success_count, errorRecordCount=error_count)
-    fake_client.get_model_invocation_job.return_value = counted_response
+    fake_client.get_model_invocation_job.return_value = {
+        **_fake_boto3_response(),
+        "totalRecordCount": 100,
+        "successRecordCount": success_count,
+        "errorRecordCount": error_count,
+    }
 
     batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
 
@@ -182,6 +185,29 @@ def test_missing_record_counts_leave_request_counts_none(patched_boto3):
     assert batch.request_counts is None
 
 
+def test_total_without_success_count_leaves_request_counts_none(patched_boto3):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "totalRecordCount": 100}
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.request_counts is None
+
+
+def test_missing_error_count_maps_to_zero_failed(patched_boto3):
+    fake_client, _ = patched_boto3
+    fake_client.get_model_invocation_job.return_value = {
+        **_fake_boto3_response(),
+        "totalRecordCount": 100,
+        "successRecordCount": 100,
+    }
+
+    batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN)
+
+    assert batch.request_counts is not None
+    assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == (100, 100, 0)
+
+
 @pytest.mark.parametrize(
     "bedrock_status,openai_status",
     [