mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(responses): claim a background response before the read that bills it
Every pod and uvicorn worker schedules its own CheckResponsesCost against the shared LiteLLM_ManagedObjectTable. The poller selected eligible rows, performed the billed retrieval, and only then marked them completed in one bulk write, so two pollers could select the same terminal response and both record a charge before either completion update landed. Each row is now claimed with a compare-and-swap on batch_processed before the read, because the read is what prices the job: aget_responses stamped with the poll origin writes the spend log itself, so there is no later point at which to serialize. A row whose read raised, or whose provider status is still non-terminal, releases its claim so a later cycle retries it rather than retiring it unbilled. That is the failure #37050 fixed on the batch side. A pod that dies between winning the claim and billing would otherwise strand the row: it holds a claim nobody will release and its status never reaches terminal, so every later cycle re-selects it and loses. The updated_at arm of the claim takes such a row back after three poll cycles, and since updated_at is @updatedAt a healthy in-flight claim written moments ago is never stolen. The poller now also persists the finished response onto its managed row instead of writing status alone, so the row carries the generation's usage rather than the stale queued copy stored at create time. Reuses the existing batch_processed column, so no migration. It already sits on the shared table defaulted to false and was unused by response rows. Claude-Session: https://claude.ai/code/session_01Hn5E8Jz1LjGLFyiYxBRcBW
This commit is contained in:
ryan-crabbe-berri
jesus
parent
bb81a9f9f1
commit
0235dbd7f2
3 changed files with 635 additions and 112 deletions
|
|
@ -6,7 +6,7 @@ same route are non-inference and free.
|
|||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Dict, Optional, cast
|
||||
from typing import TYPE_CHECKING, Dict, Final, Optional, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -14,6 +14,7 @@ from litellm.constants import (
|
|||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
|
||||
MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
PROXY_BATCH_POLLING_INTERVAL,
|
||||
STALE_OBJECT_CLEANUP_BATCH_SIZE,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
|
@ -21,11 +22,14 @@ from litellm.types.llms.openai import ResponsesAPIResponse
|
|||
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import LiteLLM_ManagedObjectTable
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
|
||||
TERMINAL_RESPONSE_STATUSES = frozenset({"completed", "failed", "cancelled", "incomplete"})
|
||||
|
||||
CLAIM_ABANDONED_AFTER_POLL_CYCLES: Final = 3
|
||||
|
||||
|
||||
class CheckResponsesCost:
|
||||
def __init__(
|
||||
|
|
@ -112,6 +116,92 @@ class CheckResponsesCost:
|
|||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
|
||||
message: Final = str(err).lower()
|
||||
return "batch_processed" in message or "unknown column" in message or "does not exist" in message
|
||||
|
||||
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 CheckResponsesCost against the shared table,
|
||||
so without this compare-and-swap two of them select the same queued response in one window
|
||||
and both bill it. The claim is taken before the read because the read is what prices the
|
||||
job: ``aget_responses`` stamped with the poll origin writes the spend log itself, so there
|
||||
is no later point at which to serialize. Schemas without the column can't be claimed, so
|
||||
they keep the pre-existing behavior rather than silently billing nothing.
|
||||
|
||||
A pod that dies between winning the claim and billing would otherwise strand the row:
|
||||
it holds a claim nobody will release, and its status never reaches terminal, so every
|
||||
later cycle re-selects it and loses. The ``updated_at`` arm takes such a claim back once
|
||||
it has gone unbilled for longer than any live cycle could hold it. ``updated_at`` is
|
||||
``@updatedAt``, so a healthy in-flight claim refreshed moments ago is never stolen.
|
||||
"""
|
||||
abandoned_before: Final = datetime.now(timezone.utc) - timedelta(
|
||||
seconds=CLAIM_ABANDONED_AFTER_POLL_CYCLES * PROXY_BATCH_POLLING_INTERVAL
|
||||
)
|
||||
try:
|
||||
claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"id": job.id,
|
||||
"OR": [
|
||||
{"batch_processed": False},
|
||||
{"updated_at": {"lt": abandoned_before}},
|
||||
],
|
||||
},
|
||||
data={"batch_processed": True},
|
||||
)
|
||||
except Exception as db_err:
|
||||
if self._is_missing_batch_processed_column_error(db_err):
|
||||
verbose_proxy_logger.warning(
|
||||
"CheckResponsesCost: batch_processed column not found, billing without a claim"
|
||||
)
|
||||
return True
|
||||
verbose_proxy_logger.error(f"CheckResponsesCost: 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 when the read did not bill it, so a later poll cycle retries it.
|
||||
|
||||
A response still queued at the provider, or whose read raised, has no spend to record yet.
|
||||
Holding the claim would retire it permanently, which is the failure #37050 hit on batches.
|
||||
"""
|
||||
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"CheckResponsesCost: failed to release the claim on job {job.id}, "
|
||||
f"so its cost will not be retried: {db_err}"
|
||||
)
|
||||
|
||||
async def _persist_terminal_response(
|
||||
self, job: "LiteLLM_ManagedObjectTable", response: ResponsesAPIResponse
|
||||
) -> None:
|
||||
"""Store the finished response on its managed row and retire the row from polling.
|
||||
|
||||
The row is the only copy of a background generation's usage that outlives the poll, so
|
||||
``GET /v1/responses/{id}`` can serve a terminal job from here instead of re-reading it
|
||||
from the provider. Every provider re-read replays the same usage and hands back a
|
||||
freshly encoded id, which is what made this route bill per read and defeated id-based
|
||||
dedup in the first place.
|
||||
|
||||
``status`` stays the literal "completed" for every terminal provider status, matching
|
||||
what this poller has always written, so stale-row expiry keeps skipping these rows.
|
||||
"""
|
||||
try:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": job.id},
|
||||
data={"status": "completed", "file_object": response.model_dump_json()},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckResponsesCost: failed to persist terminal response for job {job.id}: {db_err}"
|
||||
)
|
||||
|
||||
async def check_responses_cost(self):
|
||||
"""
|
||||
Check if background responses are complete and track their cost.
|
||||
|
|
@ -168,33 +258,47 @@ class CheckResponsesCost:
|
|||
litellm_metadata["model"] = model_name
|
||||
litellm_metadata["model_group"] = model_name # Use same value for model_group
|
||||
|
||||
response = await self._get_response(
|
||||
response_id=responses_id_security,
|
||||
litellm_metadata=litellm_metadata,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
if response.status in TERMINAL_RESPONSE_STATUSES:
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
|
||||
if not await self._claim_job_for_costing(job):
|
||||
verbose_proxy_logger.debug(
|
||||
f"Response {unified_object_id} is already claimed for costing, leaving it to the claim holder"
|
||||
)
|
||||
completed_jobs.append(job)
|
||||
continue
|
||||
|
||||
# Mark completed jobs in the database
|
||||
if len(completed_jobs) > 0:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": {"in": [job.id for job in completed_jobs]}},
|
||||
data={"status": "completed"},
|
||||
try:
|
||||
response = await self._get_response(
|
||||
response_id=responses_id_security,
|
||||
litellm_metadata=litellm_metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
await self._release_job_claim(job)
|
||||
verbose_proxy_logger.warning(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
|
||||
)
|
||||
|
||||
if response.status not in TERMINAL_RESPONSE_STATUSES:
|
||||
await self._release_job_claim(job)
|
||||
continue
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} has terminal status {response.status}, marking as complete"
|
||||
)
|
||||
completed_jobs.append((job, response))
|
||||
|
||||
for job, response in completed_jobs:
|
||||
await self._persist_terminal_response(job, response)
|
||||
|
||||
if len(completed_jobs) > 0:
|
||||
verbose_proxy_logger.info(
|
||||
f"Marked {len(completed_jobs)} response jobs as completed"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,8 @@ Unit tests for CheckResponsesCost class
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -12,6 +13,40 @@ from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
|||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
|
||||
|
||||
def _update_many_calls_writing(mock_prisma_client, matches_data):
|
||||
return [
|
||||
call
|
||||
for call in mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
if matches_data(call.kwargs["data"])
|
||||
]
|
||||
|
||||
|
||||
def _completion_calls(mock_prisma_client):
|
||||
return _update_many_calls_writing(
|
||||
mock_prisma_client, lambda data: data.get("status") == "completed"
|
||||
)
|
||||
|
||||
|
||||
def _completed_job_ids(mock_prisma_client):
|
||||
return [call.kwargs["where"]["id"] for call in _completion_calls(mock_prisma_client)]
|
||||
|
||||
|
||||
def _persisted_response(completion_call):
|
||||
return json.loads(completion_call.kwargs["data"]["file_object"])
|
||||
|
||||
|
||||
def _claim_calls(mock_prisma_client):
|
||||
return _update_many_calls_writing(
|
||||
mock_prisma_client, lambda data: data == {"batch_processed": True}
|
||||
)
|
||||
|
||||
|
||||
def _release_calls(mock_prisma_client):
|
||||
return _update_many_calls_writing(
|
||||
mock_prisma_client, lambda data: data == {"batch_processed": False}
|
||||
)
|
||||
|
||||
|
||||
class TestCheckResponsesCost:
|
||||
"""Test suite for CheckResponsesCost class"""
|
||||
|
||||
|
|
@ -135,7 +170,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check with mocked litellm.aget_responses
|
||||
|
|
@ -144,14 +179,8 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# update_many should only contain the job completion call
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
completion_call = calls[0]
|
||||
assert completion_call[1]["data"]["status"] == "completed"
|
||||
assert completion_call[1]["where"]["id"]["in"] == ["job-123"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-123"]
|
||||
assert _release_calls(mock_prisma_client) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_failed_response(
|
||||
|
|
@ -180,7 +209,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check
|
||||
|
|
@ -189,12 +218,8 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# update_many should only contain the job completion call
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"]["status"] == "completed"
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-456"]
|
||||
assert _release_calls(mock_prisma_client) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_cancelled_response(
|
||||
|
|
@ -223,7 +248,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check
|
||||
|
|
@ -232,12 +257,8 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# update_many should only contain the job completion call
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"]["status"] == "completed"
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-789"]
|
||||
assert _release_calls(mock_prisma_client) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_in_progress_response(
|
||||
|
|
@ -266,7 +287,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check
|
||||
|
|
@ -276,10 +297,7 @@ class TestCheckResponsesCost:
|
|||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# No job completion update_many — response is still in progress
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 0
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
# Stale cleanup still ran via _expire_stale_rows
|
||||
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
||||
|
||||
|
|
@ -310,7 +328,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check
|
||||
|
|
@ -320,10 +338,7 @@ class TestCheckResponsesCost:
|
|||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# No job completion update_many — response is still queued
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 0
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
# Stale cleanup still ran via _expire_stale_rows
|
||||
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
||||
|
||||
|
|
@ -344,7 +359,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check with mocked exception
|
||||
|
|
@ -357,10 +372,7 @@ class TestCheckResponsesCost:
|
|||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# No job completion update_many — exception skipped the job
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 0
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
# Stale cleanup still ran via _expire_stale_rows
|
||||
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
||||
|
||||
|
|
@ -429,7 +441,7 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Run the check
|
||||
|
|
@ -438,16 +450,7 @@ class TestCheckResponsesCost:
|
|||
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
# update_many should only contain the job completion call
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
completion_call = calls[0]
|
||||
assert len(completion_call[1]["where"]["id"]["in"]) == 2
|
||||
assert "job-1" in completion_call[1]["where"]["id"]["in"]
|
||||
assert "job-3" in completion_call[1]["where"]["id"]["in"]
|
||||
assert "job-2" not in completion_call[1]["where"]["id"]["in"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-1", "job-3"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encoded_response_id_is_fetched_through_router(
|
||||
|
|
@ -480,7 +483,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
|
|
@ -511,12 +514,7 @@ class TestCheckResponsesCost:
|
|||
== encoded_response_id
|
||||
)
|
||||
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"]["status"] == "completed"
|
||||
assert calls[0][1]["where"]["id"]["in"] == ["job-router"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-router"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encrypted_response_id_is_fetched_through_router(
|
||||
|
|
@ -556,7 +554,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
|
|
@ -584,11 +582,7 @@ class TestCheckResponsesCost:
|
|||
mock_llm_router.aget_responses.call_args[1]["response_id"]
|
||||
== encoded_response_id
|
||||
)
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["where"]["id"]["in"] == ["job-encrypted"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-encrypted"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_id_without_model_id_uses_sdk(
|
||||
|
|
@ -605,7 +599,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
side_effect=AssertionError("router cannot route an id without a model_id")
|
||||
|
|
@ -654,7 +648,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
mock_llm_router.get_deployment = MagicMock(return_value=None)
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
|
|
@ -679,12 +673,7 @@ class TestCheckResponsesCost:
|
|||
mock_sdk_aget.assert_called_once()
|
||||
assert mock_sdk_aget.call_args[1]["response_id"] == encoded_response_id
|
||||
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"]["status"] == "completed"
|
||||
assert calls[0][1]["where"]["id"]["in"] == ["job-missing-deployment"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-missing-deployment"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_with_incomplete_response(
|
||||
|
|
@ -701,7 +690,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
|
|
@ -717,12 +706,7 @@ class TestCheckResponsesCost:
|
|||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
assert len(calls) == 1
|
||||
assert calls[0][1]["data"]["status"] == "completed"
|
||||
assert calls[0][1]["where"]["id"]["in"] == ["job-incomplete"]
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-incomplete"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_no_model_in_file_object(
|
||||
|
|
@ -741,7 +725,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -783,7 +767,7 @@ class TestCheckResponsesCost:
|
|||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = MagicMock()
|
||||
|
|
@ -800,3 +784,430 @@ class TestCheckResponsesCost:
|
|||
assert metadata["user_api_key_hash"] == "sk-billed"
|
||||
assert is_unbilled_non_inference_call("aget_responses", metadata) is False
|
||||
assert is_unbilled_non_inference_call("aget_responses", None) is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_job_claimed_by_another_pod_is_never_read_or_completed(
|
||||
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""Every pod and uvicorn worker polls the same table, and the read is what writes the
|
||||
spend log, so losing the claim has to skip the read entirely or the job is billed twice."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_claimed_elsewhere"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-claimed-elsewhere"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_claimed_elsewhere"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=0
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
mock_sdk_aget.assert_not_awaited()
|
||||
mock_llm_router.aget_responses.assert_not_awaited()
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
assert _release_calls(mock_prisma_client) == []
|
||||
|
||||
claim_calls = _claim_calls(mock_prisma_client)
|
||||
assert len(claim_calls) == 1
|
||||
claim_where = claim_calls[0].kwargs["where"]
|
||||
assert claim_where["id"] == "job-claimed-elsewhere"
|
||||
assert {"batch_processed": False} in claim_where["OR"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_is_taken_back_from_a_pod_that_died_holding_it(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""A pod that dies between claiming and billing releases nothing, and the row's status
|
||||
never reaches terminal, so without a lease every later cycle re-selects it and loses.
|
||||
The window has to be longer than a live cycle can hold a claim and short enough that the
|
||||
row is retried well before stale expiry gives up on it unbilled."""
|
||||
from litellm.constants import PROXY_BATCH_POLLING_INTERVAL
|
||||
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
||||
CLAIM_ABANDONED_AFTER_POLL_CYCLES,
|
||||
)
|
||||
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_abandoned"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-abandoned"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_abandoned"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_abandoned",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
claim_where = _claim_calls(mock_prisma_client)[0].kwargs["where"]
|
||||
abandoned_arm = next(arm for arm in claim_where["OR"] if "updated_at" in arm)
|
||||
lease = timedelta(
|
||||
seconds=CLAIM_ABANDONED_AFTER_POLL_CYCLES * PROXY_BATCH_POLLING_INTERVAL
|
||||
)
|
||||
untouched_for = datetime.now(timezone.utc) - abandoned_arm["updated_at"]["lt"]
|
||||
assert lease <= untouched_for < lease + timedelta(seconds=30)
|
||||
assert lease > timedelta(seconds=PROXY_BATCH_POLLING_INTERVAL)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_is_taken_before_the_billing_read_and_kept_on_a_terminal_status(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""The read prices the job, so the claim has to be taken before it, and keeping the claim
|
||||
afterwards is what stops a second pod reading and billing the same row again."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_ordering"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-ordering"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_ordering"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
|
||||
writes_and_reads = []
|
||||
|
||||
async def record_update_many(**kwargs):
|
||||
writes_and_reads.append(kwargs["data"])
|
||||
return 1
|
||||
|
||||
async def record_read(**kwargs):
|
||||
writes_and_reads.append("provider_read")
|
||||
return ResponsesAPIResponse(
|
||||
id="resp_ordering",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(
|
||||
input_tokens=100, output_tokens=50, total_tokens=150
|
||||
),
|
||||
)
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=record_update_many
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.aget_responses", new_callable=AsyncMock, side_effect=record_read
|
||||
):
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
assert len(writes_and_reads) == 3
|
||||
assert writes_and_reads[0] == {"batch_processed": True}
|
||||
assert writes_and_reads[1] == "provider_read"
|
||||
assert writes_and_reads[2]["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("provider_status", ["queued", "in_progress"])
|
||||
async def test_non_terminal_status_releases_the_claim(
|
||||
self, check_responses_cost_instance, mock_prisma_client, provider_status
|
||||
):
|
||||
"""A response the provider has not finished yet has no spend to record, so its row must go
|
||||
back to batch_processed=False; holding the claim retires it before it is ever billed."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_still_running"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-still-running"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_still_running"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_still_running",
|
||||
object="response",
|
||||
status=provider_status,
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=None,
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
release_calls = _release_calls(mock_prisma_client)
|
||||
assert len(release_calls) == 1
|
||||
assert release_calls[0].kwargs["where"] == {
|
||||
"id": "job-still-running",
|
||||
"batch_processed": True,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_provider_read_releases_the_claim(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""A read that raised billed nothing, so the claim has to be handed back or the row is
|
||||
retired unbilled and no later poll cycle ever retries it."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_read_error"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-read-error"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_read_error"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.aget_responses",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=Exception("Provider error"),
|
||||
):
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
assert _completion_calls(mock_prisma_client) == []
|
||||
release_calls = _release_calls(mock_prisma_client)
|
||||
assert len(release_calls) == 1
|
||||
assert release_calls[0].kwargs["where"] == {
|
||||
"id": "job-read-error",
|
||||
"batch_processed": True,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_claimed_elsewhere_does_not_block_the_next_job(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""Losing one row to another pod must skip only that row: the rest of the poll page still
|
||||
has to be read and billed in the same cycle."""
|
||||
mock_job1 = MagicMock()
|
||||
mock_job1.unified_object_id = "resp_test_first"
|
||||
mock_job1.created_by = "user1"
|
||||
mock_job1.id = "job-first"
|
||||
mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_first"}
|
||||
|
||||
mock_job2 = MagicMock()
|
||||
mock_job2.unified_object_id = "resp_test_second"
|
||||
mock_job2.created_by = "user2"
|
||||
mock_job2.id = "job-second"
|
||||
mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_second"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job1, mock_job2]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=[0, 1, 1]
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_second",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
mock_aget.assert_awaited_once()
|
||||
assert mock_aget.await_args.kwargs["response_id"] == "resp_test_second"
|
||||
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-second"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"db_error_message",
|
||||
[
|
||||
"column LiteLLM_ManagedObjectTable.batch_processed does not exist",
|
||||
"Unknown column in where clause",
|
||||
"The column P2022 does not exist in the current database",
|
||||
],
|
||||
)
|
||||
async def test_claim_fails_open_on_a_schema_without_the_claim_column(
|
||||
self, check_responses_cost_instance, mock_prisma_client, db_error_message
|
||||
):
|
||||
"""A deployment that never ran the batch_processed migration cannot claim anything, so it
|
||||
keeps the pre-claim behavior of billing rather than silently billing nothing."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.id = "job-old-schema"
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=Exception(db_error_message)
|
||||
)
|
||||
|
||||
assert (
|
||||
await check_responses_cost_instance._claim_job_for_costing(mock_job) is True
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_claim_is_lost_when_the_database_fails_for_any_other_reason(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""A dropped connection is no proof the row is free, so the read that would bill it is
|
||||
not allowed to run."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.id = "job-db-down"
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=Exception("connection to server was lost")
|
||||
)
|
||||
|
||||
assert (
|
||||
await check_responses_cost_instance._claim_job_for_costing(mock_job) is False
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_old_schema_without_the_claim_column_still_bills_and_completes(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""End to end on a pre-migration schema: the claim write fails, the response is still read
|
||||
(which is what bills it) and the row is still marked completed."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_old_schema"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-old-schema"
|
||||
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_old_schema"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
|
||||
async def reject_batch_processed_writes(**kwargs):
|
||||
if "batch_processed" in kwargs["data"]:
|
||||
raise Exception(
|
||||
'column "batch_processed" of relation '
|
||||
'"LiteLLM_ManagedObjectTable" does not exist'
|
||||
)
|
||||
return 1
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=reject_batch_processed_writes
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_old_schema",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
mock_aget.assert_awaited_once()
|
||||
assert _completed_job_ids(mock_prisma_client) == ["job-old-schema"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_response_replaces_the_queued_file_object_on_the_row(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""The row has to become the durable copy of the finished generation. Leaving the queued
|
||||
placeholder there forces the retrieve endpoint back to the provider, and every re-read
|
||||
replays the same usage under a fresh id, which is what billed the job again per read."""
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = "resp_test_persisted"
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-persisted"
|
||||
mock_job.file_object = {
|
||||
"model": "gpt-5",
|
||||
"id": "resp_test_persisted",
|
||||
"status": "queued",
|
||||
"usage": None,
|
||||
}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_finished_upstream",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(
|
||||
input_tokens=100, output_tokens=50, total_tokens=150
|
||||
),
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
completion_calls = _completion_calls(mock_prisma_client)
|
||||
assert len(completion_calls) == 1
|
||||
assert completion_calls[0].kwargs["where"] == {"id": "job-persisted"}
|
||||
|
||||
persisted = _persisted_response(completion_calls[0])
|
||||
assert persisted["id"] == "resp_finished_upstream"
|
||||
assert persisted["status"] == "completed"
|
||||
assert persisted["usage"]["total_tokens"] == 150
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_persist_does_not_abort_the_rest_of_the_poll_cycle(
|
||||
self, check_responses_cost_instance, mock_prisma_client
|
||||
):
|
||||
"""One row's write failing must not take the whole cycle down with it: the jobs behind it
|
||||
are already read and billed, so losing their write loses their usage for good."""
|
||||
mock_job1 = MagicMock()
|
||||
mock_job1.unified_object_id = "resp_test_persist_fails"
|
||||
mock_job1.created_by = "user1"
|
||||
mock_job1.id = "job-persist-fails"
|
||||
mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_persist_fails"}
|
||||
|
||||
mock_job2 = MagicMock()
|
||||
mock_job2.unified_object_id = "resp_test_persist_works"
|
||||
mock_job2.created_by = "user2"
|
||||
mock_job2.id = "job-persist-works"
|
||||
mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_persist_works"}
|
||||
|
||||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job1, mock_job2]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
side_effect=[1, 1, Exception("deadlock detected"), 1]
|
||||
)
|
||||
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="resp_persisted",
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
|
||||
)
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
||||
mock_aget.return_value = mock_response
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
||||
assert mock_aget.await_count == 2
|
||||
assert _completed_job_ids(mock_prisma_client) == [
|
||||
"job-persist-fails",
|
||||
"job-persist-works",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -369,7 +369,9 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
|
||||
# Mock update_many
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Create a completed response
|
||||
completed_response = ResponsesAPIResponse(
|
||||
|
|
@ -398,17 +400,17 @@ class TestCheckResponsesCost:
|
|||
await checker.check_responses_cost()
|
||||
|
||||
# Verify update_many was called to mark job as completed
|
||||
# (stale cleanup also calls update_many, so check the specific completion call)
|
||||
# (the costing claim also calls update_many, so check the specific completion call)
|
||||
update_many_calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
completion_calls = [
|
||||
c
|
||||
for c in update_many_calls
|
||||
if c.kwargs.get("where", {}).get("id") is not None
|
||||
if c.kwargs["data"].get("status") == "completed"
|
||||
]
|
||||
assert len(completion_calls) == 1
|
||||
assert completion_calls[0].kwargs["where"]["id"]["in"] == ["job-123"]
|
||||
assert completion_calls[0].kwargs["where"]["id"] == "job-123"
|
||||
assert completion_calls[0].kwargs["data"]["status"] == "completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -429,7 +431,9 @@ class TestCheckResponsesCost:
|
|||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Create a failed response
|
||||
failed_response = ResponsesAPIResponse(
|
||||
|
|
@ -453,14 +457,14 @@ class TestCheckResponsesCost:
|
|||
await checker.check_responses_cost()
|
||||
|
||||
# Verify job was marked as completed even though it failed
|
||||
# (stale cleanup also calls update_many, so check the specific completion call)
|
||||
# (the costing claim also calls update_many, so check the specific completion call)
|
||||
update_many_calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
completion_calls = [
|
||||
c
|
||||
for c in update_many_calls
|
||||
if c.kwargs.get("where", {}).get("id") is not None
|
||||
if c.kwargs["data"].get("status") == "completed"
|
||||
]
|
||||
assert len(completion_calls) == 1
|
||||
|
||||
|
|
@ -482,7 +486,9 @@ class TestCheckResponsesCost:
|
|||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
# Create an in-progress response
|
||||
in_progress_response = ResponsesAPIResponse(
|
||||
|
|
@ -506,14 +512,14 @@ class TestCheckResponsesCost:
|
|||
await checker.check_responses_cost()
|
||||
|
||||
# Verify no completion update_many was called (job still in progress)
|
||||
# (stale cleanup may still call update_many, so filter for completion calls)
|
||||
# (the claim and its release also call update_many, so filter for completion calls)
|
||||
update_many_calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
completion_calls = [
|
||||
c
|
||||
for c in update_many_calls
|
||||
if c.kwargs.get("where", {}).get("id") is not None
|
||||
if c.kwargs["data"].get("status") == "completed"
|
||||
]
|
||||
assert len(completion_calls) == 0
|
||||
|
||||
|
|
@ -535,7 +541,9 @@ class TestCheckResponsesCost:
|
|||
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
return_value=[mock_job]
|
||||
)
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
||||
return_value=1
|
||||
)
|
||||
|
||||
checker = CheckResponsesCost(
|
||||
proxy_logging_obj=mock_proxy_logging_obj,
|
||||
|
|
@ -553,13 +561,13 @@ class TestCheckResponsesCost:
|
|||
await checker.check_responses_cost()
|
||||
|
||||
# Verify no completion update_many was called (error occurred)
|
||||
# (stale cleanup may still call update_many, so filter for completion calls)
|
||||
# (the claim and its release also call update_many, so filter for completion calls)
|
||||
update_many_calls = (
|
||||
mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
)
|
||||
completion_calls = [
|
||||
c
|
||||
for c in update_many_calls
|
||||
if c.kwargs.get("where", {}).get("id") is not None
|
||||
if c.kwargs["data"].get("status") == "completed"
|
||||
]
|
||||
assert len(completion_calls) == 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue