From 141ada1118bc64c011b67ef9fd7ee2f15854865c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:25:55 -0400 Subject: [PATCH 01/16] feat(batches): aggregate reasoning tokens and per-line pass/fail counts Batch retrieval already computed cost/usage on completion, but silently dropped reasoning tokens and never counted per-line success/failure. Adds BatchCostUsageResult (replacing bare cost/usage/models tuples) with successful_requests/failed_requests, and threads reasoning_tokens through the aggregated Usage. Both surface on SpendLogs the same way batch_models already does. --- .../proxy/common_utils/check_batch_cost.py | 22 +- litellm/batches/batch_utils.py | 92 +++++-- litellm/litellm_core_utils/litellm_logging.py | 24 +- litellm/proxy/_types.py | 2 + .../spend_tracking/spend_tracking_utils.py | 16 ++ litellm/types/utils.py | 4 +- .../test_batch_custom_pricing.py | 16 +- tests/batches_tests/test_batch_rate_limits.py | 14 +- .../test_batches_logging_unit_tests.py | 58 ++-- .../proxy_unit_tests/test_check_batch_cost.py | 32 ++- .../test_litellm/batches/test_batch_utils.py | 259 ++++++++++++------ .../test_vertex_ai_batch_passthrough.py | 78 +++--- 12 files changed, 406 insertions(+), 211 deletions(-) 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 25b00597355..85b9bd77bc3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -651,16 +651,14 @@ class CheckBatchCost: # Pass deployment model_info so custom batch pricing # (input_cost_per_token_batches etc.) is used for cost calc deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {} - batch_cost, batch_usage, batch_models = ( - await calculate_batch_cost_and_usage( - file_content_dictionary=file_content_as_dict, - custom_llm_provider=llm_provider, # type: ignore - model_name=model_name, - model_info=deployment_model_info, # type: ignore[arg-type] - ) + batch_result = await calculate_batch_cost_and_usage( + file_content_dictionary=file_content_as_dict, + custom_llm_provider=llm_provider, # type: ignore + model_name=model_name, + model_info=deployment_model_info, # type: ignore[arg-type] ) logging_obj = LiteLLMLogging( - model=batch_models[0], + model=batch_result.models[0], messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", @@ -684,9 +682,11 @@ class CheckBatchCost: await logging_obj.async_success_handler( result=response, - batch_cost=batch_cost, - batch_usage=batch_usage, - batch_models=batch_models, + batch_cost=batch_result.cost, + batch_usage=batch_result.usage, + batch_models=batch_result.models, + batch_successful_requests=batch_result.successful_requests, + batch_failed_requests=batch_result.failed_requests, ) # Record batch duration (completed_at - created_at) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 9681d64f656..20aa3c755bd 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -12,12 +12,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import token_counter +@dataclass(frozen=True, slots=True) +class BatchCostUsageResult: + """Aggregate cost, usage, and per-line pass/fail counts for a completed batch.""" + + cost: float + usage: Usage + models: list[str] + successful_requests: int + failed_requests: int + + async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, model_info: ModelInfo | None = None, -) -> tuple[float, Usage, list[str]]: +) -> BatchCostUsageResult: """ Calculate the cost and usage of a batch. @@ -32,8 +43,7 @@ async def calculate_batch_cost_and_usage( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) - return batch_cost, batch_usage, [model_name] + return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, @@ -48,7 +58,7 @@ async def _handle_completed_batch( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, litellm_params: dict | None = None, -) -> tuple[float, Usage, list[str]]: +) -> BatchCostUsageResult: """Fetch a completed batch's output file and aggregate its cost, usage, and models in a single pass over the JSONL lines, so the parsed file content is never materialized in memory. @@ -66,10 +76,7 @@ async def _handle_completed_batch( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage( - _get_file_content_as_dictionary(file_content), model_name - ) - return batch_cost, batch_usage, [model_name] + return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) return _aggregate_batch_cost_usage_models( entries=_iter_batch_input_entries(file_content), @@ -86,19 +93,24 @@ class _BatchOutputLineStats: total_tokens: int cache_read_tokens: int cache_creation_tokens: int + reasoning_tokens: int model: str | None -def _iter_successful_output_line_stats( +def _classify_output_line_stats( entries: Iterable[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, -) -> Iterator[_BatchOutputLineStats]: +) -> Iterator[_BatchOutputLineStats | None]: + """Classify every output line in a single pass: yields stats for a + successful line, ``None`` for a failed one (per ``_batch_response_was_successful``). + Counting failures this way avoids a second pass over a potentially huge output file.""" from litellm.cost_calculator import batch_cost_calculator for entry in entries: if not _batch_response_was_successful(entry, custom_llm_provider): + yield None 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) @@ -123,6 +135,7 @@ def _iter_successful_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, @@ -130,6 +143,7 @@ def _iter_successful_output_line_stats( total_tokens=usage.total_tokens, cache_read_tokens=prompt_details["cache_hit_tokens"], cache_creation_tokens=prompt_details["cache_creation_tokens"], + reasoning_tokens=reasoning_tokens or 0, model=response_model, ) @@ -139,10 +153,14 @@ def _aggregate_batch_cost_usage_models( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None = None, model_info: ModelInfo | None = None, -) -> tuple[float, Usage, list[str]]: - """Aggregate cost, usage, and models from batch output entries in a single - pass, holding one small stats record per line instead of the parsed file.""" - line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info)) +) -> BatchCostUsageResult: + """Aggregate cost, usage, models, and pass/fail counts from batch output + entries in a single pass, holding one small stats record per line instead + of the parsed file.""" + all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info)) + line_stats: Final = tuple(stats for stats in all_results if stats is not None) + successful_requests: Final = len(line_stats) + failed_requests: Final = len(all_results) - successful_requests cache_token_params: Final = { key: tokens @@ -156,18 +174,32 @@ def _aggregate_batch_cost_usage_models( total_tokens=sum(stats.total_tokens for stats in line_stats), prompt_tokens=sum(stats.prompt_tokens for stats in line_stats), completion_tokens=sum(stats.completion_tokens for stats in line_stats), + reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats), **cache_token_params, ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) - verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models) - return total_cost, batch_usage, batch_models + verbose_logger.debug( + "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d", + total_cost, + batch_usage, + batch_models, + successful_requests, + failed_requests, + ) + return BatchCostUsageResult( + cost=total_cost, + usage=batch_usage, + models=batch_models, + successful_requests=successful_requests, + failed_requests=failed_requests, + ) def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, -) -> tuple[float, Usage]: +) -> BatchCostUsageResult: """ Calculate both cost and usage from raw Vertex AI batch responses. @@ -178,6 +210,10 @@ def calculate_vertex_ai_batch_cost_and_usage( {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. + + A row with no ``response`` is counted as failed - the same signal already + used to skip it from cost/usage aggregation, since Vertex batch prediction + output doesn't establish a distinct error shape in this (non-default) path. """ from litellm.cost_calculator import batch_cost_calculator @@ -185,12 +221,16 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 + successful_requests = 0 + failed_requests = 0 actual_model_name: Final = model_name or "gemini-2.0-flash-001" for response in vertex_ai_batch_responses: response_body = response.get("response") if response_body is None: + failed_requests += 1 continue + successful_requests += 1 usage_metadata = response_body.get("usageMetadata", {}) _prompt = usage_metadata.get("promptTokenCount", 0) or 0 @@ -218,17 +258,25 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens += _total verbose_logger.info( - "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d", + "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", total_cost, prompt_tokens, completion_tokens, total_tokens, + successful_requests, + failed_requests, ) - return total_cost, Usage( - total_tokens=total_tokens, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + return BatchCostUsageResult( + cost=total_cost, + usage=Usage( + total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ), + models=[actual_model_name], + successful_requests=successful_requests, + failed_requests=failed_requests, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a72d46e3fe8..9d035afbde2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,6 +2574,8 @@ class Logging(LiteLLMLoggingBaseClass): batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) + batch_successful_requests = kwargs.get("batch_successful_requests", None) + batch_failed_requests = kwargs.get("batch_failed_requests", 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 = ( @@ -2582,22 +2584,22 @@ class Logging(LiteLLMLoggingBaseClass): if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models + result._hidden_params["batch_successful_requests"] = batch_successful_requests + result._hidden_params["batch_failed_requests"] = batch_failed_requests result.usage = batch_usage elif should_compute_batch_data: - ( - response_cost, - batch_usage, - batch_models, - ) = await _handle_completed_batch( + batch_result = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, ) - result._hidden_params["response_cost"] = response_cost - result._hidden_params["batch_models"] = batch_models - result.usage = batch_usage + result._hidden_params["response_cost"] = batch_result.cost + result._hidden_params["batch_models"] = batch_result.models + result._hidden_params["batch_successful_requests"] = batch_result.successful_requests + result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + result.usage = batch_result.usage start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, @@ -5062,6 +5064,8 @@ class StandardLoggingPayloadSetup: additional_headers=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) @@ -5451,6 +5455,8 @@ def _extract_response_obj_and_hidden_params( response_cost=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) @@ -5819,6 +5825,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: additional_headers=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a566d491597..297fc1bd201 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3462,6 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None + batch_successful_requests: int | None + batch_failed_requests: int | None error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3146d8bccfb..f08236adab2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -74,6 +74,8 @@ def _get_spend_logs_metadata( metadata: dict | None, applied_guardrails: list[str] | None = None, batch_models: list[str] | None = None, + batch_successful_requests: int | None = None, + batch_failed_requests: int | None = None, mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None, vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None, guardrail_information: list[StandardLoggingGuardrailInformation] | None = None, @@ -102,6 +104,8 @@ def _get_spend_logs_metadata( error_information=None, proxy_server_request=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, mcp_tool_call_metadata=None, vector_store_request_metadata=None, model_map_information=None, @@ -128,6 +132,8 @@ def _get_spend_logs_metadata( clean_metadata["user_api_key"] = _hash_api_key_for_spend_log(raw_user_api_key) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models + clean_metadata["batch_successful_requests"] = batch_successful_requests + clean_metadata["batch_failed_requests"] = batch_failed_requests clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload( vector_store_request_metadata @@ -310,6 +316,16 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + batch_successful_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None) + if standard_logging_payload is not None + else None + ), + batch_failed_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None) + if standard_logging_payload is not None + else None + ), mcp_tool_call_metadata=( standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None) if standard_logging_payload is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 272fbabf807..dc8adf9b88e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -39,7 +39,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -2880,6 +2880,8 @@ class StandardLoggingHiddenParams(TypedDict): litellm_overhead_time_ms: float | None additional_headers: StandardLoggingAdditionalHeaders | None batch_models: list[str] | None + batch_successful_requests: ReadOnly[int | None] + batch_failed_requests: ReadOnly[int | None] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index c2159b564a8..b76b865862a 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -116,16 +116,16 @@ def test_aggregate_batch_cost_uses_custom_model_info(): """_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - cost, _, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {cost}" + ), f"Expected total cost {expected}, got {result.cost}" @pytest.mark.parametrize("data_residency", ["eu", "us"]) @@ -164,15 +164,15 @@ async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): """calculate_batch_cost_and_usage should thread model_info.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - batch_cost, batch_usage, batch_models = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert batch_cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {batch_cost}" - assert batch_usage.prompt_tokens == 10 - assert batch_usage.completion_tokens == 5 + ), f"Expected total cost {expected}, got {result.cost}" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 2c804d21ace..0baff6c17be 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -1027,7 +1027,7 @@ async def test_batch_logging_azure_credentials_regression(): with patch( "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, @@ -1039,13 +1039,13 @@ async def test_batch_logging_azure_credentials_regression(): ], "REGRESSION: Credentials not passed through _handle_completed_batch" # Verify cost and usage were calculated - assert cost > 0, "Cost should be calculated" - assert usage.total_tokens == 40, "Usage should be calculated correctly" + assert result.cost > 0, "Cost should be calculated" + assert result.usage.total_tokens == 40, "Usage should be calculated correctly" print(" ✓ Credentials passed through full flow") - print(f" ✓ Cost: {cost}") - print(f" ✓ Usage: {usage.total_tokens} tokens") - print(f" ✓ Models: {models}") + print(f" ✓ Cost: {result.cost}") + print(f" ✓ Usage: {result.usage.total_tokens} tokens") + print(f" ✓ Models: {result.models}") # Test 4: Verify error prevention print("\n4. Testing 'Missing credentials' error prevention...") @@ -1064,7 +1064,7 @@ async def test_batch_logging_azure_credentials_regression(): "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): try: - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 62b6f5b08e4..9c26514f872 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -138,12 +138,12 @@ def test_get_file_content_as_dictionary(sample_file_content): def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): with patch("litellm.completion_cost", return_value=0.0): - _, usage, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai" ) - assert usage.total_tokens == 62 # 30 + 32 - assert usage.prompt_tokens == 42 # 20 + 22 - assert usage.completion_tokens == 20 # 10 + 10 + assert result.usage.total_tokens == 62 # 30 + 32 + assert result.usage.prompt_tokens == 42 # 20 + 22 + assert result.usage.completion_tokens == 20 # 10 + 10 @pytest.mark.asyncio @@ -156,11 +156,11 @@ async def test_batch_cost_calculator(sample_file_content_dict): so we expect the cost to be 0.5 * 2 = 1.0 """ with patch("litellm.completion_cost", return_value=0.5): - cost, _, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai", ) - assert cost == 1.0 # 0.5 * 2 successful responses + assert result.cost == 1.0 # 0.5 * 2 successful responses def test_get_response_from_batch_job_output_file(sample_file_content_dict): @@ -226,6 +226,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 + expected_cost = 0.05 expected_usage = litellm.Usage( prompt_tokens=100, @@ -236,7 +238,15 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos with patch( "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", - new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)), + new=AsyncMock( + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=10, + failed_requests=0, + ) + ), ) as mock_handle_batch: # Call async_success_handler await logging_obj.async_success_handler( @@ -251,6 +261,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos # Verify cost and usage were set on the batch result assert mock_batch._hidden_params["response_cost"] == expected_cost assert mock_batch._hidden_params["batch_models"] == expected_models + assert mock_batch._hidden_params["batch_successful_requests"] == 10 + assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage @@ -284,7 +296,7 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( "litellm.batches.batch_utils._fetch_batch_output_file_content", new=AsyncMock(return_value=sample_file_content_bytes), ): - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=batch, custom_llm_provider="openai" ) @@ -294,16 +306,18 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( + 20 * pricing["output_cost_per_token_batches"] ) - assert cost == pytest.approx(expected_cost) - assert cost > 0 + assert result.cost == pytest.approx(expected_cost) + assert result.cost > 0 assert ( - cost + result.cost < 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"] ) - assert usage.prompt_tokens == 42 - assert usage.completion_tokens == 20 - assert usage.total_tokens == 62 - assert models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"] + assert result.usage.prompt_tokens == 42 + assert result.usage.completion_tokens == 20 + assert result.usage.total_tokens == 62 + assert result.models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"] + assert result.successful_requests == 2 + assert result.failed_requests == 0 @pytest.mark.asyncio @@ -542,9 +556,19 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): ) expected_models = ["gpt-5-mini"] + from litellm.batches.batch_utils import BatchCostUsageResult + with patch( "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", - new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)), + new=AsyncMock( + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=8, + failed_requests=0, + ) + ), ) as mock_handle_batch: # Call async_success_handler with partial explicit data await logging_obj.async_success_handler( @@ -560,4 +584,6 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): # Verify computed cost data was used (not partial explicit data) assert mock_batch._hidden_params["response_cost"] == expected_cost assert mock_batch._hidden_params["batch_models"] == expected_models + assert mock_batch._hidden_params["batch_successful_requests"] == 8 + assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 72f8b87dd16..79739669ff4 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -13,6 +13,20 @@ import pytest _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" +def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): + """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, + for mocking it in tests that only care about cost/usage/models.""" + from litellm.batches.batch_utils import BatchCostUsageResult + + return BatchCostUsageResult( + cost=cost, + usage=usage, + models=models, + successful_requests=successful_requests, + failed_requests=failed_requests, + ) + + def _unmanaged_vertex_file_object( input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl", status="validating", @@ -321,7 +335,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -426,7 +440,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -526,7 +540,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -656,7 +670,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1142,7 +1156,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1271,7 +1285,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1529,7 +1543,7 @@ class TestUnmanagedVertexRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gemini-2.5-flash"], @@ -1759,7 +1773,7 @@ class TestUnmanagedBedrockRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.02, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-sonnet-4"], @@ -1951,7 +1965,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), ), patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls, ): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index d2074853f2b..46969dfc033 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -210,10 +210,10 @@ def test_estimate_tokens_never_zero_for_short_rows(): def test_output_models_uses_model_name_override(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) - _, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model" ) - assert models == ["forced-model"] + assert result.models == ["forced-model"] def test_output_models_collects_from_successful_only(monkeypatch): @@ -223,15 +223,15 @@ def test_output_models_collects_from_successful_only(monkeypatch): _failed_row(model="should-be-skipped"), _success_row(model="claude-3"), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == ["gpt-4o", "claude-3"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == ["gpt-4o", "claude-3"] def test_output_models_skips_successful_without_model(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) rows = [{"response": {"status_code": 200, "body": {}}}] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == [] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == [] # =========================================================================== # @@ -398,8 +398,8 @@ def test_total_usage_sums_successful_only(monkeypatch): _failed_row(), # excluded _success_row(usage=_usage(20, 10)), # 30 ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, @@ -417,7 +417,7 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): ) chat_row = _success_row(usage=_usage(10, 5)) - cost, usage, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[responses_row, chat_row], custom_llm_provider="openai", model_info={ @@ -426,22 +426,79 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): }, ) - assert usage.prompt_tokens == 30 - assert usage.completion_tokens == 12 - assert usage.total_tokens == 42 - assert usage.cache_read_input_tokens == 3 - assert cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) + assert result.usage.prompt_tokens == 30 + assert result.usage.completion_tokens == 12 + assert result.usage.total_tokens == 42 + assert result.usage.cache_read_input_tokens == 3 + assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) def test_total_usage_empty_is_zero(): - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") - assert cost == 0.0 - assert models == [] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") + assert result.cost == 0.0 + assert result.models == [] + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 0, 0, 0, ) + assert result.successful_requests == 0 + assert result.failed_requests == 0 + + +def test_total_usage_includes_reasoning_tokens(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row( + usage={ + "prompt_tokens": 10, + "completion_tokens": 50, + "total_tokens": 60, + "completion_tokens_details": {"reasoning_tokens": 30}, + } + ), + _success_row( + usage={ + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + "completion_tokens_details": {"reasoning_tokens": 8}, + } + ), + _failed_row(), # excluded, must not contribute reasoning tokens either + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.usage.completion_tokens_details is not None + assert result.usage.completion_tokens_details.reasoning_tokens == 38 + + +def test_aggregate_counts_successful_and_failed_requests(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row(usage=_usage(10, 5)), + _failed_row(), + _success_row(usage=_usage(20, 10)), + _failed_row(), + _failed_row(), + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.successful_requests == 2 + assert result.failed_requests == 3 + assert result.successful_requests + result.failed_requests == len(rows) + + +def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" + ) + assert isinstance(result, bu.BatchCostUsageResult) + assert (result.cost, result.models, result.successful_requests, result.failed_requests) == ( + 1.0, + ["gpt-4o"], + 1, + 0, + ) # =========================================================================== # @@ -464,10 +521,12 @@ def test_cost_from_content_completion_cost_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert total == 1.0 # 2 successful * 0.5 + assert result.cost == 1.0 # 2 successful * 0.5 assert len(calls) == 2 # failed row not costed + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_cost_from_content_model_info_path(monkeypatch): @@ -480,13 +539,13 @@ def test_cost_from_content_model_info_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="openai", model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path ) - assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): @@ -496,11 +555,13 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") - assert cost == 1.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gpt-4o", "gpt-4o"] + assert result.cost == 1.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gpt-4o", "gpt-4o"] + assert result.successful_requests == 2 + assert result.failed_requests == 1 # =========================================================================== # @@ -514,7 +575,13 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)), + lambda content, model: bu.BatchCostUsageResult( + cost=9.9, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-2.0-flash-001"], + successful_requests=1, + failed_requests=0, + ), ) # generic path must NOT be taken monkeypatch.setattr( @@ -523,12 +590,12 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): lambda **kw: pytest.fail("generic path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001" ) - assert cost == 9.9 - assert usage.total_tokens == 3 - assert models == ["gemini-2.0-flash-001"] + assert result.cost == 9.9 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-2.0-flash-001"] @pytest.mark.asyncio @@ -542,12 +609,12 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai" ) - assert cost == 0.0 - assert usage.total_tokens == 0 - assert models == [] + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.models == [] # =========================================================================== # @@ -580,14 +647,16 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, ) + assert result.successful_requests == 2 + assert result.failed_requests == 0 def test_vertex_cost_skips_none_response_body(monkeypatch): @@ -607,10 +676,12 @@ def test_vertex_cost_skips_none_response_body(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(1.0) # only one line costed - assert usage.total_tokens == 10 + assert result.cost == pytest.approx(1.0) # only one line costed + assert result.usage.total_tokens == 10 + assert result.successful_requests == 1 + assert result.failed_requests == 1 def test_vertex_usage_total_token_fallback(monkeypatch): @@ -620,8 +691,8 @@ def test_vertex_usage_total_token_fallback(monkeypatch): monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0)) responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}] - _, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert usage.total_tokens == 12 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.usage.total_tokens == 12 def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @@ -644,9 +715,9 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): } ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == 0.0 - assert usage.total_tokens == 10 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.cost == 0.0 + assert result.usage.total_tokens == 10 # =========================================================================== # @@ -659,13 +730,11 @@ async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) - cost, usage, models = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="openai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") - assert cost == 2.5 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 2.5 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] # =========================================================================== # @@ -883,16 +952,18 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"}, ) - assert cost > 0 - assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.cost > 0 + assert result.cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.successful_requests == 2 + assert result.failed_requests == 0 @pytest.mark.asyncio @@ -970,11 +1041,11 @@ async def test_handle_completed_batch_orchestration(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) - cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") - assert cost == 3.3 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 3.3 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] @pytest.mark.asyncio @@ -991,19 +1062,25 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) def fake_vertex_calc(content, model): seen["content"] = content seen["model"] = model - return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) + return bu.BatchCostUsageResult( + cost=7.7, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-x"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", model_name="gemini-x", ) - assert cost == 7.7 - assert usage.total_tokens == 3 - assert models == ["gemini-x"] + assert result.cost == 7.7 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-x"] assert seen["content"] == raw_rows assert seen["model"] == "gemini-x" @@ -1105,14 +1182,14 @@ def test_bedrock_cost_uses_deployment_model_name(): "recordId": "1", "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } - cost, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[row], custom_llm_provider="bedrock", model_name="us.anthropic.claude-sonnet-4-6", model_info={}, ) - assert cost > 0 - assert models == ["us.anthropic.claude-sonnet-4-6"] + assert result.cost > 0 + assert result.models == ["us.anthropic.claude-sonnet-4-6"] def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): @@ -1124,8 +1201,10 @@ def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145) + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145) + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): @@ -1137,11 +1216,11 @@ def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert usage.prompt_tokens_details.cached_tokens == 8700 - assert usage.prompt_tokens_details.cache_creation_tokens == 2300 - assert usage.cache_read_input_tokens == 8700 - assert usage.cache_creation_input_tokens == 2300 + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.usage.prompt_tokens_details.cached_tokens == 8700 + assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300 + assert result.usage.cache_read_input_tokens == 8700 + assert result.usage.cache_creation_input_tokens == 2300 def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): @@ -1152,9 +1231,9 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, } ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert usage.prompt_tokens_details is None + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.usage.prompt_tokens_details is None def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): @@ -1165,14 +1244,14 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): _anthropic_errored_row(), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="anthropic", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2 - assert total == pytest.approx(expected_half_price) + assert result.cost == pytest.approx(expected_half_price) def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch): @@ -1191,11 +1270,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - total, _, _ = bu._aggregate_batch_cost_usage_models( - entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" - ) + result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") - assert total == pytest.approx(0.3) + assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" assert seen[0]["custom_llm_provider"] == "anthropic" assert seen[0]["usage"].prompt_tokens == 10 @@ -1209,8 +1286,8 @@ def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch): _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"), _anthropic_errored_row(), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert models == ["claude-sonnet-4-5-20250929"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.models == ["claude-sonnet-4-5-20250929"] @pytest.mark.asyncio @@ -1220,16 +1297,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): _anthropic_errored_row(), ] - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=rows, custom_llm_provider="anthropic", model_name="claude-sonnet-4-5", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) - assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) - assert models == ["claude-sonnet-4-5"] + assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert result.models == ["claude-sonnet-4-5"] def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index ac79c183ca3..1d2d7d4d5c3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -478,14 +478,14 @@ class TestVertexAIBatchPassthroughHandler: } ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses, model_name="gemini-2.0-flash-001" ) - assert usage.total_tokens == 15 - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.total_tokens == 15 + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_batch_response_transformation(self): """Test transformation of Vertex AI batch responses to OpenAI format""" @@ -664,14 +664,14 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_should_skip_responses_with_null_response_body(self): """Failed lines (response: None) are skipped without error.""" @@ -699,27 +699,29 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0 + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0 + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_should_return_zeros_for_empty_response_list(self): """Empty input → zero cost and zero usage.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( [], model_name="gemini-2.0-flash-001" ) - assert total_cost == 0.0 - assert usage.total_tokens == 0 - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 def test_should_handle_missing_usage_metadata_gracefully(self): """Response without usageMetadata → 0 tokens, 0 cost for that line.""" @@ -729,13 +731,13 @@ class TestVertexAIBatchCostCalculation: {"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}}, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 - assert usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 0 @pytest.mark.asyncio async def test_openai_shaped_output_records_nonzero_cost_and_usage(self): @@ -813,7 +815,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = False - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=openai_shaped_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -822,17 +824,17 @@ class TestVertexAIBatchCostCalculation: litellm.disable_vertex_batch_output_transformation = original_flag assert ( - usage.prompt_tokens == 18 - ), f"expected 18 prompt tokens, got {usage.prompt_tokens}" + result.usage.prompt_tokens == 18 + ), f"expected 18 prompt tokens, got {result.usage.prompt_tokens}" assert ( - usage.completion_tokens == 8 - ), f"expected 8 completion tokens, got {usage.completion_tokens}" + result.usage.completion_tokens == 8 + ), f"expected 8 completion tokens, got {result.usage.completion_tokens}" assert ( - usage.total_tokens == 26 - ), f"expected 26 total tokens, got {usage.total_tokens}" + result.usage.total_tokens == 26 + ), f"expected 26 total tokens, got {result.usage.total_tokens}" assert ( - cost > 0 - ), f"expected non-zero cost for completed Vertex batch, got {cost}" + result.cost > 0 + ), f"expected non-zero cost for completed Vertex batch, got {result.cost}" @pytest.mark.asyncio async def test_raw_vertex_output_still_works_when_transformation_disabled(self): @@ -865,7 +867,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = True - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=raw_vertex_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -873,7 +875,7 @@ class TestVertexAIBatchCostCalculation: finally: litellm.disable_vertex_batch_output_transformation = original_flag - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert usage.total_tokens == 15 - assert cost > 0, "raw Vertex shape should also produce non-zero cost" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.usage.total_tokens == 15 + assert result.cost > 0, "raw Vertex shape should also produce non-zero cost" From 2bfa1613b4837db3bc01e48297588a123c0c08e9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:48:04 -0400 Subject: [PATCH 02/16] fix(batches): satisfy LIT010/LIT011/reportPrivateUsage budgets for new fields Suppress the mutation/private-access lints the new batch_successful_requests/batch_failed_requests plumbing triggers, matching the existing suppressed pattern already used for response_cost/batch_models on the same lines. --- litellm/batches/batch_utils.py | 4 ++-- litellm/litellm_core_utils/litellm_logging.py | 14 +++++++------- litellm/proxy/_types.py | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 20aa3c755bd..6f6de18b04e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -221,8 +221,8 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 - successful_requests = 0 - failed_requests = 0 + successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above + failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above actual_model_name: Final = model_name or "gemini-2.0-flash-001" for response in vertex_ai_batch_responses: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9d035afbde2..c795b181064 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,8 +2574,8 @@ class Logging(LiteLLMLoggingBaseClass): batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) - batch_successful_requests = kwargs.get("batch_successful_requests", None) - batch_failed_requests = kwargs.get("batch_failed_requests", None) + batch_successful_requests: Final = kwargs.get("batch_successful_requests", None) + batch_failed_requests: Final = kwargs.get("batch_failed_requests", 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 = ( @@ -2584,12 +2584,12 @@ class Logging(LiteLLMLoggingBaseClass): if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models - result._hidden_params["batch_successful_requests"] = batch_successful_requests - result._hidden_params["batch_failed_requests"] = batch_failed_requests + result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above + result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_usage elif should_compute_batch_data: - batch_result = await _handle_completed_batch( + batch_result: Final = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, @@ -2597,8 +2597,8 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["response_cost"] = batch_result.cost result._hidden_params["batch_models"] = batch_result.models - result._hidden_params["batch_successful_requests"] = batch_result.successful_requests - result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above + result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage start_time, end_time, result = self._success_handler_helper_fn( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 297fc1bd201..cbc8c003868 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -15,7 +15,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import NotRequired, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._uuid import uuid from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS @@ -3462,8 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: int | None - batch_failed_requests: int | None + batch_successful_requests: ReadOnly[int | None] + batch_failed_requests: ReadOnly[int | None] error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None From 02cf319b483d8c6961b5464f0a72875c0ee6ae4a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 14:58:14 -0400 Subject: [PATCH 03/16] fix(batches): count failures reported only in the batch error file Live verification against a real OpenAI batch showed per-request failures (e.g. a rejected param) land in error_file_id, never in the output file, so failed_requests silently undercounted them (0 instead of the real 1). _handle_completed_batch now also fetches error_file_id when present and folds its line count into failed_requests. --- litellm/batches/batch_utils.py | 124 +++++++++++++----- .../test_litellm/batches/test_batch_utils.py | 63 +++++++++ 2 files changed, 152 insertions(+), 35 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6f6de18b04e..4889be3c0f6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,6 +1,7 @@ import json from collections.abc import Iterable, Iterator from dataclasses import dataclass +from dataclasses import replace as dataclasses_replace from typing import Any, Final, Literal import litellm @@ -70,18 +71,28 @@ async def _handle_completed_batch( litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) """ 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 + ) - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) - ): - return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) + output_file_result: Final = ( + calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) + if ( + custom_llm_provider == "vertex_ai" + and model_name + and getattr(litellm, "disable_vertex_batch_output_transformation", False) + ) + else _aggregate_batch_cost_usage_models( + entries=_iter_batch_input_entries(file_content), + custom_llm_provider=custom_llm_provider, + model_name=model_name, + ) + ) - return _aggregate_batch_cost_usage_models( - entries=_iter_batch_input_entries(file_content), - custom_llm_provider=custom_llm_provider, - model_name=model_name, + 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 ) @@ -280,6 +291,50 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +async def _fetch_batch_managed_file_content( + file_id: str, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + litellm_params: dict | None = None, +) -> bytes: + """ + Fetch a batch's output or error file and return its raw JSONL bytes. + + Args: + file_id: The provider or unified (litellm-managed) file id to fetch + custom_llm_provider: The LLM provider + litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) + Required for Azure and other providers that need authentication + """ + from litellm.files.main import afile_content + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) + + resolved_file_id = file_id + is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) + if is_base64_unified_file_id: + try: + resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e + ) + + # Build kwargs for afile_content with credentials from litellm_params + file_content_kwargs: Final = { + "file_id": resolved_file_id, + "custom_llm_provider": custom_llm_provider, + } + + # Extract and add credentials for file access + credentials: Final = _extract_file_access_credentials(litellm_params) + file_content_kwargs.update(credentials) + + _file_content: Final = await afile_content(**file_content_kwargs) + return _file_content.content + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -294,37 +349,36 @@ async def _fetch_batch_output_file_content( litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) Required for Azure and other providers that need authentication """ - from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + return await _fetch_batch_managed_file_content( + batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) - # Build kwargs for afile_content with credentials from litellm_params - file_content_kwargs: Final = { - "file_id": file_id, - "custom_llm_provider": custom_llm_provider, - } - # Extract and add credentials for file access - credentials: Final = _extract_file_access_credentials(litellm_params) - file_content_kwargs.update(credentials) +async def _count_error_file_failed_requests( + batch: Batch, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + litellm_params: dict | None, +) -> int: + """Count failed requests reported only in the batch's separate error file. - _file_content: Final = await afile_content(**file_content_kwargs) - return _file_content.content + OpenAI-shaped batch providers write successful lines to ``output_file_id`` + and per-request failures (e.g. a rejected param) to a distinct + ``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 + ) + 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/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 46969dfc033..06f42ce7b51 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1048,6 +1048,69 @@ async def test_handle_completed_batch_orchestration(monkeypatch): assert result.models == ["gpt-4o"] +@pytest.mark.asyncio +async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): + """Regression test: OpenAI writes per-request failures (e.g. a rejected param) + to a separate error_file_id, never into the output file - so failed_requests + must include them or it silently undercounts real batch failures.""" + from litellm.types.llms.openai import Batch + + rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))] + error_rows = [ + { + "id": "batch_req_err1", + "custom_id": "req-2-bad", + "response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}}, + "error": None, + } + ] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + async def fake_afile_content(**kw): + return type("R", (), {"content": _vertex_jsonl(error_rows)})() + + 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(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id="ef", + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 1 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 0 + + @pytest.mark.asyncio async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch): raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}] From 682032dd1b23ba49aab8ba486b44716805619ec1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:29:25 -0400 Subject: [PATCH 04/16] fix(batches): fix CI failures from lint format and merged upstream guard - ruff format litellm/batches/batch_utils.py - reconcile the upstream output_file_id=None guard (merged in from litellm_internal_staging) with BatchCostUsageResult, and count that batch's error_file_id failures instead of always reporting 0 - fix test_handle_completed_batch_no_output_file_is_zero's tuple unpacking, which predated the BatchCostUsageResult refactor - commit the batch_successful_requests/batch_failed_requests fixture fix to test_spend_management_endpoints.py that was left uncommitted --- litellm/batches/batch_utils.py | 14 ++++++++++---- tests/test_litellm/batches/test_batch_utils.py | 10 ++++++---- .../test_spend_management_endpoints.py | 6 +++--- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8a42efd6dc5..4eb9c7a5dfa 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -79,7 +79,15 @@ async def _handle_completed_batch( # The generic retrieval helper keeps raising for callers that explicitly ask # for a missing output file. if batch.output_file_id is None: - return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), [] + return BatchCostUsageResult( + cost=0.0, + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + models=[], + successful_requests=0, + failed_requests=await _count_error_file_failed_requests( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ), + ) 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( @@ -328,9 +336,7 @@ async def _fetch_batch_managed_file_content( resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e - ) + verbose_logger.error("Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index daf3b91a26b..3a69911dfe6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1126,11 +1126,13 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch) - cost, usage, models = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") - assert cost == 0.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) - assert models == [] + assert result.cost == 0.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0) + assert result.models == [] + assert result.successful_requests == 0 + assert result.failed_requests == 0 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 7052e050806..b8ed2b04b9d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2633,7 +2633,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -2729,7 +2729,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -2823,7 +2823,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, From ff833a5872cc167a7166a58fdf78aeb2d175b570 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:37:07 -0400 Subject: [PATCH 05/16] fix(batches): address Greptile type-discipline feedback Add Final to the reasoning_tokens local var and type the _batch_cost_result test helper's parameters, per review feedback on PR #37208. --- litellm/batches/batch_utils.py | 4 +- .../proxy_unit_tests/test_check_batch_cost.py | 519 ++++++------------ 2 files changed, 167 insertions(+), 356 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4eb9c7a5dfa..057c9978879 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -165,7 +165,9 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + reasoning_tokens: Final = ( + usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + ) yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 139bf583e6a..ff2dce498f0 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -6,14 +6,24 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, ARN unified_object_id) batches with no managed unified id. """ +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest +if TYPE_CHECKING: + from litellm.batches.batch_utils import BatchCostUsageResult + _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" -def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): +def _batch_cost_result( + cost: float, + usage: dict, + models: list[str], + successful_requests: int = 1, + failed_requests: int = 0, +) -> "BatchCostUsageResult": """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, for mocking it in tests that only care about cost/usage/models.""" from litellm.batches.batch_utils import BatchCostUsageResult @@ -90,9 +100,7 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance( - self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router - ): + def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -104,23 +112,15 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,9 +129,7 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -142,9 +140,7 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -168,18 +164,12 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() @@ -205,9 +195,7 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -215,9 +203,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -228,32 +214,20 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - ) - fallback_where = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ - "where" - ] - ) + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -267,13 +241,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -282,22 +252,16 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -345,9 +309,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -356,15 +318,11 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - "batch_processed" not in update_data - ), "update() must NOT include batch_processed when column is absent" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -440,7 +398,9 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -458,9 +418,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance( - snapshot, MappingProxyType - ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert isinstance(snapshot, MappingProxyType), ( + "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + ) assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -474,13 +434,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -488,21 +444,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -550,9 +500,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -560,15 +508,13 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - update_data["batch_processed"] is True - ), "update() must include batch_processed=True when column is present" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True, ( + "update() must include batch_processed=True when column is present" + ) assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -712,22 +658,16 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" @@ -761,9 +701,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a failed cost tracking attempt must not mark the job processed" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a failed cost tracking attempt must not mark the job processed" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -780,13 +720,9 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -796,31 +732,25 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), f"Expected update() to be called exactly once for a {terminal_status} job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + f"Expected update() to be called exactly once for a {terminal_status} job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == terminal_status - assert ( - update_data["batch_processed"] is True - ), "terminal-status update() must set batch_processed=True so polling stops" + assert update_data["batch_processed"] is True, ( + "terminal-status update() must set batch_processed=True so polling stops" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -855,13 +785,9 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -871,9 +797,7 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file - ) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -882,9 +806,7 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) response = LiteLLMBatch( id="batch-456", @@ -902,9 +824,7 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook await check_batch_cost_instance.check_batch_cost() @@ -955,13 +875,9 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -971,24 +887,18 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{completed_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with patch( "litellm.files.main.afile_content", @@ -996,22 +906,18 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a completed batch with no output file must be marked processed exactly once" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a completed batch with no output file must be marked processed exactly once" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == completed_status - assert ( - update_data["batch_processed"] is True - ), "completed-without-output update() must set batch_processed=True so polling stops" - assert ( - mock_afile_content.await_count == 0 - ), "a batch with no output file must not be billed" - assert ( - mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 - ), "a batch with no output file must not enter the cost-tracking path" + assert update_data["batch_processed"] is True, ( + "completed-without-output update() must set batch_processed=True so polling stops" + ) + assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" + assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( + "a batch with no output file must not enter the cost-tracking path" + ) @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -1022,9 +928,7 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -1032,9 +936,7 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "in_progress" @@ -1060,9 +962,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a non-terminal batch must not be written back (would stop polling prematurely)" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a non-terminal batch must not be written back (would stop polling prematurely)" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -1079,13 +981,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -1093,21 +991,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1155,9 +1047,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1165,20 +1055,16 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_afile_content.await_count == 1 - ), f"{terminal_status} batch with an output file must fetch results and be billed" - mock_logging_obj.async_success_handler.assert_awaited_once() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + assert mock_afile_content.await_count == 1, ( + f"{terminal_status} batch with an output file must fetch results and be billed" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["batch_processed"] is True - assert ( - update_data["status"] == terminal_status - ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert update_data["status"] == terminal_status, ( + f"billed {terminal_status} batch must keep its real terminal status in the DB" + ) @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1195,13 +1081,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1211,23 +1093,17 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"failed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with ( patch( @@ -1248,12 +1124,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a terminal batch with a 404ing output file must be retired, not retried forever" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1266,13 +1140,9 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1281,9 +1151,7 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1294,14 +1162,10 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1316,9 +1180,7 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1361,9 +1223,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1414,9 +1274,7 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = ( - file_object if file_object is not None else _unmanaged_vertex_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1454,9 +1312,7 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with( - "gemini-2.5-flash" - ) + router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1476,9 +1332,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1526,9 +1380,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1536,9 +1388,7 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1562,9 +1412,7 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - '{"id":"8823717160934178816","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1586,9 +1434,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1619,9 +1465,7 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1662,9 +1506,7 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = ( - file_object if file_object is not None else _unmanaged_bedrock_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() return job def _bedrock_deployment(self): @@ -1719,9 +1561,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1759,9 +1599,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1769,9 +1607,7 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1794,13 +1630,9 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"{self._ARN}","status":"completed"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"aws_region_name": "us-east-1"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1816,9 +1648,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1849,9 +1679,7 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1975,9 +1803,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1986,8 +1812,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = ( - lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( + _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2056,9 +1882,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2072,9 +1896,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2082,9 +1904,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth( - api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] - ), + valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), llm_router=None, ) is True @@ -2101,6 +1921,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + + class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2196,9 +2018,7 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=Exception("db down") - ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2294,9 +2114,7 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma( - [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2334,9 +2152,7 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "batch_processed": True - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2389,17 +2205,13 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma( - [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "status": "stale_expired" - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2454,14 +2266,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [ - call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list - ] + retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] assert retired == ["job-no-model", "job-gone"] - assert ( - llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" - ), "the newer healthy batch must still be polled in the same cycle" + assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( + "the newer healthy batch must still be polled in the same cycle" + ) @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 3bfeaa78ddd1bb7cae9ab6576e21a699a466b642 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:39:28 -0400 Subject: [PATCH 06/16] revert: undo accidental whole-file ruff-format of test_check_batch_cost.py The prior commit ran ruff format on the whole file to type the _batch_cost_result helper, reflowing hundreds of unrelated pre-existing lines that were never ruff-format-clean to begin with (confirmed at the PR's base commit, before any of these changes). CI's ruff-format gate only checks litellm/**/*.py, not tests/, so this reformatting served no CI purpose and only bloated the diff. Restores everything except the intended TYPE_CHECKING import and _batch_cost_result annotations. --- .../proxy_unit_tests/test_check_batch_cost.py | 507 ++++++++++++------ 1 file changed, 354 insertions(+), 153 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff2dce498f0..cd7f28007af 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -100,7 +100,9 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): + def check_batch_cost_instance( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -112,15 +114,23 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): + async def test_cleanup_scoped_to_batch_file_purpose( + self, check_batch_cost_instance, mock_prisma_client + ): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + ) stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,7 +139,9 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_confirms_batch_processed_support( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -140,7 +152,9 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_marks_column_absent( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -164,12 +178,18 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): + async def test_find_many_uses_pagination_and_excludes_stale( + self, check_batch_cost_instance, mock_prisma_client + ): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() @@ -195,7 +215,9 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -203,7 +225,9 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + ) assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -214,20 +238,32 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): + async def test_column_absence_cached_across_cycles( + self, check_batch_cost_instance, mock_prisma_client + ): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + ) + fallback_where = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ + "where" + ] + ) assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -241,9 +277,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -252,16 +292,22 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -309,7 +355,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -318,11 +366,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + "batch_processed" not in update_data + ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -398,9 +450,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result( - 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] - ), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -418,9 +468,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance(snapshot, MappingProxyType), ( - "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" - ) + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -434,9 +484,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -444,15 +498,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -500,7 +560,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -508,13 +570,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True, ( - "update() must include batch_processed=True when column is present" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + update_data["batch_processed"] is True + ), "update() must include batch_processed=True when column is present" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -658,16 +722,22 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" @@ -701,9 +771,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a failed cost tracking attempt must not mark the job processed" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a failed cost tracking attempt must not mark the job processed" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -720,9 +790,13 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -732,25 +806,31 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - f"Expected update() to be called exactly once for a {terminal_status} job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), f"Expected update() to be called exactly once for a {terminal_status} job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == terminal_status - assert update_data["batch_processed"] is True, ( - "terminal-status update() must set batch_processed=True so polling stops" - ) + assert ( + update_data["batch_processed"] is True + ), "terminal-status update() must set batch_processed=True so polling stops" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -785,9 +865,13 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -797,7 +881,9 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -806,7 +892,9 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) response = LiteLLMBatch( id="batch-456", @@ -824,7 +912,9 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) await check_batch_cost_instance.check_batch_cost() @@ -875,9 +965,13 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -887,18 +981,24 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{completed_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with patch( "litellm.files.main.afile_content", @@ -906,18 +1006,22 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a completed batch with no output file must be marked processed exactly once" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a completed batch with no output file must be marked processed exactly once" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == completed_status - assert update_data["batch_processed"] is True, ( - "completed-without-output update() must set batch_processed=True so polling stops" - ) - assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" - assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( - "a batch with no output file must not enter the cost-tracking path" - ) + assert ( + update_data["batch_processed"] is True + ), "completed-without-output update() must set batch_processed=True so polling stops" + assert ( + mock_afile_content.await_count == 0 + ), "a batch with no output file must not be billed" + assert ( + mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 + ), "a batch with no output file must not enter the cost-tracking path" @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -928,7 +1032,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -936,7 +1042,9 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "in_progress" @@ -962,9 +1070,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a non-terminal batch must not be written back (would stop polling prematurely)" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a non-terminal batch must not be written back (would stop polling prematurely)" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -981,9 +1089,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -991,15 +1103,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1047,7 +1165,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1055,16 +1175,20 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_afile_content.await_count == 1, ( - f"{terminal_status} batch with an output file must fetch results and be billed" - ) + assert ( + mock_afile_content.await_count == 1 + ), f"{terminal_status} batch with an output file must fetch results and be billed" mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True - assert update_data["status"] == terminal_status, ( - f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["batch_processed"] is True + assert ( + update_data["status"] == terminal_status + ), f"billed {terminal_status} batch must keep its real terminal status in the DB" @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1081,9 +1205,13 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1093,17 +1221,23 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with ( patch( @@ -1124,10 +1258,12 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a terminal batch with a 404ing output file must be retired, not retried forever" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1140,9 +1276,13 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1151,7 +1291,9 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1162,10 +1304,14 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1180,7 +1326,9 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1223,7 +1371,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1274,7 +1424,9 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1312,7 +1464,9 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1332,7 +1486,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1380,7 +1536,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1388,7 +1546,9 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1412,7 +1572,9 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1434,7 +1596,9 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1465,7 +1629,9 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1506,7 +1672,9 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) return job def _bedrock_deployment(self): @@ -1561,7 +1729,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1599,7 +1769,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1607,7 +1779,9 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1630,9 +1804,13 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1648,7 +1826,9 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1679,7 +1859,9 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1803,7 +1985,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1812,8 +1996,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( - _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -1882,7 +2066,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -1896,7 +2082,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -1904,7 +2092,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), llm_router=None, ) is True @@ -1921,8 +2111,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] - - class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2018,7 +2206,9 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2114,7 +2304,9 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2152,7 +2344,9 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "batch_processed": True + } @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2205,13 +2399,17 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2266,11 +2464,14 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] assert retired == ["job-no-model", "job-gone"] - assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( - "the newer healthy batch must still be polled in the same cycle" - ) + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 98447e00436443f26ac5538bb4ad646cc0b32f75 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:58:48 -0400 Subject: [PATCH 07/16] fix(batches): reconcile merge with upstream batch pricing changes - Resolved merge conflicts in _handle_completed_batch and CheckBatchCost._track_completed_batch_cost, keeping BatchCostUsageResult while adopting upstream's model_info threading and improved deployment-pricing lookup - Fixed two upstream tests and one mock that still expected the old tuple(cost, usage, models) return shape - Suppressed the one new LIT002 violation from the empty-output-file BatchCostUsageResult literal --- litellm/batches/batch_utils.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 20 +++++++++---------- .../test_litellm_logging.py | 12 +++++++++-- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index bae8e198223..d0874b5aff5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -86,7 +86,7 @@ async def _handle_completed_batch( return BatchCostUsageResult( cost=0.0, usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), - models=[], + 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 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 0adebf2506b..64b3df4180b 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1501,24 +1501,24 @@ async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monke monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - cost, usage, _ = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name="bedrock/global.anthropic.claude-sonnet-4-6", ) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (1800, 1000, 2800) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800) # 3e-06 / 1.5e-05 on-demand, halved for batch. - assert cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) + assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) # The response model alone cannot price a bedrock batch: this is the $0 bug. - zero_cost, zero_usage, _ = await bu._handle_completed_batch( + zero_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name=None, ) - assert zero_cost == 0.0 - assert zero_usage.total_tokens == 2800 + assert zero_result.cost == 0.0 + assert zero_result.usage.total_tokens == 2800 @pytest.mark.asyncio @@ -1531,7 +1531,7 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - free_cost, _, _ = await bu._handle_completed_batch( + free_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", @@ -1542,15 +1542,15 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> "output_cost_per_token_batches": 0.0, }, ) - assert free_cost == 0.0 + assert free_result.cost == 0.0 - billed_cost, _, _ = await bu._handle_completed_batch( + billed_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", model_info=None, ) - assert billed_cost > 0.0 + assert billed_result.cost > 0.0 # =========================================================================== # diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0d54680fa81..e46057fe11a 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -586,9 +586,17 @@ class TestRetrieveBatchCostPassesModelIdentity: captured: dict[str, object] = {} - async def fake_handle_completed_batch(**kwargs: object) -> tuple[float, Usage, list[str]]: + from litellm.batches.batch_utils import BatchCostUsageResult + + async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: captured.update(kwargs) - return 1.25, Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), ["m"] + return BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["m"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) From d989f172e9d78ff63c5867f45850d5065b1ad7c8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 20:32:11 -0400 Subject: [PATCH 08/16] fix(batches): revert invalid Final on a loop-scoped local basedpyright rejects a Final variable assigned inside a loop body (reportGeneralTypeIssues); the type-discipline checker doesn't flag this line without Final either, so the annotation only bought a basedpyright budget regression. --- litellm/batches/batch_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index d0874b5aff5..4e62c4a0130 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -170,9 +170,7 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens: Final = ( - usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None - ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, From 5a7edc9c77836f8f47634d8d47719bb44059a1fc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:26:52 -0400 Subject: [PATCH 09/16] fix(spend-tracking): make the new batch count keys writable SpendLogsMetadata is built by assigning each key in turn, so ReadOnly on the two new ones breached the basedpyright reportTypedDictNotRequiredAccess ceiling. Every sibling key in this TypedDict is writable for the same reason. --- litellm/proxy/_types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bf0dbc97482..377eb7dce0e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3528,8 +3528,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: ReadOnly[int | None] - batch_failed_requests: ReadOnly[int | None] + batch_successful_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict + batch_failed_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None From 4c2f0f3632c39ae8015f136b02eb002ee24ad248 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:34:16 -0400 Subject: [PATCH 10/16] test(batches): cover decoding a model-encoded error file id The error file now resolves through _provider_output_file_id like the output file does. Sending the encoded id straight to the provider 404s, and the swallowed fetch failure would silently report zero failures. --- .../test_litellm/batches/test_batch_utils.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 1540a587349..c86c7c4df03 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1164,6 +1164,53 @@ async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): assert result.failed_requests == 1 +@pytest.mark.asyncio +async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch): + """A model-encoded error file id must be decoded to the raw provider id before + the fetch, exactly like the output file id. Sending the encoded id straight to + the provider 404s, and the swallowed fetch failure silently reports 0 failures.""" + import base64 + + from litellm.types.llms.openai import Batch + + provider_error_file_id = "file-real-error-id" + encoded_error_file_id = "file-" + base64.urlsafe_b64encode( + f"litellm:{provider_error_file_id};model,model-abc".encode() + ).decode().rstrip("=") + + requested_file_ids = [] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))]) + + async def fake_afile_content(**kw): + requested_file_ids.append(kw["file_id"]) + return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})() + + 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(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id=encoded_error_file_id, + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert requested_file_ids == [provider_error_file_id] + assert result.failed_requests == 1 + + @pytest.mark.asyncio async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] From fee86197081f1fef09b0d3071c893e20036e45a1 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 11:26:39 -0700 Subject: [PATCH 11/16] test: add batch request count keys to gcs pub sub spend logs fixture --- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 5789f19aa55..1838fb16e91 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From 4ef012627ddbbc08d87fc85f01367bc66029b7ae Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 12:06:42 -0700 Subject: [PATCH 12/16] fix: count error-file failures in the batch cost poller path --- .../proxy/common_utils/check_batch_cost.py | 27 +++- .../proxy_unit_tests/test_check_batch_cost.py | 125 ++++++++++++++++++ 2 files changed, 149 insertions(+), 3 deletions(-) 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 163b1592dd9..d34635fa253 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,9 +2,10 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ +from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -626,6 +627,7 @@ 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, ) @@ -761,12 +763,31 @@ class CheckBatchCost: model_id=model_id, deployment_model=litellm_model_name, ) - batch_result = await calculate_batch_cost_and_usage( + batch_file_provider: Final = cast( + Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], llm_provider + ) + output_file_result: Final = await calculate_batch_cost_and_usage( file_content_dictionary=file_content_as_dict, - custom_llm_provider=llm_provider, # type: ignore + custom_llm_provider=batch_file_provider, model_name=model_name, model_info=deployment_model_info, ) + error_file_failed_requests: Final = await _count_error_file_failed_requests( + response, + custom_llm_provider=batch_file_provider, + litellm_params={ + **credentials, + "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), + }, + ) + batch_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, + ) + ) logging_obj = LiteLLMLogging( model=batch_result.models[0], messages=[{"role": "user", "content": ""}], diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 7f4fac4e6e9..bf5add54139 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1373,6 +1373,131 @@ class TestCheckBatchCost: update_data["status"] == terminal_status ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + @pytest.mark.asyncio + async def test_error_file_failures_add_to_failed_request_count( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """OpenAI-shaped providers report per-request failures only in a separate + error file. The poller prices from the output file, so without also counting + the error file's lines, batch_failed_requests on the spend log undercounts: + regression test for the poller path merging error-file failures. + """ + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + mock_job = MagicMock() + mock_job.id = "job-error-file-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.error_file_id = "file-error-456" + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-4" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + output_file_content = MagicMock() + output_file_content.content = b'{"id":"req-1"}' + error_file_content = MagicMock() + error_file_content.content = ( + b'{"id":"err-1","error":{"message":"rejected"}}\n' + b'{"id":"err-2","error":{"message":"rejected"}}\n\n' + ) + + def _file_content_for(**kwargs): + if kwargs["file_id"] == "file-error-456": + return error_file_content + return output_file_content + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + side_effect=_file_content_for, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=_batch_cost_result( + 0.01, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["gpt-4"], + successful_requests=3, + failed_requests=1, + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("gpt-4", "openai", None, None), + ), + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + assert mock_afile_content.await_count == 2, ( + "the poller must fetch the error file in addition to the output file" + ) + fetched_file_ids = { + call.kwargs["file_id"] for call in mock_afile_content.await_args_list + } + assert fetched_file_ids == {"file-output-123", "file-error-456"} + + mock_logging_obj.async_success_handler.assert_awaited_once() + handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs + assert handler_kwargs["batch_successful_requests"] == 3 + assert handler_kwargs["batch_failed_requests"] == 3, ( + "2 error-file lines must add to the output file's 1 failed request" + ) + assert handler_kwargs["batch_cost"] == 0.01 + @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router From 99884f0eaa42cd5d8dfd98ee8e411d1e203b88eb Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 12:53:52 -0700 Subject: [PATCH 13/16] test: fake the provider file boundary in the poller error-file regression test --- .../proxy_unit_tests/test_check_batch_cost.py | 338 +++++++----------- 1 file changed, 122 insertions(+), 216 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index bf5add54139..4f9fcf1952b 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1382,37 +1382,31 @@ class TestCheckBatchCost: the error file's lines, batch_failed_requests on the spend log undercounts: regression test for the poller path merging error-file failures. """ + import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + import httpx + import respx + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-error-file-1" - mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() mock_job.created_by = "user-1" - - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" mock_response.error_file_id = "file-error-456" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) - + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1420,83 +1414,73 @@ class TestCheckBatchCost: mock_deployment.model_info.model_dump.return_value = {} mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) - output_file_content = MagicMock() - output_file_content.content = b'{"id":"req-1"}' - error_file_content = MagicMock() - error_file_content.content = ( - b'{"id":"err-1","error":{"message":"rejected"}}\n' - b'{"id":"err-2","error":{"message":"rejected"}}\n\n' + succeeded_line = json.dumps( + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + }, + "error": None, + } + ) + rejected_line = json.dumps( + { + "custom_id": "req-2", + "response": { + "status_code": 400, + "body": {"error": {"message": "bad request"}}, + }, + "error": None, + } + ) + error_file_lines = "\n".join( + json.dumps({"custom_id": custom_id, "error": {"message": "rejected"}}) for custom_id in ("req-3", "req-4") ) - def _file_content_for(**kwargs): - if kwargs["file_id"] == "file-error-456": - return error_file_content - return output_file_content - - decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" - with ( - patch( - "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", - side_effect=[decoded_id, None, None], - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", - return_value="model-123", - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", - return_value="batch-456", - ), - patch( - "litellm.files.main.afile_content", - new_callable=AsyncMock, - side_effect=_file_content_for, - ) as mock_afile_content, - patch( - "litellm.batches.batch_utils._get_file_content_as_dictionary", - return_value=[{"id": "req-1"}], - ), - patch( - "litellm.batches.batch_utils.calculate_batch_cost_and_usage", - new_callable=AsyncMock, - return_value=_batch_cost_result( - 0.01, - {"prompt_tokens": 10, "completion_tokens": 5}, - ["gpt-4"], - successful_requests=3, - failed_requests=1, - ), - ), - patch( - "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", - return_value=("gpt-4", "openai", None, None), - ), - patch( + respx.mock(assert_all_called=True) as provider, + patch( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs "litellm.litellm_core_utils.litellm_logging.Logging" ) as mock_logging_cls, ): + 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( + return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode()) + ) mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() mock_logging_cls.return_value = mock_logging_obj await check_batch_cost_instance.check_batch_cost() - assert mock_afile_content.await_count == 2, ( - "the poller must fetch the error file in addition to the output file" - ) - fetched_file_ids = { - call.kwargs["file_id"] for call in mock_afile_content.await_args_list - } - assert fetched_file_ids == {"file-output-123", "file-error-456"} - mock_logging_obj.async_success_handler.assert_awaited_once() handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs - assert handler_kwargs["batch_successful_requests"] == 3 + assert handler_kwargs["batch_successful_requests"] == 1 assert handler_kwargs["batch_failed_requests"] == 3, ( - "2 error-file lines must add to the output file's 1 failed request" + "2 error-file lines must add to the output file's 1 rejected request" ) - assert handler_kwargs["batch_cost"] == 0.01 + assert handler_kwargs["batch_models"] == ["gpt-4"] + assert handler_kwargs["batch_usage"].total_tokens == 15 + assert handler_kwargs["batch_cost"] > 0 @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1513,13 +1497,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1529,23 +1509,17 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"failed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with ( patch( @@ -1566,12 +1540,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a terminal batch with a 404ing output file must be retired, not retried forever" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1584,13 +1556,9 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1599,9 +1567,7 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1612,14 +1578,10 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1634,9 +1596,7 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1679,9 +1639,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1732,9 +1690,7 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = ( - file_object if file_object is not None else _unmanaged_vertex_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1772,9 +1728,7 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with( - "gemini-2.5-flash" - ) + router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1794,9 +1748,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1844,9 +1796,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1854,9 +1804,7 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1880,9 +1828,7 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - '{"id":"8823717160934178816","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1904,9 +1850,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1937,9 +1881,7 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1980,9 +1922,7 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = ( - file_object if file_object is not None else _unmanaged_bedrock_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() return job def _bedrock_deployment(self): @@ -2037,9 +1977,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -2077,9 +2015,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -2087,9 +2023,7 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -2112,13 +2046,9 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"{self._ARN}","status":"completed"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"aws_region_name": "us-east-1"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -2134,9 +2064,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -2167,9 +2095,7 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -2293,9 +2219,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -2304,8 +2228,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = ( - lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( + _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2374,9 +2298,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2390,9 +2312,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2400,9 +2320,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth( - api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] - ), + valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), llm_router=None, ) is True @@ -2419,6 +2337,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + + class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2514,9 +2434,7 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=Exception("db down") - ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2612,9 +2530,7 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma( - [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2652,9 +2568,7 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "batch_processed": True - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2707,17 +2621,13 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma( - [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "status": "stale_expired" - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2772,14 +2682,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [ - call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list - ] + retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] assert retired == ["job-no-model", "job-gone"] - assert ( - llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" - ), "the newer healthy batch must still be polled in the same cycle" + assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( + "the newer healthy batch must still be polled in the same cycle" + ) @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): @@ -2808,6 +2715,7 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2824,8 +2732,12 @@ class _FakeManagedObjectRow: self.request_tags = None self.created_at = 1700000000 self.file_object = json.dumps( - {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1", - "output_file_id": _CLAIM_OUTPUT_FILE_ID} + { + "id": "batch-456", + "status": "in_progress", + "input_file_id": "file-input-1", + "output_file_id": _CLAIM_OUTPUT_FILE_ID, + } ) @@ -2925,9 +2837,7 @@ class TestMultiPodBatchCostClaim: router = MagicMock() router.aretrieve_batch = AsyncMock(return_value=response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) router.get_deployment = MagicMock(return_value=deployment) return router @@ -3087,9 +2997,7 @@ class TestMultiPodBatchCostClaim: await asyncio.Event().wait() with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj: - interrupted = asyncio.create_task( - self._instance(prisma, self._router()).check_batch_cost() - ) + interrupted = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) await asyncio.wait_for(reached_fetch.wait(), timeout=5) assert row.batch_processed is False, "an in-flight costing must not mark the row processed" interrupted.cancel() @@ -3124,9 +3032,7 @@ class TestMultiPodBatchCostClaim: await finish_fetch.wait() with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt): - costing = asyncio.create_task( - self._instance(prisma, self._router()).check_batch_cost() - ) + costing = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) await asyncio.wait_for(reached_fetch.wait(), timeout=5) with pytest.raises(HTTPException) as blocked: From 38145c20826697c13e3b1955e41d53061376eb94 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 13:17:43 -0700 Subject: [PATCH 14/16] test: undo the drive-by reformat below the poller error-file regression test --- .../proxy_unit_tests/test_check_batch_cost.py | 196 ++++++++++++------ 1 file changed, 137 insertions(+), 59 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 4f9fcf1952b..757762eac87 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1497,9 +1497,13 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1509,17 +1513,23 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with ( patch( @@ -1540,10 +1550,12 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a terminal batch with a 404ing output file must be retired, not retried forever" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1556,9 +1568,13 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1567,7 +1583,9 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1578,10 +1596,14 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1596,7 +1618,9 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1639,7 +1663,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1690,7 +1716,9 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1728,7 +1756,9 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1748,7 +1778,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1796,7 +1828,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1804,7 +1838,9 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1828,7 +1864,9 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1850,7 +1888,9 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1881,7 +1921,9 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1922,7 +1964,9 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) return job def _bedrock_deployment(self): @@ -1977,7 +2021,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -2015,7 +2061,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -2023,7 +2071,9 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -2046,9 +2096,13 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -2064,7 +2118,9 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -2095,7 +2151,9 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -2219,7 +2277,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -2228,8 +2288,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( - _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2298,7 +2358,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2312,7 +2374,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2320,7 +2384,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), llm_router=None, ) is True @@ -2337,8 +2403,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] - - class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2434,7 +2498,9 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2530,7 +2596,9 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2568,7 +2636,9 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "batch_processed": True + } @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2621,13 +2691,17 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2682,11 +2756,14 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] assert retired == ["job-no-model", "job-gone"] - assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( - "the newer healthy batch must still be polled in the same cycle" - ) + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): @@ -2715,7 +2792,6 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() - class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2732,12 +2808,8 @@ class _FakeManagedObjectRow: self.request_tags = None self.created_at = 1700000000 self.file_object = json.dumps( - { - "id": "batch-456", - "status": "in_progress", - "input_file_id": "file-input-1", - "output_file_id": _CLAIM_OUTPUT_FILE_ID, - } + {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1", + "output_file_id": _CLAIM_OUTPUT_FILE_ID} ) @@ -2837,7 +2909,9 @@ class TestMultiPodBatchCostClaim: router = MagicMock() router.aretrieve_batch = AsyncMock(return_value=response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) router.get_deployment = MagicMock(return_value=deployment) return router @@ -2997,7 +3071,9 @@ class TestMultiPodBatchCostClaim: await asyncio.Event().wait() with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj: - interrupted = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) + interrupted = asyncio.create_task( + self._instance(prisma, self._router()).check_batch_cost() + ) await asyncio.wait_for(reached_fetch.wait(), timeout=5) assert row.batch_processed is False, "an in-flight costing must not mark the row processed" interrupted.cancel() @@ -3032,7 +3108,9 @@ class TestMultiPodBatchCostClaim: await finish_fetch.wait() with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt): - costing = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) + costing = asyncio.create_task( + self._instance(prisma, self._router()).check_batch_cost() + ) await asyncio.wait_for(reached_fetch.wait(), timeout=5) with pytest.raises(HTTPException) as blocked: From 128d7e52784b4a8fee4b171e7528ff6841cd2a08 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 14:09:18 -0700 Subject: [PATCH 15/16] refactor(batches): make count_error_file_failed_requests public for the poller import --- .../proxy/common_utils/check_batch_cost.py | 4 ++-- litellm/batches/batch_utils.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) 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 d34635fa253..3b09dc9272e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -627,7 +627,7 @@ class CheckBatchCost: later poll. """ from litellm.batches.batch_utils import ( - _count_error_file_failed_requests, + count_error_file_failed_requests, _get_file_content_as_dictionary, calculate_batch_cost_and_usage, ) @@ -772,7 +772,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_failed_requests: Final = await count_error_file_failed_requests( response, custom_llm_provider=batch_file_provider, litellm_params={ diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index ea2a95c1717..3831f57a10d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -89,13 +89,13 @@ async def _handle_completed_batch( 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( + failed_requests=await count_error_file_failed_requests( batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ), ) 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( + error_file_failed_requests: Final = await count_error_file_failed_requests( batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ) @@ -450,7 +450,7 @@ async def _fetch_batch_output_file_content( ) -async def _count_error_file_failed_requests( +async def count_error_file_failed_requests( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], litellm_params: dict | None, From 2963b47cda666efc363cd497463e827cfe4d4706 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 14:09:33 -0700 Subject: [PATCH 16/16] test: patch the Logging handler instead of the class in the poller error-file test --- tests/proxy_unit_tests/test_check_batch_cost.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 757762eac87..ff5e8f89d64 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1388,6 +1388,8 @@ class TestCheckBatchCost: import httpx import respx + from litellm.litellm_core_utils.litellm_logging import Logging + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1456,9 +1458,9 @@ class TestCheckBatchCost: with ( respx.mock(assert_all_called=True) as provider, - patch( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs + Logging, "async_success_handler", new_callable=AsyncMock + ) as success_handler, ): 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()) @@ -1466,14 +1468,11 @@ class TestCheckBatchCost: 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()) ) - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - mock_logging_cls.return_value = mock_logging_obj - await check_batch_cost_instance.check_batch_cost() - mock_logging_obj.async_success_handler.assert_awaited_once() - handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs + spend_log_calls = [call.kwargs for call in success_handler.await_args_list if "batch_cost" in call.kwargs] + assert len(spend_log_calls) == 1 + handler_kwargs = spend_log_calls[0] assert handler_kwargs["batch_successful_requests"] == 1 assert handler_kwargs["batch_failed_requests"] == 3, ( "2 error-file lines must add to the output file's 1 rejected request"