scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

perf(batches): reuse fetched batch result files for line-item logging

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-30 23:35:59 +00:00
parent a891065868
commit ff77170749
9 changed files with 433 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -705,9 +705,10 @@ class CheckBatchCost:
later poll.
"""
from litellm.batches.batch_utils import (
count_error_file_failed_requests,
_get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
count_error_file_failed_requests_from_content,
fetch_batch_error_file_content,
)
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
@ -850,7 +851,7 @@ class CheckBatchCost:
model_name=model_name,
model_info=deployment_model_info,
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
error_file_content = await fetch_batch_error_file_content(
response,
custom_llm_provider=batch_file_provider,
litellm_params={
@ -858,6 +859,7 @@ class CheckBatchCost:
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
},
)
error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content)
batch_result: Final = (
output_file_result
if not error_file_failed_requests
@ -916,6 +918,8 @@ class CheckBatchCost:
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
batch_output_file_content=content_bytes,
batch_error_file_content=error_file_content,
)
except Exception:
await self._release_job_claim(job)

18
litellm/batches/batch_line_item_logging.py
View file

@ -9,6 +9,7 @@ from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.batches.batch_utils import (
BatchResultFiles,
_batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design
_fetch_batch_managed_file_content, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream
_get_response_from_batch_job_output_file, # pyright: ignore[reportPrivateUsage] # same reuse
@ -422,6 +423,7 @@ async def log_batch_line_items(
model_name: str | None,
litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict
model_info: ModelInfo | None,
result_files: BatchResultFiles | None = None,
claim_cache: DualCache = batch_line_item_claim_cache,
) -> int:
"""Emit one callback event per JSONL line of a completed batch (request
@ -429,7 +431,9 @@ async def log_batch_line_items(
``litellm.store_batch_line_items_in_callbacks`` flag. The aggregate
aretrieve_batch event still bills the batch, so per-line events carry
``batch_parent_id`` and never update spend themselves. Any failure here
is logged and swallowed: aggregate accounting must be unaffected."""
is logged and swallowed: aggregate accounting must be unaffected.
``result_files`` carries the output/error bytes the aggregate path already
fetched, so they are reused instead of refetched."""
line_provider: Final = _supported_line_provider(custom_llm_provider)
if line_provider is None:
verbose_logger.warning(
@ -468,7 +472,11 @@ async def log_batch_line_items(
input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params)
requests_by_id: Final = _requests_by_custom_id(input_file_content)
output_content: Final = await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params)
output_content: Final = (
result_files.output
if result_files is not None and result_files.output is not None
else await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params)
)
first_row: Final = next(_output_entries(output_content), None)
if _uses_native_vertex_output(line_provider, model_name, first_row):
verbose_logger.warning(
@ -476,7 +484,11 @@ async def log_batch_line_items(
batch.id,
)
return 0
error_content: Final = await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params)
error_content: Final = (
result_files.error
if result_files is not None and result_files.error is not None
else await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params)
)
for content in (output_content, error_content):
for entry in _output_entries(content):
try:

104
litellm/batches/batch_utils.py
View file

@ -33,6 +33,14 @@ class BatchCostUsageResult:
completion_cost: float = 0.0
@dataclass(frozen=True, slots=True)
class BatchResultFiles:
"""Raw JSONL bytes of a completed batch's result files; None means not fetched."""
output: bytes | None
error: bytes | None
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
@ -115,29 +123,44 @@ async def _handle_completed_batch(
threaded through so a deployment's configured rates win over the
global cost map.
"""
# A completed batch whose request lines all failed has no output file - the
# results are written to a separate error_file_id and output_file_id is None.
# There is nothing to price or measure, so report an empty result set instead
# of calling _fetch_batch_output_file_content, which raises on a missing
# output file. Without this guard the logging worker crashes on every
# aretrieve_batch poll and the completed batch's zero-cost accounting is lost.
# The generic retrieval helper keeps raising for callers that explicitly ask
# for a missing output file.
return (
await _handle_completed_batch_with_files(
batch,
custom_llm_provider,
model_name=model_name,
litellm_params=litellm_params,
model_info=model_info,
)
)[0]
async def _handle_completed_batch_with_files(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None = None,
litellm_params: dict | None = None,
model_info: ModelInfo | None = None,
) -> "tuple[BatchCostUsageResult, BatchResultFiles]":
"""_handle_completed_batch plus the raw output/error bytes it fetched, so
downstream line-item logging reuses them instead of refetching."""
error_file_content: Final = await fetch_batch_error_file_content(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content)
if batch.output_file_id is None:
return BatchCostUsageResult(
cost=0.0,
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str]
successful_requests=0,
failed_requests=await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
return (
BatchCostUsageResult(
cost=0.0,
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
models=[], # mutable-ok: same contract as _handle_completed_batch above
successful_requests=0,
failed_requests=error_file_failed_requests,
),
BatchResultFiles(output=None, error=error_file_content),
)
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
error_file_failed_requests: Final = await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
output_file_result: Final = (
calculate_vertex_ai_batch_cost_and_usage(
@ -154,11 +177,14 @@ async def _handle_completed_batch(
)
)
if not error_file_failed_requests:
return output_file_result
return dataclasses_replace(
output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests
result: Final = (
output_file_result
if not error_file_failed_requests
else dataclasses_replace(
output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests
)
)
return result, BatchResultFiles(output=file_content, error=error_file_content)
class _LineOutcome(Enum):
@ -488,6 +514,27 @@ async def _fetch_batch_output_file_content(
)
async def fetch_batch_error_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
litellm_params: dict | None,
) -> bytes | None:
"""Fetch the batch's separate error file bytes; None when it has none or the fetch fails."""
if batch.error_file_id is None:
return None
try:
return await _fetch_batch_managed_file_content(
batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch
verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e)
return None
def count_error_file_failed_requests_from_content(content: bytes | None) -> int:
return 0 if content is None else sum(1 for _ in _iter_batch_input_lines(content))
async def count_error_file_failed_requests(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
@ -500,16 +547,11 @@ async def count_error_file_failed_requests(
``error_file_id`` - they never appear in the output file at all, so
counting failures from the output file alone silently undercounts them.
"""
if batch.error_file_id is None:
return 0
try:
error_file_content = await _fetch_batch_managed_file_content(
batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
return count_error_file_failed_requests_from_content(
await fetch_batch_error_file_content(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch
verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e)
return 0
return sum(1 for _ in _iter_batch_input_lines(error_file_content))
)
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:

40
litellm/litellm_core_utils/litellm_logging.py
View file

@ -32,7 +32,11 @@ from litellm._logging import (
verbose_logger,
)
from litellm._uuid import uuid
from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final
from litellm.batches.batch_utils import (
BatchResultFiles,
_handle_completed_batch_with_files,
batch_cost_is_final,
)
from litellm.caching.caching import DualCache
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.caching.redis_batch import flush_post_call_redis_batches
@ -3219,9 +3223,12 @@ class Logging(LiteLLMLoggingBaseClass):
batch_models = kwargs.get("batch_models", None)
batch_successful_requests: Final = kwargs.get("batch_successful_requests", None)
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
batch_output_file_content: Final = kwargs.get("batch_output_file_content", None)
batch_error_file_content: Final = kwargs.get("batch_error_file_content", None)
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result)
result_files: BatchResultFiles | None = None # rebind-ok: one batch-data branch below supplies the bytes
if has_explicit_batch_data:
result._hidden_params["response_cost"] = batch_cost
result._hidden_params["batch_models"] = batch_models
@ -3241,15 +3248,21 @@ class Logging(LiteLLMLoggingBaseClass):
total_cost=batch_cost,
cost_for_built_in_tools_cost_usd_dollar=0.0,
)
if batch_output_file_content is not None or batch_error_file_content is not None:
result_files = BatchResultFiles(
output=batch_output_file_content if isinstance(batch_output_file_content, bytes) else None,
error=batch_error_file_content if isinstance(batch_error_file_content, bytes) else None,
)
elif should_compute_batch_data:
batch_result: Final = await _handle_completed_batch(
batch_result, fetched_result_files = await _handle_completed_batch_with_files(
batch=result,
custom_llm_provider=self.custom_llm_provider,
model_name=self.get_deployment_model_for_cost(),
litellm_params=self.litellm_params,
model_info=self.get_router_deployment_model_info(),
)
result_files = fetched_result_files
result._hidden_params["response_cost"] = batch_result.cost
result._hidden_params["batch_models"] = batch_result.models
@ -3266,14 +3279,21 @@ class Logging(LiteLLMLoggingBaseClass):
if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data):
from litellm.batches.batch_line_item_logging import log_batch_line_items
await log_batch_line_items(
batch=result,
custom_llm_provider=self.custom_llm_provider,
parent=self,
model_name=self.get_deployment_model_for_cost(),
litellm_params=self.litellm_params,
model_info=self.get_router_deployment_model_info(),
)
try:
await log_batch_line_items(
batch=result,
custom_llm_provider=self.custom_llm_provider,
parent=self,
model_name=self.get_deployment_model_for_cost(),
litellm_params=self.litellm_params,
model_info=self.get_router_deployment_model_info(),
result_files=result_files,
)
except Exception: # noqa: BLE001 # line-item logging (claim step included) must never reach the aggregate aretrieve_batch path
verbose_logger.exception(
"batch line item logging failed for batch_id=%s; aggregate logging unaffected",
result.id,
)
self.truncated_messages_for_logging = await truncate_base64_in_messages_async(
StandardLoggingPayloadSetup.append_system_prompt_messages(

44
tests/batches_tests/test_batches_logging_unit_tests.py
View file

@ -222,8 +222,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
)
logging_obj.custom_llm_provider = "openai"
# Mock _handle_completed_batch to return cost data
from litellm.batches.batch_utils import BatchCostUsageResult
# Mock _handle_completed_batch_with_files to return cost data
from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles
expected_cost = 0.05
expected_usage = litellm.Usage(
@ -234,14 +234,17 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
expected_models = ["gpt-5-mini"]
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files",
new=AsyncMock(
return_value=BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=10,
failed_requests=0,
return_value=(
BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=10,
failed_requests=0,
),
BatchResultFiles(output=None, error=None),
)
),
) as mock_handle_batch:
@ -382,7 +385,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data():
explicit_models = ["gpt-5-mini", "gpt-5.5"]
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files",
new=AsyncMock(),
) as mock_handle_batch:
# Call async_success_handler with explicit cost data
@ -517,7 +520,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc
logging_obj.custom_llm_provider = "openai"
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files",
new=AsyncMock(),
) as mock_handle_batch:
# Call async_success_handler with in_progress batch (unified file ID)
@ -603,17 +606,20 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
)
expected_models = ["gpt-5-mini"]
from litellm.batches.batch_utils import BatchCostUsageResult
from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles
with patch(
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files",
new=AsyncMock(
return_value=BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=8,
failed_requests=0,
return_value=(
BatchCostUsageResult(
cost=expected_cost,
usage=expected_usage,
models=expected_models,
successful_requests=8,
failed_requests=0,
),
BatchResultFiles(output=None, error=None),
)
),
) as mock_handle_batch:

47
tests/unit/batches/test_batch_line_item_logging.py
View file

@ -26,6 +26,7 @@ from litellm.batches.batch_line_item_logging import (
batch_line_item_claim_cache,
log_batch_line_items,
)
from litellm.batches.batch_utils import BatchResultFiles
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.integrations.custom_logger import CustomLogger
@ -1240,3 +1241,49 @@ async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(re
assert fake._store.get(key) == "other-worker"
assert len(recorder.success_events) == 0
assert len(recorder.failure_events) == 0
@pytest.mark.asyncio
async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder):
file_mock: Final = AsyncMock(side_effect=_file_content)
with patch(
"litellm.files.main.afile_content", file_mock
): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch
emitted: Final = await log_batch_line_items(
batch=_batch(),
custom_llm_provider="openai",
parent=_parent_logging(),
model_name="gpt-4o",
litellm_params=None,
model_info=None,
result_files=BatchResultFiles(output=OUTPUT_JSONL, error=ERROR_JSONL),
claim_cache=DualCache(),
)
assert emitted == 2
assert len(recorder.success_events) == 1
assert len(recorder.failure_events) == 1
assert file_mock.await_count == 1
assert file_mock.await_args.kwargs["file_id"] == "input-file-1"
@pytest.mark.asyncio
async def test_empty_result_files_fall_back_to_fetching_everything(recorder):
file_mock: Final = AsyncMock(side_effect=_file_content)
with patch(
"litellm.files.main.afile_content", file_mock
): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch
emitted: Final = await log_batch_line_items(
batch=_batch(),
custom_llm_provider="openai",
parent=_parent_logging(),
model_name="gpt-4o",
litellm_params=None,
model_info=None,
result_files=BatchResultFiles(output=None, error=None),
claim_cache=DualCache(),
)
assert emitted == 2
fetched: Final = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list)
assert fetched == ["error-file-1", "input-file-1", "output-file-1"]

54
tests/unit/batches/test_batch_utils.py
View file

@ -2451,3 +2451,57 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_
assert calls == []
assert (result.successful_requests, result.failed_requests) == (0, 1)
@pytest.mark.asyncio
async def test_handle_completed_batch_with_files_returns_bytes_and_fetches_error_once(monkeypatch):
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
output_bytes = _vertex_jsonl(rows)
error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}])
fetched = []
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return output_bytes
async def fake_afile_content(**kw):
fetched.append(kw["file_id"])
return type("R", (), {"content": error_bytes})()
import litellm.cost_calculator as cc
import litellm.files.main as files_main
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3))
batch = _batch("of").model_copy(update={"error_file_id": "ef"})
result, files = await bu._handle_completed_batch_with_files(batch, custom_llm_provider="openai")
assert result.cost == 3.3
assert result.failed_requests == 1
assert files.output == output_bytes
assert files.error == error_bytes
assert fetched == ["ef"]
@pytest.mark.asyncio
async def test_handle_completed_batch_with_files_no_output_returns_error_bytes(monkeypatch):
error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}])
async def fake_afile_content(**kw):
return type("R", (), {"content": error_bytes})()
import litellm.files.main as files_main
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
batch = _batch(None).model_copy(update={"error_file_id": "ef"})
result, files = await bu._handle_completed_batch_with_files(batch, custom_llm_provider="openai")
assert result.failed_requests == 1
assert result.usage.total_tokens == 0
assert files.output is None
assert files.error == error_bytes

195
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -1185,19 +1185,24 @@ class TestRetrieveBatchCostPassesModelIdentity:
captured: dict[str, object] = {}
from litellm.batches.batch_utils import BatchCostUsageResult
from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles
async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult:
async def fake_handle_completed_batch(
**kwargs: object,
) -> tuple[BatchCostUsageResult, BatchResultFiles]:
captured.update(kwargs)
return BatchCostUsageResult(
cost=1.25,
usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800),
models=["m"],
successful_requests=1,
failed_requests=0,
return (
BatchCostUsageResult(
cost=1.25,
usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800),
models=["m"],
successful_requests=1,
failed_requests=0,
),
BatchResultFiles(output=None, error=None),
)
monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch)
monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", fake_handle_completed_batch)
obj = LitellmLogging(
model="bedrock/global.anthropic.claude-sonnet-4-6",
@ -1280,7 +1285,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
from litellm.litellm_core_utils import litellm_logging as logging_module
handle_completed_batch = AsyncMock()
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", handle_completed_batch)
batch = self._batch(status, output_file_id)
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
@ -1290,20 +1295,23 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
@pytest.mark.asyncio
async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None:
from litellm.batches.batch_utils import BatchCostUsageResult
from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles
from litellm.litellm_core_utils import litellm_logging as logging_module
from litellm.types.utils import Usage
handle_completed_batch = AsyncMock(
return_value=BatchCostUsageResult(
cost=8e-06,
usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35),
models=["gpt-5.6-luna"],
successful_requests=2,
failed_requests=0,
return_value=(
BatchCostUsageResult(
cost=8e-06,
usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35),
models=["gpt-5.6-luna"],
successful_requests=2,
failed_requests=0,
),
BatchResultFiles(output=None, error=None),
)
)
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", handle_completed_batch)
batch = self._batch("completed", "file-out")
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
@ -9081,3 +9089,154 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch):
logging_module._in_memory_loggers.clear()
monkeypatch.delenv("LITELLM_OTEL_V2", raising=False)
is_otel_v2_enabled.cache_clear()
class TestRetrieveBatchReusesFetchedResultFiles:
"""The aggregate path and log_batch_line_items must share one fetch of each
result file: output and error are fetched once total (not once each per
consumer), and forwarded bytes keep line-item logging off the wire."""
@staticmethod
def _batch_file_bytes() -> dict[str, bytes]:
input_jsonl = json.dumps(
{
"custom_id": "a",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
}
).encode()
output_jsonl = json.dumps(
{
"custom_id": "a",
"response": {
"status_code": 200,
"body": {
"id": "chatcmpl-1",
"model": "gpt-4o",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi back"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
},
},
}
).encode()
error_jsonl = json.dumps({"custom_id": "b", "error": {"message": "rejected"}}).encode()
return {"file-in": input_jsonl, "file-out": output_jsonl, "file-err": error_jsonl}
@staticmethod
def _logging_obj() -> LitellmLogging:
obj = LitellmLogging(
model="gpt-4o",
messages=[{"role": "user", "content": "Hey"}],
stream=False,
call_type="aretrieve_batch",
start_time=time.time(),
litellm_call_id="batch-call-reuse",
function_id="f",
)
obj.custom_llm_provider = "openai"
obj.update_environment_variables(
litellm_params={"metadata": {}},
optional_params={},
custom_llm_provider="openai",
)
return obj
@staticmethod
def _batch():
from litellm.types.utils import LiteLLMBatch
return LiteLLMBatch(
id="batch_reuse_fetched_files",
completion_window="24h",
created_at=1,
endpoint="/v1/chat/completions",
input_file_id="file-in",
object="batch",
status="completed",
output_file_id="file-out",
error_file_id="file-err",
)
@pytest.mark.asyncio
async def test_compute_path_fetches_each_result_file_once(self, monkeypatch) -> None:
file_bytes = self._batch_file_bytes()
async def fake_afile_content(**kwargs):
from types import SimpleNamespace
return SimpleNamespace(content=file_bytes[kwargs["file_id"]])
file_mock = AsyncMock(side_effect=fake_afile_content)
monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False)
monkeypatch.setattr("litellm.files.main.afile_content", file_mock)
monkeypatch.setattr("litellm.cost_calculator.batch_cost_calculator", lambda **kw: (0.01, 0.02))
await self._logging_obj()._async_success_handler_body(
result=self._batch(), start_time=None, end_time=None
)
fetched = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list)
assert fetched == ["file-err", "file-in", "file-out"], (
"each result file must be fetched exactly once across aggregate costing and line-item logging"
)
@pytest.mark.asyncio
async def test_explicit_kwargs_path_forwards_result_files(self, monkeypatch) -> None:
from litellm.types.utils import Usage
file_bytes = self._batch_file_bytes()
async def fake_afile_content(**kwargs):
from types import SimpleNamespace
return SimpleNamespace(content=file_bytes[kwargs["file_id"]])
file_mock = AsyncMock(side_effect=fake_afile_content)
monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False)
monkeypatch.setattr("litellm.files.main.afile_content", file_mock)
await self._logging_obj().async_success_handler(
result=self._batch(),
batch_cost=1.5,
batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
batch_models=["gpt-4o"],
batch_successful_requests=1,
batch_failed_requests=1,
batch_output_file_content=file_bytes["file-out"],
batch_error_file_content=file_bytes["file-err"],
)
assert file_mock.await_count == 1
assert file_mock.await_args.kwargs["file_id"] == "file-in", (
"forwarded output/error bytes must leave only the input file to fetch"
)
@pytest.mark.asyncio
async def test_line_item_logging_failure_leaves_aggregate_result_intact(self, monkeypatch) -> None:
from litellm.types.utils import Usage
monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False)
monkeypatch.setattr(
"litellm.batches.batch_line_item_logging.log_batch_line_items",
AsyncMock(side_effect=RuntimeError("claim backend exploded")),
)
batch = self._batch()
await self._logging_obj().async_success_handler(
result=batch,
batch_cost=1.5,
batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
batch_models=["gpt-4o"],
batch_successful_requests=1,
batch_failed_requests=1,
)
assert batch._hidden_params["response_cost"] == 1.5
assert batch.usage.total_tokens == 15

8
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -1546,10 +1546,10 @@ class TestCheckBatchCost:
Logging, "async_success_handler", new_callable=AsyncMock
) as success_handler,
):
provider.get("https://api.openai.com/v1/files/file-output-123/content").mock(
output_route = provider.get("https://api.openai.com/v1/files/file-output-123/content").mock(
return_value=httpx.Response(200, content=f"{succeeded_line}\n{rejected_line}\n".encode())
)
provider.get("https://api.openai.com/v1/files/file-error-456/content").mock(
error_route = provider.get("https://api.openai.com/v1/files/file-error-456/content").mock(
return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode())
)
await check_batch_cost_instance.check_batch_cost()
@ -1558,6 +1558,10 @@ class TestCheckBatchCost:
assert len(spend_log_calls) == 1
handler_kwargs = spend_log_calls[0]
assert handler_kwargs["batch_successful_requests"] == 1
assert output_route.call_count == 1
assert error_route.call_count == 1
assert handler_kwargs["batch_output_file_content"] == f"{succeeded_line}\n{rejected_line}\n".encode()
assert handler_kwargs["batch_error_file_content"] == f"{error_file_lines}\n\n".encode()
assert handler_kwargs["batch_failed_requests"] == 3, (
"2 error-file lines must add to the output file's 1 rejected request"
)