mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fbc2e358bc
* fix(responses): bill a background response once, from the poll job
Reading a completed background response with GET /v1/responses/{id} wrote a
full-cost spend row on every read, on top of the row the background cost poll
writes. Reads of a Responses object (GET, DELETE, cancel) now never bill, and
the poll job is the single biller.
The poll job now bills with the creating request's attribution (key hash, team,
org, user, aliases, tags) read off the managed object row, through a helper it
shares with the batch cost poll. The background create persists the creating
key and request tags so that attribution survives until the poll runs.
* test(proxy): a background response read reports zero cost
* fix(responses): bill a background response once, on the first read or poll that sees it finished
A conditional update moves the stored row out of queued or in_progress so only one worker or pod bills it, through one re-read stamped with the creator's attribution. A failed billing read puts the row back for the next poll
* fix(responses): keep the caller's response id out of the billing failure warning
* fix(responses): keep the creating client's User-Agent tags on the one background charge
* fix(responses): keep managed background reads free and bill each response once
The managed files hook marks reads and cancels of proxy-issued response ids, so only those stay free. SDK and database-less reads of a background response still bill, extra_body background creates are tracked, and the poll's status read no longer cools down a deployment on a 404.
* fix(responses): bill finished background responses from the fetched body and store rows Postgres accepts
* fix(responses): give the background billing event its own spend log row and bill responses the create left untracked
* fix(responses): narrow the managed files hook with a protocol instead of an enterprise import
* fix(responses): leave a read the caller retargets with client-side credentials to the cost poll
* fix(responses): bill background responses on the creator's end user and report the charge on the settling read
The read that settles the bill returns the charge in x-litellm-response-cost, the billed event
carries the creator's end user and the key, team, and user budget fields, and its spend row is
named after the id the client got. Override reads, cancels, and deletes settle from the
configured deployment, a streamed create marks its row billed, stored rows keep no model
output, and model insights counts the billed row's spend without a request.
* test(responses): cover background response billing across proxies, reads and chaos
* fix(responses): settle background billing only for callers allowed to address the response
A delete now settles the bill only after the response id access check passes, a streamed
create is marked billed when it sets background inside extra_body or its completed event
arrives after the last chunk, the billing event keeps the deployment's litellm_params
base_model, a key with no user and no team bills no user instead of its owner scope, and the
enterprise cost check import is optional
* fix(responses): record the bound on the pending status IN filters
* fix(responses): send the billed cost split on the settling read and keep the batch cost event as it was
The read that settles a background response now sends the whole cost header
family from the billed event, so the breakdown headers add up to the charge.
Batch cost events go back to main's attribution (creator id as the user, no
budget fields); only background responses get the end user and budgets.
Integration tests expect the charge on the settling read (or none when the
poll billed first), one charge for a streamed create, and SDK flows that poll
with the id each read returns.
* fix(responses): bill each background response once across upgrades, streams, and failed settles
Rows stored before this billing model are found by the issued id and rekeyed by the
provider id, and the cost poll retires one a read already billed. A streamed background
create is free on its own and bills the finished response through the same managed row
before the client sees response.completed, so a racing read cannot bill it again. A
settle that fails retries in the background, a retargeted read with no row tracks and
bills from the configured deployment, and a delete keeps its billed cost header.
Reverts the unrelated json.dumps narrowing in the completion transformation.
* fix(responses): bill cache-polled background streams and match poll rows by the issued id
* test(responses): cover deletes, failed settles, budgets, base models, and cache-polled background billing
Adds integration cells for the background response billing paths the audit
inventory lists: a DELETE that settles the bill and refuses a stranger, a
settle the database refuses that retries once it recovers, the Prometheus
remaining budget, base_model pricing from model_info and litellm_params, the
daily model usage rollup, the end user, the settling read's cost split, stored
rows without model output, a forged managed marker on the create, streamed
creates with background in extra_body, and the cache-polled create
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
567 lines
21 KiB
Python
567 lines
21 KiB
Python
"""
|
|
Integration tests for responses API background cost tracking
|
|
"""
|
|
|
|
import asyncio
|
|
import os
|
|
from datetime import datetime
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
import pytest
|
|
|
|
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
|
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
|
|
|
|
|
class TestResponsesBackgroundCostTracking:
|
|
"""Integration tests for responses API background cost tracking"""
|
|
|
|
@pytest.fixture
|
|
def mock_managed_files_obj(self):
|
|
"""Create a mock managed files object"""
|
|
managed_files = MagicMock()
|
|
managed_files.store_unified_object_id = AsyncMock()
|
|
return managed_files
|
|
|
|
@pytest.fixture
|
|
def mock_proxy_logging_obj(self, mock_managed_files_obj):
|
|
"""Create a mock proxy logging object"""
|
|
logging_obj = MagicMock()
|
|
logging_obj.get_proxy_hook = MagicMock(return_value=mock_managed_files_obj)
|
|
return logging_obj
|
|
|
|
@pytest.fixture
|
|
def mock_llm_router(self):
|
|
"""Create a mock LLM router"""
|
|
router = MagicMock()
|
|
return router
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_response_in_managed_objects_table(
|
|
self, mock_managed_files_obj, mock_proxy_logging_obj, mock_llm_router
|
|
):
|
|
"""Test that background responses are stored in managed objects table"""
|
|
# Create a mock response with queued status and hidden params
|
|
response = ResponsesAPIResponse(
|
|
id="resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpncHQtNDtsbGxfcmVzcG9uc2VfaWQ6cmVzcF8xMjM",
|
|
object="response",
|
|
status="queued",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
# Add hidden params with model_id (simulating what base_process_llm_request does)
|
|
response._hidden_params = {"model_id": "model-deployment-id-123"}
|
|
|
|
# Mock request data
|
|
data = {
|
|
"model": "gpt-4",
|
|
"input": "Test input",
|
|
"background": True,
|
|
}
|
|
|
|
# Mock user_api_key_dict
|
|
user_api_key_dict = MagicMock()
|
|
user_api_key_dict.user_id = "test-user"
|
|
|
|
# Simulate the storage logic from endpoints.py
|
|
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
|
if response.status in ["queued", "in_progress"]:
|
|
# Get model_id from hidden params
|
|
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
|
model_id = hidden_params.get("model_id", None)
|
|
|
|
if model_id:
|
|
# Store in managed objects table using response.id directly
|
|
await mock_managed_files_obj.store_unified_object_id(
|
|
unified_object_id=response.id,
|
|
file_object=response,
|
|
litellm_parent_otel_span=None,
|
|
model_object_id=response.id,
|
|
file_purpose="response",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
# Verify store_unified_object_id was called
|
|
mock_managed_files_obj.store_unified_object_id.assert_called_once()
|
|
call_args = mock_managed_files_obj.store_unified_object_id.call_args
|
|
|
|
# Verify the arguments - unified_object_id should be response.id
|
|
assert call_args[1]["unified_object_id"] == response.id
|
|
assert call_args[1]["model_object_id"] == response.id
|
|
assert call_args[1]["file_purpose"] == "response"
|
|
assert call_args[1]["user_api_key_dict"] == user_api_key_dict
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_storage_for_non_background_requests(
|
|
self, mock_managed_files_obj, mock_proxy_logging_obj
|
|
):
|
|
"""Test that non-background requests are not stored"""
|
|
# Create a mock response
|
|
response = ResponsesAPIResponse(
|
|
id="resp_456",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
total_tokens=150,
|
|
),
|
|
)
|
|
|
|
# Mock request data without background flag
|
|
data = {
|
|
"model": "gpt-4",
|
|
"input": "Test input",
|
|
"background": False,
|
|
}
|
|
|
|
# Simulate the storage logic
|
|
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
|
if response.status in ["queued", "in_progress"]:
|
|
await mock_managed_files_obj.store_unified_object_id()
|
|
|
|
# Verify store_unified_object_id was NOT called
|
|
mock_managed_files_obj.store_unified_object_id.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_storage_for_completed_responses(
|
|
self, mock_managed_files_obj, mock_proxy_logging_obj
|
|
):
|
|
"""Test that completed responses are not stored"""
|
|
# Create a mock response with completed status
|
|
response = ResponsesAPIResponse(
|
|
id="resp_789",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
total_tokens=150,
|
|
),
|
|
)
|
|
|
|
# Mock request data with background flag
|
|
data = {
|
|
"model": "gpt-4",
|
|
"input": "Test input",
|
|
"background": True,
|
|
}
|
|
|
|
# Simulate the storage logic
|
|
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
|
if response.status in ["queued", "in_progress"]:
|
|
await mock_managed_files_obj.store_unified_object_id()
|
|
|
|
# Verify store_unified_object_id was NOT called (status is completed)
|
|
mock_managed_files_obj.store_unified_object_id.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_storage_without_model_id(
|
|
self, mock_managed_files_obj, mock_proxy_logging_obj
|
|
):
|
|
"""Test that responses without model_id in hidden params are not stored"""
|
|
# Create a mock response without hidden params
|
|
response = ResponsesAPIResponse(
|
|
id="resp_no_model",
|
|
object="response",
|
|
status="queued",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
# Mock request data with background flag
|
|
data = {
|
|
"model": "gpt-4",
|
|
"input": "Test input",
|
|
"background": True,
|
|
}
|
|
|
|
user_api_key_dict = MagicMock()
|
|
|
|
# Simulate the storage logic
|
|
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
|
if response.status in ["queued", "in_progress"]:
|
|
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
|
model_id = hidden_params.get("model_id", None)
|
|
|
|
if model_id: # This will be False
|
|
await mock_managed_files_obj.store_unified_object_id(
|
|
unified_object_id=response.id,
|
|
file_object=response,
|
|
litellm_parent_otel_span=None,
|
|
model_object_id=response.id,
|
|
file_purpose="response",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
|
|
# Verify store_unified_object_id was NOT called (no model_id)
|
|
mock_managed_files_obj.store_unified_object_id.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_error_handling_in_storage(
|
|
self, mock_managed_files_obj, mock_proxy_logging_obj
|
|
):
|
|
"""Test that errors during storage are handled gracefully"""
|
|
# Mock store_unified_object_id to raise an exception
|
|
mock_managed_files_obj.store_unified_object_id = AsyncMock(
|
|
side_effect=Exception("Database error")
|
|
)
|
|
|
|
response = ResponsesAPIResponse(
|
|
id="resp_error",
|
|
object="response",
|
|
status="queued",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
response._hidden_params = {"model_id": "test-model-id"}
|
|
|
|
data = {
|
|
"model": "gpt-4",
|
|
"input": "Test input",
|
|
"background": True,
|
|
}
|
|
|
|
user_api_key_dict = MagicMock()
|
|
user_api_key_dict.user_id = "test-user"
|
|
|
|
# Try to store - should not raise (error is caught in endpoints.py)
|
|
try:
|
|
if data.get("background") and isinstance(response, ResponsesAPIResponse):
|
|
if response.status in ["queued", "in_progress"]:
|
|
hidden_params = getattr(response, "_hidden_params", {}) or {}
|
|
model_id = hidden_params.get("model_id", None)
|
|
|
|
if model_id:
|
|
await mock_managed_files_obj.store_unified_object_id(
|
|
unified_object_id=response.id,
|
|
file_object=response,
|
|
litellm_parent_otel_span=None,
|
|
model_object_id=response.id,
|
|
file_purpose="response",
|
|
user_api_key_dict=user_api_key_dict,
|
|
)
|
|
except Exception:
|
|
# Exception should be caught and logged, not raised
|
|
pass
|
|
|
|
# Verify the method was called (even though it raised)
|
|
assert mock_managed_files_obj.store_unified_object_id.called
|
|
|
|
|
|
def _check_responses_cost_module_available():
|
|
"""Check if litellm_enterprise.proxy.common_utils.check_responses_cost module is available"""
|
|
try:
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import ( # noqa: F401
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not _check_responses_cost_module_available(),
|
|
reason="litellm_enterprise.proxy.common_utils.check_responses_cost module not available (enterprise-only feature)",
|
|
)
|
|
class TestCheckResponsesCost:
|
|
"""Tests for the CheckResponsesCost polling class"""
|
|
|
|
@pytest.fixture
|
|
def mock_prisma_client(self):
|
|
"""Create a mock Prisma client"""
|
|
client = MagicMock()
|
|
client.db = MagicMock()
|
|
client.db.litellm_managedobjecttable = MagicMock()
|
|
return client
|
|
|
|
@pytest.fixture
|
|
def mock_proxy_logging_obj(self):
|
|
"""Create a mock proxy logging object"""
|
|
return MagicMock()
|
|
|
|
@pytest.fixture
|
|
def mock_llm_router(self):
|
|
"""Create a mock LLM router"""
|
|
return MagicMock()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_initialization(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test CheckResponsesCost initialization"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
assert checker.proxy_logging_obj == mock_proxy_logging_obj
|
|
assert checker.prisma_client == mock_prisma_client
|
|
assert checker.llm_router == mock_llm_router
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_no_jobs(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test polling when there are no jobs"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
# Mock find_many to return empty list
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[]
|
|
)
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
# Should not raise any errors
|
|
await checker.check_responses_cost()
|
|
|
|
# Verify find_many was called with correct parameters (includes pagination)
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
|
where={
|
|
"status": {"in": ["queued", "in_progress"]},
|
|
"file_purpose": "response",
|
|
},
|
|
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
|
order={"created_at": "asc"},
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_completed_job(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test polling with a completed job"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
# Create a mock job
|
|
mock_job = MagicMock()
|
|
mock_job.id = "job-123"
|
|
mock_job.unified_object_id = "resp_test_id"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_id"}
|
|
|
|
# Mock find_many to return the job
|
|
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)
|
|
|
|
# Create a completed response
|
|
completed_response = ResponsesAPIResponse(
|
|
id="resp_test_id",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
total_tokens=150,
|
|
),
|
|
)
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
# Mock litellm.aget_responses to return completed response
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
|
mock_aget.return_value = completed_response
|
|
|
|
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)
|
|
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
|
|
]
|
|
assert len(completion_calls) == 1
|
|
assert completion_calls[0].kwargs["where"] == {
|
|
"id": "job-123",
|
|
"status": {"in": ["queued", "in_progress"]},
|
|
}
|
|
assert completion_calls[0].kwargs["data"]["status"] == "completed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_failed_job(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test polling with a failed job"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
# Create a mock job
|
|
mock_job = MagicMock()
|
|
mock_job.id = "job-456"
|
|
mock_job.unified_object_id = "resp_failed"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_failed"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
|
|
|
# Create a failed response
|
|
failed_response = ResponsesAPIResponse(
|
|
id="resp_failed",
|
|
object="response",
|
|
status="failed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
|
mock_aget.return_value = failed_response
|
|
|
|
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)
|
|
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
|
|
]
|
|
assert len(completion_calls) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_in_progress_job(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test polling with a job still in progress"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
# Create a mock job
|
|
mock_job = MagicMock()
|
|
mock_job.id = "job-789"
|
|
mock_job.unified_object_id = "resp_in_progress"
|
|
mock_job.created_by = "test-user"
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
|
|
|
# Create an in-progress response
|
|
in_progress_response = ResponsesAPIResponse(
|
|
id="resp_in_progress",
|
|
object="response",
|
|
status="in_progress",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
|
mock_aget.return_value = in_progress_response
|
|
|
|
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)
|
|
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
|
|
]
|
|
assert len(completion_calls) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_error_handling(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test that errors when querying responses are handled gracefully"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
# Create a mock job
|
|
mock_job = MagicMock()
|
|
mock_job.id = "job-error"
|
|
mock_job.unified_object_id = "resp_error"
|
|
mock_job.created_by = "test-user"
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
|
|
|
|
checker = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
|
|
# Mock litellm.aget_responses to raise an exception
|
|
with patch(
|
|
"litellm.aget_responses",
|
|
new_callable=AsyncMock,
|
|
side_effect=Exception("API error"),
|
|
):
|
|
# Should not raise - errors are caught and logged
|
|
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)
|
|
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
|
|
]
|
|
assert len(completion_calls) == 0
|