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 a8e46349917..4bb00408fc3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -255,6 +255,52 @@ class CheckBatchCost: "so it will no longer be polled" ) + async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool: + """ + Atomically flip batch_processed from false to true, returning whether this pod won + the row. Every pod and uvicorn worker schedules its own poller against the shared + table, so without this compare-and-swap two of them can select the same completed + batch in one window and both emit an aretrieve_batch spend log for it. Schemas + without the column can't be claimed, so they keep the pre-existing behavior. + + Called immediately before the spend log is written rather than before the results + fetch, because batch_processed is also what holds off deletion of the files that + fetch reads and what keeps an unbilled row selectable by the next poll cycle. + """ + if not self._has_batch_processed_column: + return True + try: + claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={"id": job.id, "batch_processed": False}, + data={"batch_processed": True}, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to claim job {job.id} for cost tracking: {db_err}" + ) + return False + return claimed > 0 + + async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None: + """Give a claimed row back once billing it failed, so a later poll cycle retries it. + + Safe to match on batch_processed=True: while this poller is active the retrieve + path leaves the column alone (batch_cost_poller_is_active), so a true value here + is always this pod's own claim. + """ + if not self._has_batch_processed_column: + return + try: + await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={"id": job.id, "batch_processed": True}, + data={"batch_processed": False}, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to release the claim on job {job.id}, " + f"so its cost will not be retried: {db_err}" + ) + @staticmethod def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool: """A unified id that decodes but carries no model_id can never be routed.""" @@ -572,9 +618,10 @@ class CheckBatchCost: """ Fetch a completed batch's results, compute cost/usage, and emit the aretrieve_batch spend log. Returns (model_name, llm_provider) on - success, None when the job can't be routed to a deployment. Raises on - results-fetch or cost-computation failures so the caller can leave the - job unprocessed and retry it on a later poll. + success, None when the job can't be routed to a deployment or when + another pod claimed it. Raises on results-fetch or cost-computation + failures so the caller can leave the job unprocessed and retry it on a + later poll. """ from litellm.batches.batch_utils import ( _get_file_content_as_dictionary, @@ -743,12 +790,23 @@ class CheckBatchCost: optional_params={}, ) - await logging_obj.async_success_handler( - result=response, - batch_cost=batch_cost, - batch_usage=batch_usage, - batch_models=batch_models, - ) + if not await self._claim_job_for_costing(job): + verbose_proxy_logger.info( + f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod " + "in this window, so its cost is already being tracked there" + ) + return None + + try: + await logging_obj.async_success_handler( + result=response, + batch_cost=batch_cost, + batch_usage=batch_usage, + batch_models=batch_models, + ) + except Exception: + await self._release_job_claim(job) + raise # Record batch duration (completed_at - created_at) if prom_logger and response.completed_at and response.created_at: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 1dbbbfc43a0..0065dbebc59 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -6,11 +6,17 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, ARN unified_object_id) batches with no managed unified id. """ +import asyncio +import json +from contextlib import contextmanager from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" +_CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA==" +_CLAIM_OUTPUT_FILE_ID = "file-output-123" def _unmanaged_vertex_file_object( @@ -95,7 +101,7 @@ class TestCheckBatchCost: ): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) # Return empty so the main poll loop exits immediately mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( @@ -161,7 +167,7 @@ class TestCheckBatchCost: from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[] @@ -192,7 +198,7 @@ class TestCheckBatchCost: from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( @@ -221,7 +227,7 @@ class TestCheckBatchCost: from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False @@ -254,7 +260,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -563,7 +569,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -679,7 +685,7 @@ class TestCheckBatchCost: import litellm from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + 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) @@ -801,7 +807,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -869,7 +875,7 @@ class TestCheckBatchCost: import base64 mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -944,7 +950,7 @@ class TestCheckBatchCost: ).decode() mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -1044,7 +1050,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -1111,7 +1117,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() @@ -1168,7 +1174,7 @@ class TestCheckBatchCost: from unittest.mock import patch mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -1284,7 +1290,7 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -1355,7 +1361,7 @@ class TestCheckBatchCost: 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 + return_value=1 ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( @@ -1672,7 +1678,7 @@ class TestUnmanagedVertexRouting: prisma = instance.prisma_client prisma.db = MagicMock() prisma.db.litellm_managedobjecttable = MagicMock() - prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + 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()] @@ -1902,7 +1908,7 @@ class TestUnmanagedBedrockRouting: prisma = instance.prisma_client prisma.db = MagicMock() prisma.db.litellm_managedobjecttable = MagicMock() - prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + 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()] @@ -2577,3 +2583,353 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + +class _FakeManagedObjectRow: + """One managed batch row the provider has finished but nothing has costed yet.""" + + def __init__(self): + self.id = "job-claim-1" + self.unified_object_id = _CLAIM_UNIFIED_BATCH_ID + self.model_object_id = "batch-456" + self.file_purpose = "batch" + self.status = "in_progress" + self.batch_processed = False + self.created_by = "user-1" + self.team_id = None + self.api_key = None + 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} + ) + + +class _FakeManagedObjectTable: + """A LiteLLM_ManagedObjectTable double backed by one real, mutable row. + + It honours the batch_processed and status filters, so the poller's compare-and-swap + and the managed-files deletion guard both read the same state a shared Postgres row + would give them. Staleness sweeps (the only queries scoped by created_at) never match. + """ + + def __init__(self, row: _FakeManagedObjectRow, journal: list): + self.row = row + self.journal = journal + self.update_many = AsyncMock(side_effect=self._update_many) + self.update = AsyncMock(side_effect=self._update) + self.find_many = AsyncMock(side_effect=self._find_many) + self.find_first = AsyncMock(return_value=None) + + def _matches(self, where: dict) -> bool: + for key, value in where.items(): + if key == "created_at": + return False + if key == "status": + if self.row.status in value.get("not_in", []): + return False + if "in" in value and self.row.status not in value["in"]: + return False + elif getattr(self.row, key) != value: + return False + return True + + async def _update_many(self, *, where: dict, data: dict) -> int: + if not self._matches(where): + return 0 + if "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) + return 1 + + async def _update(self, *, where: dict, data: dict) -> None: + self.journal.append("finalize") + for key, value in data.items(): + setattr(self.row, key, value) + + async def _find_many(self, *, where: dict, take=None, order=None) -> list: + return [self.row] if self._matches(where) else [] + + +class TestMultiPodBatchCostClaim: + """LIT-4827 regression: every pod and uvicorn worker schedules its own poller against + the shared LiteLLM_ManagedObjectTable, so a completed batch must be claimed atomically + before its cost is logged. Without the claim two pods select the same row in one window + and both write an aretrieve_batch spend log for it, double counting the spend. + + The claim sits immediately before the spend-log write rather than before the results + fetch, because batch_processed is also what keeps an unbilled row selectable by later + poll cycles and what blocks deletion of the files the fetch reads.""" + + @staticmethod + def _instance(prisma, llm_router): + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.return_value = None + return CheckBatchCost( + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma, + llm_router=llm_router, + ) + + @staticmethod + def _prisma(row: _FakeManagedObjectRow, journal: list): + prisma = MagicMock() + prisma.db.litellm_managedobjecttable = _FakeManagedObjectTable(row, journal) + prisma.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[]) + prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + return prisma + + @staticmethod + def _router(): + response = MagicMock() + response.status = "completed" + response.output_file_id = _CLAIM_OUTPUT_FILE_ID + response.error_file_id = None + response.created_at = 1 + response.completed_at = 2 + response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}' + + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "openai" + deployment.litellm_params.model = "gpt-4" + deployment.model_info.model_dump.return_value = {} + + 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 = MagicMock(return_value=deployment) + return router + + @staticmethod + @contextmanager + def _billing_patches(journal: list, during_fetch=None, bill_error=None): + """Patch the cost path a batch runs through, journalling the results fetch and the + spend-log write. during_fetch runs while the output file is being read, which is + the window an interrupted worker or a concurrent file deletion lands in.""" + file_content = MagicMock() + file_content.content = b'{"id":"req-1"}' + + async def _afile_content(**kwargs): + journal.append("fetch") + if during_fetch is not None: + await during_fetch() + return file_content + + async def _bill(**kwargs): + journal.append("bill") + if bill_error is not None: + raise bill_error + + def _is_b64(file_id): + if file_id == _CLAIM_UNIFIED_BATCH_ID: + return "llm_model_id,model-123;llm_batch_id,batch-456;" + return False + + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock(side_effect=_bill) + + with ( + patch(_IS_B64, side_effect=_is_b64), + 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=AsyncMock(side_effect=_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=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]), + ), + 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", return_value=logging_obj), + ): + yield logging_obj + + @staticmethod + def _claim_calls(prisma) -> list: + return [ + call.kwargs + for call in prisma.db.litellm_managedobjecttable.update_many.call_args_list + if "id" in call.kwargs["where"] + ] + + @staticmethod + async def _run_deletion_guard(prisma, file_id: str) -> None: + """Run the real managed-files deletion guard against the row the poller is costing.""" + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=None) + cache.async_set_cache = AsyncMock() + guard = _PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma) + + scheduler = MagicMock() + scheduler.get_job.return_value = MagicMock() + with patch("litellm.proxy.proxy_server.scheduler", scheduler): + await guard._check_file_deletion_allowed(file_id) + + @pytest.mark.asyncio + async def test_winning_pod_claims_the_row_between_fetching_and_billing(self): + """The claim flips batch_processed false -> true after the results are in hand and + before the spend log is written, so a concurrent pod's claim finds no matching row.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + + with self._billing_patches(journal) as logging_obj: + await self._instance(prisma, self._router()).check_batch_cost() + + assert journal == ["fetch", "claim", "bill", "finalize"] + assert self._claim_calls(prisma) == [ + { + "where": {"id": "job-claim-1", "batch_processed": False}, + "data": {"batch_processed": True}, + } + ] + logging_obj.async_success_handler.assert_awaited_once() + assert row.batch_processed is True + + @pytest.mark.asyncio + async def test_a_pod_that_loses_the_claim_after_fetching_does_not_bill(self): + """Both pods select the row and fetch its results in the same window. The one whose + compare-and-swap finds the row already taken must not write a second spend log.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + + async def _other_pod_wins_the_row(): + row.batch_processed = True + + with self._billing_patches(journal, during_fetch=_other_pod_wins_the_row) as logging_obj: + await self._instance(prisma, self._router()).check_batch_cost() + + assert journal == ["fetch"] + logging_obj.async_success_handler.assert_not_awaited() + assert self._claim_calls(prisma) == [ + { + "where": {"id": "job-claim-1", "batch_processed": False}, + "data": {"batch_processed": True}, + } + ] + + @pytest.mark.asyncio + async def test_a_failed_spend_log_write_releases_the_claim(self): + """A transient failure while billing a claimed batch must hand the row back, or its + spend is silently lost instead of being retried on the next cycle.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + + with self._billing_patches(journal, bill_error=Exception("spend log write failed")): + await self._instance(prisma, self._router()).check_batch_cost() + + assert journal == ["fetch", "claim", "bill", "release"] + assert row.batch_processed is False + assert self._claim_calls(prisma)[-1] == { + "where": {"id": "job-claim-1", "batch_processed": True}, + "data": {"batch_processed": False}, + } + + @pytest.mark.asyncio + async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self): + """A pod killed while reading a batch's results must leave the row for a later + cycle. Claiming before the fetch marked the batch processed for good, so the pod + that died took that batch's spend with it and no other pod ever selected it.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + reached_fetch = asyncio.Event() + + async def _never_returns(): + reached_fetch.set() + 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() + ) + 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() + with pytest.raises(asyncio.CancelledError): + await interrupted + + assert journal == ["fetch"] + logging_obj.async_success_handler.assert_not_awaited() + + survivor_journal = [] + survivor_prisma = self._prisma(row, survivor_journal) + with self._billing_patches(survivor_journal) as survivor_logging: + await self._instance(survivor_prisma, self._router()).check_batch_cost() + + assert survivor_journal == ["fetch", "claim", "bill", "finalize"] + survivor_logging.async_success_handler.assert_awaited_once() + assert row.batch_processed is True + + @pytest.mark.asyncio + async def test_costing_in_flight_keeps_the_referenced_file_undeletable(self): + """The deletion guard only holds files whose batch still has batch_processed false, + so claiming the row before the fetch let a concurrent delete remove the very output + file the in-flight costing was about to read.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + reached_fetch = asyncio.Event() + finish_fetch = asyncio.Event() + + async def _wait_for_the_delete_attempt(): + reached_fetch.set() + 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() + ) + await asyncio.wait_for(reached_fetch.wait(), timeout=5) + + with pytest.raises(HTTPException) as blocked: + await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + assert blocked.value.status_code == 400 + assert _CLAIM_OUTPUT_FILE_ID in blocked.value.detail + + finish_fetch.set() + await asyncio.wait_for(costing, timeout=5) + + assert journal == ["fetch", "claim", "bill", "finalize"] + assert row.batch_processed is True + await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + + @pytest.mark.asyncio + async def test_schema_without_batch_processed_still_bills(self): + """Older schemas have no column to claim, so they keep the pre-fix behavior instead + of losing every batch's cost.""" + row = _FakeManagedObjectRow() + journal = [] + prisma = self._prisma(row, journal) + instance = self._instance(prisma, self._router()) + instance._has_batch_processed_column = False + + with self._billing_patches(journal) as logging_obj: + await instance.check_batch_cost() + + assert self._claim_calls(prisma) == [] + assert journal == ["fetch", "bill", "finalize"] + logging_obj.async_success_handler.assert_awaited_once() diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index d9a0b275392..c75c8099ea1 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -192,6 +192,7 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti return_value=[mock_job] ) mock_prisma.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) # Mock proxy_logging_obj — should NOT be called for file content mock_proxy_logging = MagicMock()