From 5c6391d7d2590f19e6cd5bc2f5d89c5aa4ffe5f1 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 19 Aug 2026 02:15:22 -0700
Subject: [PATCH] refactor(batches): parameterize the batch output Mapping
 annotations

---
 litellm/batches/batch_utils.py                 | 18 +++++++++++-------
 tests/test_litellm/batches/test_batch_utils.py |  5 ++---
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py
index feb84ccd8a6..0f8acce3379 100644
--- a/litellm/batches/batch_utils.py
+++ b/litellm/batches/batch_utils.py
@@ -118,7 +118,7 @@ def _iter_successful_output_line_stats(
 
 
 def _safe_output_line_stats(
-    entry: Mapping,
+    entry: Mapping[str, Any],
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
     model_name: str | None,
     model_info: ModelInfo | None,
@@ -142,7 +142,7 @@ def _safe_output_line_stats(
 
 
 def _compute_output_line_stats(
-    entry: Mapping,
+    entry: Mapping[str, Any],
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
     model_name: str | None,
     model_info: ModelInfo | None,
@@ -171,7 +171,7 @@ def _compute_output_line_stats(
 
 
 def _output_line_cost(
-    response_body: Mapping,
+    response_body: Mapping[str, Any],
     usage: Usage,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
     model_name: str | None,
@@ -503,7 +503,9 @@ def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
     return 0
 
 
-def _get_batch_job_usage_from_response_body(response_body: Mapping, custom_llm_provider: str = "openai") -> Usage:
+def _get_batch_job_usage_from_response_body(
+    response_body: Mapping[str, Any], custom_llm_provider: str = "openai"
+) -> Usage:
     """
     Get the tokens of a batch job from the response body
     """
@@ -535,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: Mapping, custom_llm_p
     return usage
 
 
-def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping) -> dict:
+def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
     """
     Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
 
@@ -546,7 +548,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping) -
 
 
 def _get_response_from_batch_job_output_file(
-    batch_job_output_file: Mapping, custom_llm_provider: str = "openai"
+    batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
 ) -> Any:
     """
     Get the response from the batch job output file
@@ -560,7 +562,9 @@ def _get_response_from_batch_job_output_file(
     return _response_body
 
 
-def _batch_response_was_successful(batch_job_output_file: Mapping, custom_llm_provider: str = "openai") -> bool:
+def _batch_response_was_successful(
+    batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
+) -> bool:
     """
     Check if the batch job response was successful
 
diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py
index 254d663af93..a30420e6224 100644
--- a/tests/test_litellm/batches/test_batch_utils.py
+++ b/tests/test_litellm/batches/test_batch_utils.py
@@ -470,9 +470,8 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
 
 
 def test_empty_body_line_does_not_zero_whole_batch():
-    # Regression: a status-200 row with an empty body made the real
-    # litellm.completion_cost raise ValueError, aborting the aggregation so the
-    # entire batch was booked at $0. The bad line must be skipped instead.
+    """A status-200 row with an empty body makes litellm.completion_cost raise;
+    that line must be skipped instead of zeroing the whole batch."""
     rows = [
         _success_row(usage=_usage(10, 5)),
         {