From 2f3d44b4694a3dc40df39d24cb03c0c7f530ffd5 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 16:32:09 -0700 Subject: [PATCH] fix(batch): preserve recoverable costs and prevent poll starvation (#47190) Backport of #47190 to stable/1.103.x. Cherry-picked from 0ede0e10e7ee694c46f14afd1606885d8c369342 (main). Ported from the stable/1.104.x hand-port (3b7b5ab319). Carries the small typed table helpers that port relies on (_ManagedObjectRow protocol and the _managed_object_table/_user_table/_token_table/_team_table accessors, from the stable/1.104.x base). The poll-starvation integration test is dropped since its base and harness (#42622, integration/_support/prisma_clients) are not on this line. The 404-recovery unit test resets the litellm callback lists so a callback left by an earlier test cannot shadow it. --- .../proxy/common_utils/check_batch_cost.py | 171 ++++++++++---- .../proxy_unit_tests/test_check_batch_cost.py | 210 ++++++++++++------ 2 files changed, 263 insertions(+), 118 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 e7acb4c0b86..95d2f6d3856 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,11 +3,12 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ import json -from collections.abc import Sequence +import re +from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast +from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -23,12 +24,18 @@ if TYPE_CHECKING: from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import LiteLLM_ManagedObjectTable from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.repositories.prisma_protocols import TableActions from litellm.router import Router from litellm.types.router import Deployment from litellm.types.utils import LiteLLMBatch +def _error_names_resource(error: Exception, resource_id: str) -> bool: + return resource_id in re.findall(r"""[^\s"'`]+""", str(error)) + + CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" +RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", @@ -44,6 +51,48 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( ) +class _ManagedObjectRow(Protocol): + @property + def created_at(self) -> datetime: ... + + @property + def updated_at(self) -> datetime: ... + + @property + def id(self) -> str: ... + + @property + def unified_object_id(self) -> str: ... + + @property + def created_by(self) -> str | None: ... + + @property + def file_object(self) -> object: ... + + +def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]": + table: Final[TableActions[_ManagedObjectRow]] = prisma_client.db.litellm_managedobjecttable + return table + + +def _user_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_UserTable]": + table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = prisma_client.db.litellm_usertable + return table + + +def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]": + table: Final[TableActions[prisma_models.LiteLLM_VerificationToken]] = ( + prisma_client.db.litellm_verificationtoken + ) + return table + + +def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]": + table: Final[TableActions[prisma_models.LiteLLM_TeamTable]] = prisma_client.db.litellm_teamtable + return table + + class CheckBatchCost: def __init__( self, @@ -63,6 +112,11 @@ class CheckBatchCost: # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True self.batch_processed_support_confirmed: bool = False + self._poll_cursor: tuple[datetime, str] | None = None + self._poll_before: datetime = datetime.now(timezone.utc) + self._arrival_cursor: tuple[datetime, str] | None = None + self._arrival_before: datetime = self._poll_before + self._poll_arrivals: bool = False @staticmethod def _is_missing_batch_processed_column_error(err: Exception) -> bool: @@ -244,29 +298,9 @@ class CheckBatchCost: f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" ) - if not self._has_batch_processed_column: - return - - # A row already in a terminal status is never rewritten by the sweep above, so - # without this it keeps a poll-page slot forever and starves newer batches. - retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={ - "file_purpose": "batch", - "batch_processed": False, - "status": {"in": ["complete", "completed"]}, - "created_at": {"lt": cutoff}, - }, - data={"batch_processed": True}, - ) - if retired > 0: - verbose_proxy_logger.warning( - f"CheckBatchCost: gave up on {retired} completed managed objects older than " - f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed" - ) - - async def _fallback_find_jobs(self) -> list: + async def _fallback_find_jobs(self) -> "Sequence[_ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" - return await self.prisma_client.db.litellm_managedobjecttable.find_many( + return await self._find_jobs( where={ "file_purpose": "batch", "status": { @@ -280,31 +314,76 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) - async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None: + async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]: + arrivals: Final = self._poll_arrivals + page: Final = await self._find_job_page(where, arrivals=arrivals) + self._poll_arrivals = not arrivals if page else arrivals + if page: + return page + return await self._find_job_page(where, arrivals=not arrivals) + + async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[_ManagedObjectRow]: + cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor + if cursor is None: + if arrivals: + self._arrival_before = datetime.now(timezone.utc) + else: + self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW + before: Final = self._arrival_before if arrivals else self._poll_before + page: Final = await _managed_object_table(self.prisma_client).find_many( + where={ + **where, + "created_at": { + "lte": before, + **({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}), + }, + **( + { + "OR": [ + {"created_at": {"gt": cursor[0]}}, + {"created_at": cursor[0], "id": {"gt": cursor[1]}}, + ] + } + if cursor is not None + else {} + ), + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order=[{"created_at": "asc"}, {"id": "asc"}], + ) + next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None + if arrivals: + self._arrival_cursor = next_cursor + else: + self._poll_cursor = next_cursor + return page + + def _unprocessed_snapshot(self, job: _ManagedObjectRow) -> Mapping[str, object]: + return { + "id": job.id, + "updated_at": job.updated_at, + **({"batch_processed": False} if self._has_batch_processed_column else {}), + } + + async def _retire_job(self, job: "_ManagedObjectRow", reason: str) -> None: """ Take a row that can never be costed out of the poll page. Leaving it selectable would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and once enough such rows accumulate no newer batch is ever reached. Older schemas without batch_processed can only be excluded through the status filter. """ - data: Final = ( - {"batch_processed": True} - if self._has_batch_processed_column - else {"status": "stale_expired"} - ) + data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"} try: - await self.prisma_client.db.litellm_managedobjecttable.update( - where={"id": job.id}, + retired: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=data, ) except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}" - ) + verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}") + return + if retired == 0: return verbose_proxy_logger.warning( f"CheckBatchCost: job {job.id} can never be costed ({reason}), " @@ -395,13 +474,13 @@ class CheckBatchCost: from litellm.exceptions import NotFoundError - return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id) def _batch_deployment_exists(self, model_id: str) -> bool: """A 404 only proves the batch is gone when it came from the batch's own deployment. Once that deployment leaves the router, default fallbacks can silently send the retrieve to a provider that never saw the batch, so its - 404 must not retire the row; the staleness sweep bounds it instead.""" + 404 must not retire a potentially recoverable row.""" return self.llm_router.get_deployment(model_id=model_id) is not None @staticmethod @@ -416,7 +495,7 @@ class CheckBatchCost: if not output_file_id: return False - return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error) + return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id) async def _finalize_unbilled_terminal_job( self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch" @@ -438,15 +517,17 @@ class CheckBatchCost: db_batch_object=job, unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), ) - update_data: Final[dict] = { + update_data: Final[dict[str, object]] = { "status": response.status, "file_object": response.model_dump_json(), **({"batch_processed": True} if self._has_batch_processed_column else {}), } - await self.prisma_client.db.litellm_managedobjecttable.update( - where={"id": job.id}, + finalized: Final = await _managed_object_table(self.prisma_client).update_many( + where=self._unprocessed_snapshot(job), data=update_data, ) + if finalized == 0: + return verbose_proxy_logger.info( f"CheckBatchCost: marked job {job.id} as {response.status} in DB" ) @@ -954,7 +1035,7 @@ class CheckBatchCost: # endpoint may transition a batch to "complete" before # CheckBatchCost runs. The batch_processed=False filter # already prevents reprocessing finished batches. - jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( + jobs = await self._find_jobs( where={ "file_purpose": "batch", "batch_processed": False, @@ -967,8 +1048,6 @@ class CheckBatchCost: ] }, }, - take=MAX_OBJECTS_PER_POLL_CYCLE, - order={"created_at": "asc"}, ) self.batch_processed_support_confirmed = True except Exception as query_err: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 6f093e6592d..a7886229cbd 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -9,6 +9,9 @@ ARN unified_object_id) batches with no managed unified id. import asyncio import json from contextlib import contextmanager +from copy import copy +from datetime import datetime, timezone +from hashlib import sha256 from typing import TYPE_CHECKING, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -203,7 +206,7 @@ class TestCheckBatchCost: find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args assert find_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE - assert find_call[1]["order"] == {"created_at": "asc"} + assert find_call[1]["order"] == [{"created_at": "asc"}, {"id": "asc"}] not_in = find_call[1]["where"]["status"]["not_in"] assert "stale_expired" in not_in # "complete"/"completed" are intentionally NOT excluded from the @@ -228,15 +231,13 @@ class TestCheckBatchCost: ) # 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"), []] + side_effect=[Exception("column batch_processed does not exist"), [], []] ) await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - ) - assert len(calls) == 2 + calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + assert len(calls) == 3 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where assert "stale_expired" in fallback_where["status"]["not_in"] @@ -263,15 +264,8 @@ class TestCheckBatchCost: 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 == 2 + fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -770,7 +764,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "Expected update() to be called exactly once for the completed job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -903,7 +897,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -1019,12 +1013,10 @@ class TestCheckBatchCost: 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_many.call_count == 2, ( + f"Expected update() to be called exactly once for a {terminal_status} job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == terminal_status assert ( update_data["batch_processed"] is True @@ -1130,9 +1122,13 @@ class TestCheckBatchCost: assert store_call.kwargs["user_api_key_dict"].user_id == "user-1" assert store_call.kwargs["user_api_key_dict"].team_id == "team-1" - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_call = mock_prisma_client.db.litellm_managedobjecttable.update.call_args - assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1"} + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args + assert update_call.kwargs["where"] == { + "id": "job-terminal-mint-1", + "batch_processed": False, + "updated_at": mock_job.updated_at, + } update_data = update_call.kwargs["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True @@ -1207,12 +1203,10 @@ 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_many.call_count == 2, ( + "a completed batch with no output file must be marked processed exactly once" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == completed_status assert ( update_data["batch_processed"] is True @@ -1454,7 +1448,7 @@ class TestCheckBatchCost: 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_many.call_count >= 3 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( @@ -1637,12 +1631,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_many.call_count == 2, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -2055,7 +2047,7 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2286,7 +2278,7 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2866,9 +2858,10 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() llm_router.aretrieve_batch.assert_not_awaited() - prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - call = prisma.db.litellm_managedobjecttable.update.call_args[1] - assert call["where"] == {"id": "job-no-model"} + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + call = prisma.db.litellm_managedobjecttable.update_many.call_args[1] + assert call["where"]["id"] == "job-no-model" + assert call["where"]["batch_processed"] is False assert call["data"] == {"batch_processed": True} @pytest.mark.asyncio @@ -2896,10 +2889,8 @@ 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 - } + prisma.db.litellm_managedobjecttable.update_many.assert_awaited() + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2960,26 +2951,7 @@ class TestPollPageStarvation: await instance.check_batch_cost() - 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): - """A row already in a terminal status is never rewritten by the staleness sweep, so - it needs its own bound or it starves newer batches indefinitely.""" - prisma = self._prisma([]) - - await self._instance(prisma, MagicMock()).check_batch_cost() - - calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list - assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep" - where = calls[1][1]["where"] - assert where["file_purpose"] == "batch" - assert where["batch_processed"] is False - assert where["status"] == {"in": ["complete", "completed"]} - assert "created_at" in where - assert calls[1][1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_newer_batch_is_polled_once_dead_rows_are_retired(self): @@ -3019,7 +2991,8 @@ class TestPollPageStarvation: retired = [ call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list + for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list + if "id" in call[1]["where"] ] assert retired == ["job-no-model", "job-gone"] assert ( @@ -3067,7 +3040,8 @@ class _FakeManagedObjectRow: self.team_id = None self.api_key = None self.request_tags = None - self.created_at = 1700000000 + self.created_at = datetime(9999, 1, 1, tzinfo=timezone.utc) + self.updated_at = datetime(2000, 1, 1, tzinfo=timezone.utc) self.file_object = json.dumps( {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1", "output_file_id": _CLAIM_OUTPUT_FILE_ID} @@ -3093,7 +3067,9 @@ class _FakeManagedObjectTable: def _matches(self, where: dict) -> bool: for key, value in where.items(): if key == "created_at": - return False + if "lt" in value and not self.row.created_at < value["lt"]: + return False + continue if key == "status": if self.row.status in value.get("not_in", []): return False @@ -3106,7 +3082,9 @@ class _FakeManagedObjectTable: async def _update_many(self, *, where: dict, data: dict) -> int: if not self._matches(where): return 0 - if "batch_processed" in where: + if "updated_at" in where: + self.journal.append("finalize") + elif "batch_processed" in where: self.journal.append("claim" if data.get("batch_processed") else "release") for key, value in data.items(): setattr(self.row, key, value) @@ -3364,6 +3342,94 @@ class TestMultiPodBatchCostClaim: ["batch-test-model"] if deployment_model == "openai/batch-test-model" else ["batch-test-model"] * successes ) + @pytest.mark.asyncio + @pytest.mark.parametrize("error_source", ("batch", "output")) + async def test_a_404_for_another_identifier_preserves_recovery( + self, monkeypatch: pytest.MonkeyPatch, error_source: str + ) -> None: + import base64 + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.types.utils import SpecialEnums + + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + + class SpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.costs: tuple[float, ...] = () + + 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": + cost: Final = kwargs["response_cost"] + assert isinstance(cost, float) + self.costs += (cost,) + + recorder: Final = SpendRecorder() + litellm.logging_callback_manager.add_litellm_async_success_callback(recorder) + litellm.register_model( + { + "batch-recovery-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": 0.01, + "output_cost_per_token_batches": 0.02, + } + } + ) + 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("=") + ) + router: Final = self._router() + response: Final = router.aretrieve_batch.return_value + error: Final = litellm.NotFoundError( + message=f"No object found with id '{'batch-456' if error_source == 'batch' else _CLAIM_OUTPUT_FILE_ID}0'.", + model="gpt-4", + llm_provider="openai", + ) + router.aretrieve_batch.side_effect = (error, response) if error_source == "batch" else (response, response) + prisma: Final = self._prisma(row, []) + output: Final = json.dumps( + { + "response": { + "status_code": 200, + "body": { + "model": "batch-recovery-model", + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ).encode() + with respx.mock: + route: Final = respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content") + route.mock( + side_effect=error if error_source == "output" else None, + return_value=None if error_source == "output" else httpx.Response(200, content=output), + ) + await self._instance(prisma, router).check_batch_cost() + assert row.batch_processed is False + assert recorder.costs == () + route.mock(return_value=httpx.Response(200, content=output)) + await self._instance(prisma, router).check_batch_cost() + await self._instance(prisma, router).check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + assert recorder.costs == (pytest.approx(10 * 0.01 + 5 * 0.02),) + @staticmethod def _claim_calls(prisma) -> list: return [