From 608d7499836c1aaa7fe752c473c68d5813b9f29b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:46:59 -0700 Subject: [PATCH 1/3] fix(batches): stop one bad output line from zeroing an entire batch's spend --- litellm/batches/batch_utils.py | 161 ++++++++++++------ .../test_litellm/batches/test_batch_utils.py | 40 +++-- .../proxy/hooks/test_batch_file_validation.py | 12 +- type-discipline-budget.json | 2 +- 4 files changed, 149 insertions(+), 66 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index c2cbb9604e5..feb84ccd8a6 100644 --- a/litellm/batches/batch_utils.py +++ b/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, + 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, + 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, + 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 json.JSONDecodeError 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,7 @@ 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, custom_llm_provider: str = "openai") -> Usage: """ Get the tokens of a batch job from the response body """ @@ -472,7 +535,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) -> dict: """ Get the ``result`` object from a line of an Anthropic message batch results JSONL file. @@ -482,7 +545,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, custom_llm_provider: str = "openai" +) -> Any: """ Get the response from the batch job output file """ @@ -495,7 +560,7 @@ 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, 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 573882ebfca..254d663af93 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/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,17 @@ 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}] # =========================================================================== # @@ -471,6 +469,26 @@ 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(): + # 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. + 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 diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 1ce1a2f3e51..28624254565 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/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: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f8e481dc142..43753224714 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22894 + "limit": 22891 }, "LIT002": { "limit": 26888 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 2/3] 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)), { From 5eeccf69b60b15e1f0939371d3ee18c0eb2099d4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 14:35:01 -0700 Subject: [PATCH 3/3] fix(batches): skip undecodable batch output lines when costing --- litellm/batches/batch_utils.py | 2 +- tests/test_litellm/batches/test_batch_utils.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 0f8acce3379..0cf22d82ca6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -426,7 +426,7 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]: def _parse_batch_output_line(line: bytes) -> dict | None: try: parsed: Final = json.loads(line) - except json.JSONDecodeError as e: + except ValueError as e: verbose_logger.warning("skipping malformed batch output line: %s", str(e)) return None if isinstance(parsed, dict): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index f3a7413bf03..ebe093c591c 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -186,6 +186,11 @@ def test_iter_output_entries_skips_malformed_and_non_object_lines(): 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}] + + # =========================================================================== # # _estimate_batch_entry_tokens (regression: an uncountable/malformed row must # never contribute zero tokens, or a crafted batch could evade the TPM limit)