scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

refactor(batches): parameterize the batch output Mapping annotations

Browse source
This commit is contained in:
mateo-berri 2026-08-19 02:15:22 -07:00
parent 608d749983
commit 5c6391d7d2
2 changed files with 13 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/batches/batch_utils.py
View file

@ -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

5
tests/test_litellm/batches/test_batch_utils.py
View file

@ -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)),
{