From 99341ad05de23a3689d36d800ff4939885f7a3bb Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 15:52:33 -0700 Subject: [PATCH] fix(batches): handle empty models when logging batch cost (#47155) Backport of #47155 to stable/1.106.x. Cherry-picked from 4143e0869eff3db190a5e8f73069bfe2c9428303 (main). --- .../proxy/common_utils/check_batch_cost.py | 3 +- .../common_utils/test_check_batch_cost.py | 118 +++++++++++++++++- 2 files changed, 119 insertions(+), 2 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 2b2d0a82431..1aa35c229be 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -890,8 +890,9 @@ class CheckBatchCost: failed_requests=output_file_result.failed_requests + error_file_failed_requests, ) ) + logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id logging_obj = LiteLLMLogging( - model=batch_result.models[0], + model=logging_model, messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 58d6a44e4a4..3cdc6598c69 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -9,7 +9,7 @@ ARN unified_object_id) batches with no managed unified id. import asyncio import json from contextlib import contextmanager -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -3055,6 +3055,122 @@ class TestMultiPodBatchCostClaim: ): yield logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("deployment_model", "successes", "failures", "expected_model"), + ( + ("openai/", 0, 2, "openai/"), + ("", 0, 2, "model-123"), + ("openai/", 2, 1, "batch-test-model"), + ("openai/batch-test-model", 2, 1, "batch-test-model"), + ("openai/batch-test-model", 2, 0, "batch-test-model"), + ), + ) + async def test_batch_accounting_preserves_models_and_mixed_results( + self, deployment_model: str, successes: int, failures: int, expected_model: str + ) -> None: + import base64 + from datetime import datetime + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + from litellm.types.utils import LiteLLMBatch, SpecialEnums + + input_rate: Final = 0.000002 + output_rate: Final = 0.000006 + litellm.register_model( + { + "batch-test-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": input_rate, + "output_cost_per_token_batches": output_rate, + } + } + ) + success: Final = { + "response": { + "status_code": 200, + "body": { + "model": "batch-test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120}, + }, + }, + "error": None, + } + failure: Final = {"response": None, "error": {"code": "batch_expired", "message": "expired"}} + rows: Final = (success,) * successes + (failure,) * failures + content: Final = b"\n".join(json.dumps(row).encode() for row in rows) + + class BatchSpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.calls: tuple[dict[str, object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs["call_type"] == "aretrieve_batch": + self.calls += (kwargs,) + + callback: Final = BatchSpendRecorder() + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + row: Final = _FakeManagedObjectRow() + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + prisma: Final = self._prisma(row, []) + router: Final = self._router() + router.get_deployment.return_value = Deployment( + model_name="batch-alias", + litellm_params=LiteLLM_Params(model=deployment_model, custom_llm_provider="openai"), + model_info=ModelInfo(id="model-123"), + ) + router.aretrieve_batch.return_value = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + output_file_id=_CLAIM_OUTPUT_FILE_ID, + object="batch", + status="completed", + ) + router.get_deployment_credentials_with_provider.return_value = { + "api_key": "synthetic-key", + "custom_llm_provider": "openai", + } + with respx.mock: + respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( + return_value=httpx.Response(200, content=content) + ) + instance: Final = self._instance(prisma, router) + await instance.check_batch_cost() + await instance.check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + router.aretrieve_batch.assert_awaited_once() + assert len(callback.calls) == 1 + payload: Final = callback.calls[0]["standard_logging_object"] + assert payload["model"] == expected_model + assert payload["response_cost"] == pytest.approx(successes * (100 * input_rate + 20 * output_rate)) + assert payload["prompt_tokens"] == successes * 100 + assert payload["completion_tokens"] == successes * 20 + assert payload["hidden_params"]["batch_successful_requests"] == successes + assert payload["hidden_params"]["batch_failed_requests"] == failures + assert payload["hidden_params"]["batch_models"] == ( + ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes + ) + @staticmethod def _claim_calls(prisma) -> list: return [