Merge pull request #37457 from BerriAI/litellm_batch_empty_line_cost

fix(batches): stop one bad output line from zeroing an entire batch's spend

litellm/batches/batch_utils.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import Iterable, Iterator
+from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import dataclass
 from typing import Any, Final, Literal
 
@@ -87,7 +87,7 @@ async def _handle_completed_batch(
         return batch_cost, batch_usage, [model_name]
 
     return _aggregate_batch_cost_usage_models(
-        entries=_iter_batch_input_entries(file_content),
+        entries=_iter_batch_output_entries(file_content),
         custom_llm_provider=custom_llm_provider,
         model_name=model_name,
         model_info=model_info,
@@ -111,43 +111,91 @@ def _iter_successful_output_line_stats(
     model_name: str | None,
     model_info: ModelInfo | None,
 ) -> Iterator[_BatchOutputLineStats]:
+    for entry in entries:
+        stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
+        if stats is not None:
+            yield stats
+
+
+def _safe_output_line_stats(
+    entry: Mapping[str, Any],
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
+    model_name: str | None,
+    model_info: ModelInfo | None,
+) -> _BatchOutputLineStats | None:
+    """Return the stats for one batch output line, or None for a line that is
+    unsuccessful or cannot be costed, so a single bad line never aborts the
+    whole batch's cost accounting."""
+    custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
+    try:
+        if not _batch_response_was_successful(entry, custom_llm_provider):
+            return None
+        return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
+    except Exception as e:  # noqa: BLE001  # any single line's costing failure must not abort the whole batch
+        verbose_logger.warning(
+            "batch output line could not be costed, so it is billed at $0 and the rest of the batch "
+            "is still billed. custom_id=%s error=%s",
+            custom_id,
+            str(e),
+        )
+        return None
+
+
+def _compute_output_line_stats(
+    entry: Mapping[str, Any],
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
+    model_name: str | None,
+    model_info: ModelInfo | None,
+) -> _BatchOutputLineStats:
+    response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
+    usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
+    prompt_details: Final = parse_prompt_tokens_details(usage)
+    raw_model: Final = response_body.get("model")
+    response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
+    return _BatchOutputLineStats(
+        cost=_output_line_cost(
+            response_body=response_body,
+            usage=usage,
+            custom_llm_provider=custom_llm_provider,
+            model_name=model_name,
+            response_model=response_model,
+            model_info=model_info,
+        ),
+        prompt_tokens=usage.prompt_tokens,
+        completion_tokens=usage.completion_tokens,
+        total_tokens=usage.total_tokens,
+        cache_read_tokens=prompt_details["cache_hit_tokens"],
+        cache_creation_tokens=prompt_details["cache_creation_tokens"],
+        model=response_model,
+    )
+
+
+def _output_line_cost(
+    response_body: Mapping[str, Any],
+    usage: Usage,
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
+    model_name: str | None,
+    response_model: str | None,
+    model_info: ModelInfo | None,
+) -> float:
     from litellm.cost_calculator import batch_cost_calculator
 
-    for entry in entries:
-        if not _batch_response_was_successful(entry, custom_llm_provider):
-            continue
-        response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
-        usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
-        prompt_details = parse_prompt_tokens_details(usage)
-        raw_model = response_body.get("model")
-        response_model = raw_model if isinstance(raw_model, str) and raw_model else None
-        if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
-            if custom_llm_provider == "bedrock" and model_name:
-                cost_model = model_name
-            else:
-                cost_model = response_model or model_name or ""
-            prompt_cost, completion_cost = batch_cost_calculator(
-                usage=usage,
-                model=cost_model,
-                custom_llm_provider=custom_llm_provider,
-                model_info=model_info,
-            )
-            line_cost = prompt_cost + completion_cost
-        else:
-            line_cost = litellm.completion_cost(
-                completion_response=response_body,
-                custom_llm_provider=custom_llm_provider,
-                call_type=CallTypes.aretrieve_batch.value,
-            )
-        yield _BatchOutputLineStats(
-            cost=line_cost,
-            prompt_tokens=usage.prompt_tokens,
-            completion_tokens=usage.completion_tokens,
-            total_tokens=usage.total_tokens,
-            cache_read_tokens=prompt_details["cache_hit_tokens"],
-            cache_creation_tokens=prompt_details["cache_creation_tokens"],
-            model=response_model,
+    if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
+        return litellm.completion_cost(
+            completion_response=response_body,
+            custom_llm_provider=custom_llm_provider,
+            call_type=CallTypes.aretrieve_batch.value,
         )
+    cost_model: Final = (
+        model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
+    )
+    prompt_cost, completion_cost = batch_cost_calculator(
+        usage=usage,
+        model=cost_model,
+        custom_llm_provider=custom_llm_provider,
+        model_info=model_info,
+    )
+    return prompt_cost + completion_cost
 
 
 def _aggregate_batch_cost_usage_models(
@@ -338,9 +386,10 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
 
 def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
     """
-    Get the file content as a list of dictionaries from JSON Lines format
+    Get the file content as a list of dictionaries from JSON Lines format,
+    skipping malformed lines
     """
-    return list(_iter_batch_input_entries(file_content))
+    return list(_iter_batch_output_entries(file_content))
 
 
 def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
@@ -361,15 +410,29 @@ def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
             yield line
 
 
-def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]:
+def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]:
     """
-    Yield parsed batch input JSONL entries one at a time without materializing the
-    whole file as a list, so peak memory stays bounded. Raises on a malformed line;
-    callers that must survive bad rows should iterate ``_iter_batch_input_lines``
-    and parse per-row instead.
+    Yield parsed batch output JSONL entries one at a time without materializing
+    the whole file as a list, so peak memory stays bounded. A malformed or
+    non-object line is skipped with a warning so one bad line never aborts the
+    whole batch's cost accounting.
     """
     for line in _iter_batch_input_lines(file_content):
-        yield json.loads(line)
+        entry = _parse_batch_output_line(line)
+        if entry is not None:
+            yield entry
+
+
+def _parse_batch_output_line(line: bytes) -> dict | None:
+    try:
+        parsed: Final = json.loads(line)
+    except ValueError as e:
+        verbose_logger.warning("skipping malformed batch output line: %s", str(e))
+        return None
+    if isinstance(parsed, dict):
+        return parsed
+    verbose_logger.warning("skipping non-object batch output line of type %s", type(parsed).__name__)
+    return None
 
 
 # A batch request's input tokens scale roughly with its serialized size, so this
@@ -440,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: dict, 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
     """
@@ -472,7 +537,7 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
     return usage
 
 
-def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> 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.
 
@@ -482,7 +547,9 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: dict) -> d
     return batch_results_line.get("result", None) or {}
 
 
-def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom_llm_provider: str = "openai") -> Any:
+def _get_response_from_batch_job_output_file(
+    batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
+) -> Any:
     """
     Get the response from the batch job output file
     """
@@ -495,7 +562,9 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
     return _response_body
 
 
-def _batch_response_was_successful(batch_job_output_file: dict, 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
 

tests/test_litellm/batches/test_batch_utils.py

@@ -150,13 +150,13 @@ def test_parse_jsonl_empty_content_is_empty_list():
     assert bu._get_file_content_as_dictionary(b"") == []
 
 
-def test_parse_jsonl_malformed_raises():
-    with pytest.raises(Exception):
-        bu._get_file_content_as_dictionary(b"not valid json")
+def test_parse_jsonl_malformed_lines_skipped():
+    content = b'{"a": 1}\nnot valid json\n{"b": 2}\n'
+    assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]
 
 
 # =========================================================================== #
-# _iter_batch_input_lines / _iter_batch_input_entries  (JSONL parsing)
+# _iter_batch_input_lines / _iter_batch_output_entries  (JSONL parsing)
 # =========================================================================== #
 
 
@@ -173,19 +173,22 @@ def test_iter_input_lines_empty():
     assert list(bu._iter_batch_input_lines(b"")) == []
 
 
-def test_iter_input_entries_parses_each_row():
+def test_iter_output_entries_parses_each_row():
     content = b'{"body": {"model": "gpt-4o"}}\n{"body": {"model": "claude-3"}}\n'
-    assert list(bu._iter_batch_input_entries(content)) == [
+    assert list(bu._iter_batch_output_entries(content)) == [
         {"body": {"model": "gpt-4o"}},
         {"body": {"model": "claude-3"}},
     ]
 
 
-def test_iter_input_entries_raises_on_malformed_line():
-    # _iter_batch_input_entries raises on a bad row; callers that must survive
-    # bad rows iterate _iter_batch_input_lines and parse per-row instead.
-    with pytest.raises(Exception):
-        list(bu._iter_batch_input_entries(b'{"ok":1}\nnot-json\n'))
+def test_iter_output_entries_skips_malformed_and_non_object_lines():
+    content = b'{"ok": 1}\nnot-json\n[1, 2]\n{"ok": 2}\n'
+    assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
+
+
+def test_iter_output_entries_skips_undecodable_line():
+    content = b'{"ok": 1}\n{"note": "\xff-bad"}\n{"ok": 2}\n'
+    assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
 
 
 # =========================================================================== #
@@ -471,6 +474,25 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
     assert len(calls) == 2  # failed row not costed
 
 
+def test_empty_body_line_does_not_zero_whole_batch():
+    """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)),
+        {
+            "custom_id": "request-poison-empty",
+            "response": {"status_code": 200, "request_id": "inject-empty-body", "body": {}},
+        },
+        _success_row(usage=_usage(20, 10)),
+    ]
+
+    cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
+
+    assert cost > 0.0
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
+    assert models == ["gpt-4o", "gpt-4o"]
+
+
 def test_cost_from_content_model_info_path(monkeypatch):
     # model_info set -> batch_cost_calculator(prompt_cost, completion_cost).
     import litellm.cost_calculator as cc

tests/test_litellm/proxy/hooks/test_batch_file_validation.py

@@ -1739,18 +1739,18 @@ def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes:
     return ("\n".join(rows)).encode("utf-8")
 
 
-def test_iter_batch_input_entries_matches_dict_list():
+def test_iter_batch_output_entries_matches_dict_list():
     from litellm.batches.batch_utils import (
         _get_file_content_as_dictionary,
-        _iter_batch_input_entries,
+        _iter_batch_output_entries,
     )
 
     raw = _make_batch_input_bytes(50)
-    streamed = list(_iter_batch_input_entries(raw))
+    streamed = list(_iter_batch_output_entries(raw))
     assert streamed == _get_file_content_as_dictionary(raw)
     assert streamed[0]["custom_id"] == "request-0"
     # tolerant of blank lines and a missing trailing newline
-    assert list(_iter_batch_input_entries(raw + b"\n\n")) == streamed
+    assert list(_iter_batch_output_entries(raw + b"\n\n")) == streamed
 
 
 def test_streaming_count_peak_below_dict_list():
@@ -1759,7 +1759,7 @@ def test_streaming_count_peak_below_dict_list():
 
     from litellm.batches.batch_utils import (
         _get_file_content_as_dictionary,
-        _iter_batch_input_entries,
+        _iter_batch_output_entries,
     )
 
     raw = _make_batch_input_bytes(8000)
@@ -1777,7 +1777,7 @@ def test_streaming_count_peak_below_dict_list():
     def _stream():
         count = 0
         models: set = set()
-        for entry in _iter_batch_input_entries(raw):
+        for entry in _iter_batch_output_entries(raw):
             count += 1
             model = (entry.get("body") or {}).get("model")
             if model:

type-discipline-budget.json

@@ -1,6 +1,6 @@
 {
   "LIT001": {
-    "limit": 22809
+    "limit": 22806
   },
   "LIT002": {
     "limit": 26878