diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index db28317ae2c..cb41eba32eb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 171c03d6f00..db9f0c9b6b1 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -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: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8160618f2c5..3660dfed94b 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -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: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index de573d04f98..d5b111a41cf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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( diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 73adb391481..e8a90b385f2 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -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: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 877b6b67854..a8edb920bf7 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -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"] diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index b8b922f72a7..367039cc0b4 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 60b7ed32399..9239e47ee7f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 4e7effad3bf..41470cc0d2e 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -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" )