mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
4bf40c4e8d
* fix(logging): stop billing and logging response reads as LLM calls Retrieving, deleting or cancelling a stored response, and vector store management calls, run through the same logging lifecycle as inference. A retrieved response replays the usage of the call that created it, so every read priced it again and wrote a second spend log row for the same tokens. Non-inference calls now cost 0, report no usage, log no placeholder chat message, and get a litellm.responses_management operation name instead of reading as chat. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): keep billing background response jobs after the poll Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): use an empty list for read-call messages A tuple matches no branch in the loggers that walk this value, so lunary's parse_messages falls through to clean_message and raises AttributeError on the success hook. An empty list reads as no messages everywhere: it satisfies the isinstance(list) checks in newrelic, mlflow and datadog, iterates zero times in traceloop and helicone, and is what StandardLoggingPayload.messages is typed to hold. None would be type-legal too but is not iterable, so it trades one crash for another in mlflow and traceloop. * fix(otel): stop the legacy emitter reporting replayed tokens on response reads The zeroing so far lands in the standard logging payload, which the legacy OpenTelemetry emitter does not read for usage: it takes prompt, completion and total tokens straight off the response object, so a retrieval span still carried the token counts of the call that produced the response, and the token usage histogram still recorded them. That emitter is the default, so the spend row said zero while the trace said otherwise. The background cost poller keeps its counts, the same exemption the pricing path already makes. * fix(logging): keep billing a background response when its retrieval is read A response created with background=true comes back queued and carries no usage, so its create bills nothing. The retrieval that first sees the finished job is the only place that job's tokens are ever visible, and pricing every read at zero therefore loses the spend outright rather than deduplicating it. On a proxy without the enterprise cost poller a background job ended up costing $0 end to end. is_unbilled_non_inference_call now takes the response it is deciding about and treats a background response the same way it already treats the poller's own read, which is the same exemption seen from the other side. The legacy OpenTelemetry emitter's time per output token metric picks up the read gate it was missing, so it stops dividing a read's latency by the replayed completion token count. * test(proxy): pass the read response to the non-inference predicate The poller test called is_unbilled_non_inference_call with the pre-background signature, so it broke when the predicate gained the response it classifies. It now hands the predicate a foreground read, and asserts that the same read is free without the origin stamp, so the stamp is what the test proves. * fix(otel): stop the v2 metrics recorder reporting replayed tokens on response reads The v2 span builder sources usage from the standard logging payload, so the earlier fix already zeroes it there. The metrics recorder reads response_obj directly, so a responses-management read still recorded the original generation's tokens into gen_ai.client.token.usage and divided generation time by them for gen_ai.server.time_per_output_token. The read still records operation and response duration, under the litellm.responses_management operation, so it stays observable. * fix(proxy): keep the response-cost headers on calls priced at zero Pricing responses reads and vector-store management routes at zero dropped the whole x-litellm-response-cost family off those replies. The header build reads a falsy zero as a cost this response never recorded and filters it out, and a call that returns before pricing stores no cost breakdown for the component headers to read, so a client parsing the cost off a read got a KeyError where it had previously been handed a number. Those calls now advertise the family at zero. Retrieving a background response, and the cost poller's read of one, still report their real cost. The params-taking form of the predicate moves from opentelemetry into internal_call_metadata so the proxy header build and the OTEL recorders share one copy. * fix(proxy): report a zero cost split only under a zero cost total The component headers were filled from call-type membership alone, while the total they sit beside keeps its real value when the read priced normally, so a breakdown that had not landed by the time headers were built could advertise a real total next to an all-zero split. The split is now reported as zero only when the total agrees with it, and is otherwise left absent. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yucheng Zhu <yucheng@berri.ai>
793 lines
30 KiB
Python
793 lines
30 KiB
Python
"""
|
|
Unit tests for CheckResponsesCost class
|
|
"""
|
|
|
|
import asyncio
|
|
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 TestCheckResponsesCost:
|
|
"""Test suite for CheckResponsesCost 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 ProxyLogging object"""
|
|
logging_obj = MagicMock()
|
|
logging_obj.get_proxy_hook = MagicMock(return_value=None)
|
|
return logging_obj
|
|
|
|
@pytest.fixture
|
|
def mock_llm_router(self):
|
|
"""Create a mock LLM Router"""
|
|
router = MagicMock()
|
|
router.aget_responses = AsyncMock()
|
|
router.get_deployment = MagicMock()
|
|
return router
|
|
|
|
@pytest.fixture
|
|
def check_responses_cost_instance(
|
|
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Create a CheckResponsesCost instance with mocked dependencies"""
|
|
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
|
|
CheckResponsesCost,
|
|
)
|
|
|
|
instance = CheckResponsesCost(
|
|
proxy_logging_obj=mock_proxy_logging_obj,
|
|
prisma_client=mock_prisma_client,
|
|
llm_router=mock_llm_router,
|
|
)
|
|
# Mock _expire_stale_rows (raw SQL) so _cleanup_stale_managed_objects
|
|
# succeeds without a real DB. Individual tests can override this.
|
|
instance._expire_stale_rows = AsyncMock(return_value=0)
|
|
return instance
|
|
|
|
def test_initialization(self, check_responses_cost_instance):
|
|
"""Test that CheckResponsesCost initializes correctly"""
|
|
assert check_responses_cost_instance.proxy_logging_obj is not None
|
|
assert check_responses_cost_instance.prisma_client is not None
|
|
assert check_responses_cost_instance.llm_router is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_no_jobs(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost when there are no jobs to process"""
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[]
|
|
)
|
|
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
# Verify find_many was called with pagination params
|
|
find_many_call = (
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
|
|
)
|
|
assert find_many_call[1]["where"] == {
|
|
"status": {"in": ["queued", "in_progress"]},
|
|
"file_purpose": "response",
|
|
}
|
|
assert find_many_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE
|
|
assert find_many_call[1]["order"] == {"created_at": "asc"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cleanup_stale_managed_objects(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Stale rows are expired via _expire_stale_rows before polling."""
|
|
from litellm.constants import STALE_OBJECT_CLEANUP_BATCH_SIZE
|
|
|
|
check_responses_cost_instance._expire_stale_rows = AsyncMock(return_value=5)
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[]
|
|
)
|
|
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
# _expire_stale_rows should have been called with a cutoff datetime and batch size
|
|
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
|
call_args = check_responses_cost_instance._expire_stale_rows.call_args
|
|
assert call_args[0][1] == STALE_OBJECT_CLEANUP_BATCH_SIZE
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_completed_response(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test check_responses_cost with a completed response"""
|
|
# Mock job with response ID
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_123"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-123"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_123"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
|
|
# Mock completed response
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_123",
|
|
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(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check with mocked litellm.aget_responses
|
|
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()
|
|
|
|
# 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"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_failed_response(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Test check_responses_cost with a failed response"""
|
|
# Mock job
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_456"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-456"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_456"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
|
|
# Mock failed response
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_456",
|
|
object="response",
|
|
status="failed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check
|
|
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()
|
|
|
|
# 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"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_cancelled_response(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost with a cancelled response"""
|
|
# Mock job
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_789"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-789"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_789"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
|
|
# Mock cancelled response
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_789",
|
|
object="response",
|
|
status="cancelled",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check
|
|
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()
|
|
|
|
# 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"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_in_progress_response(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost with a response still in progress"""
|
|
# Mock job
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_in_progress"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-in-progress"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_in_progress"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
|
|
# Mock in-progress response
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_in_progress",
|
|
object="response",
|
|
status="in_progress",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check
|
|
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()
|
|
|
|
# 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
|
|
# Stale cleanup still ran via _expire_stale_rows
|
|
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_queued_response(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost with a queued response"""
|
|
# Mock job
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_queued"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-queued"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_queued"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job]
|
|
)
|
|
|
|
# Mock queued response
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_queued",
|
|
object="response",
|
|
status="queued",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check
|
|
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()
|
|
|
|
# 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
|
|
# Stale cleanup still ran via _expire_stale_rows
|
|
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_exception(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost handles exceptions gracefully"""
|
|
# Mock job
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_error"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-error"
|
|
mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_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=0
|
|
)
|
|
|
|
# Run the check with mocked exception
|
|
with patch(
|
|
"litellm.aget_responses",
|
|
new_callable=AsyncMock,
|
|
side_effect=Exception("Provider error"),
|
|
):
|
|
# Should not raise, just skip the job
|
|
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
|
|
# Stale cleanup still ran via _expire_stale_rows
|
|
check_responses_cost_instance._expire_stale_rows.assert_called_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_multiple_jobs(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""Test check_responses_cost with multiple jobs"""
|
|
# Mock multiple jobs
|
|
mock_job1 = MagicMock()
|
|
mock_job1.unified_object_id = "resp_test_1"
|
|
mock_job1.created_by = "user1"
|
|
mock_job1.id = "job-1"
|
|
mock_job1.file_object = {"model": "gpt-4o", "id": "resp_test_1"}
|
|
|
|
mock_job2 = MagicMock()
|
|
mock_job2.unified_object_id = "resp_test_2"
|
|
mock_job2.created_by = "user2"
|
|
mock_job2.id = "job-2"
|
|
mock_job2.file_object = {"model": "gpt-4o", "id": "resp_test_2"}
|
|
|
|
mock_job3 = MagicMock()
|
|
mock_job3.unified_object_id = "resp_test_3"
|
|
mock_job3.created_by = "user3"
|
|
mock_job3.id = "job-3"
|
|
mock_job3.file_object = {"model": "gpt-4o", "id": "resp_test_3"}
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
|
return_value=[mock_job1, mock_job2, mock_job3]
|
|
)
|
|
|
|
# Mock responses - 2 completed, 1 in progress
|
|
mock_response1 = ResponsesAPIResponse(
|
|
id="resp_1",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
total_tokens=150,
|
|
),
|
|
)
|
|
|
|
mock_response2 = ResponsesAPIResponse(
|
|
id="resp_2",
|
|
object="response",
|
|
status="in_progress",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
mock_response3 = ResponsesAPIResponse(
|
|
id="resp_3",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=ResponseAPIUsage(
|
|
input_tokens=200,
|
|
output_tokens=100,
|
|
total_tokens=300,
|
|
),
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
|
|
return_value=0
|
|
)
|
|
|
|
# Run the check
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
|
|
mock_aget.side_effect = [mock_response1, mock_response2, mock_response3]
|
|
|
|
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"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_encoded_response_id_is_fetched_through_router(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""
|
|
Regression test for https://github.com/BerriAI/litellm/issues/35131
|
|
|
|
A background response created against a deployment whose credentials only
|
|
exist in the config (e.g. Azure api_base/api_key) must be fetched through
|
|
the router so the deployment credentials are applied. Calling
|
|
litellm.aget_responses directly only sees provider env vars, fails, and
|
|
leaves the row in "queued" forever.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
|
custom_llm_provider="azure",
|
|
model_id="deployment-abc",
|
|
response_id="resp_upstream_123",
|
|
)
|
|
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = encoded_response_id
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-router"
|
|
mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id}
|
|
|
|
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
|
|
)
|
|
|
|
mock_llm_router.aget_responses = AsyncMock(
|
|
return_value=ResponsesAPIResponse(
|
|
id=encoded_response_id,
|
|
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,
|
|
side_effect=AssertionError(
|
|
"must not bypass the router for a deployment-scoped response id"
|
|
),
|
|
) as mock_sdk_aget:
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
mock_sdk_aget.assert_not_called()
|
|
assert (
|
|
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]["data"]["status"] == "completed"
|
|
assert calls[0][1]["where"]["id"]["in"] == ["job-router"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_encrypted_response_id_is_fetched_through_router(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, monkeypatch
|
|
):
|
|
"""
|
|
Rows store the *encrypted* response id when responses id security is on.
|
|
After decryption the id still carries the deployment model_id, so the
|
|
fetch must go through the router (issue #35131).
|
|
"""
|
|
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
from litellm.types.utils import SpecialEnums
|
|
|
|
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids")
|
|
|
|
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
|
custom_llm_provider="openai",
|
|
model_id="deployment-xyz",
|
|
response_id="resp_upstream_456",
|
|
)
|
|
encrypted_response_id = "resp_" + str(
|
|
encrypt_value_helper(
|
|
value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
|
|
encoded_response_id, "test-user", "test-team"
|
|
)
|
|
)
|
|
)
|
|
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = encrypted_response_id
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-encrypted"
|
|
mock_job.file_object = {"model": "gpt-5", "id": encrypted_response_id}
|
|
|
|
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
|
|
)
|
|
|
|
mock_llm_router.aget_responses = AsyncMock(
|
|
return_value=ResponsesAPIResponse(
|
|
id=encoded_response_id,
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
)
|
|
|
|
with patch(
|
|
"litellm.aget_responses",
|
|
new_callable=AsyncMock,
|
|
side_effect=AssertionError(
|
|
"must not bypass the router for a deployment-scoped response id"
|
|
),
|
|
) as mock_sdk_aget:
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
mock_sdk_aget.assert_not_called()
|
|
assert (
|
|
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"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_response_id_without_model_id_uses_sdk(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""Ids that carry no deployment info can't be routed, so fall back to the SDK."""
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_plain_upstream_id"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-plain"
|
|
mock_job.file_object = {"model": "gpt-5", "id": "resp_plain_upstream_id"}
|
|
|
|
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
|
|
)
|
|
mock_llm_router.aget_responses = AsyncMock(
|
|
side_effect=AssertionError("router cannot route an id without a model_id")
|
|
)
|
|
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_plain_upstream_id",
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
|
|
mock_sdk_aget.return_value = mock_response
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
mock_sdk_aget.assert_called_once()
|
|
mock_llm_router.aget_responses.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_deployment_falls_back_to_sdk(
|
|
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
|
):
|
|
"""
|
|
An encoded id whose deployment was removed from the router must fall back
|
|
to the SDK so provider env credentials can still retrieve it, instead of
|
|
failing every poll cycle until stale expiration.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
|
custom_llm_provider="openai",
|
|
model_id="deployment-deleted",
|
|
response_id="resp_upstream_789",
|
|
)
|
|
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = encoded_response_id
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-missing-deployment"
|
|
mock_job.file_object = {"model": "gpt-5", "id": encoded_response_id}
|
|
|
|
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
|
|
)
|
|
mock_llm_router.get_deployment = MagicMock(return_value=None)
|
|
mock_llm_router.aget_responses = AsyncMock(
|
|
side_effect=AssertionError("router has no deployment for this model_id")
|
|
)
|
|
|
|
mock_response = ResponsesAPIResponse(
|
|
id=encoded_response_id,
|
|
object="response",
|
|
status="completed",
|
|
created_at=int(datetime.now().timestamp()),
|
|
output=[],
|
|
usage=None,
|
|
)
|
|
|
|
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
|
|
mock_sdk_aget.return_value = mock_response
|
|
await check_responses_cost_instance.check_responses_cost()
|
|
|
|
mock_llm_router.get_deployment.assert_called_once_with(model_id="deployment-deleted")
|
|
mock_llm_router.aget_responses.assert_not_called()
|
|
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"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_with_incomplete_response(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""'incomplete' is terminal in the Responses API, so the row must not stay queued."""
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_incomplete"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-incomplete"
|
|
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_incomplete"}
|
|
|
|
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
|
|
)
|
|
|
|
mock_response = ResponsesAPIResponse(
|
|
id="resp_incomplete",
|
|
object="response",
|
|
status="incomplete",
|
|
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()
|
|
|
|
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"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_check_responses_cost_no_model_in_file_object(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""When file_object has no 'model' key, model_name is None and metadata skips model fields."""
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_no_model"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-no-model"
|
|
mock_job.file_object = {} # no "model" key → model_name=None branch
|
|
|
|
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
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status = "completed"
|
|
|
|
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()
|
|
|
|
# aget_responses should be called without model metadata
|
|
call_kwargs = mock_aget.call_args[1]
|
|
assert "model" not in call_kwargs.get("litellm_metadata", {})
|
|
assert "model_group" not in call_kwargs.get("litellm_metadata", {})
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poll_stamps_internal_call_origin_so_the_read_is_billed(
|
|
self, check_responses_cost_instance, mock_prisma_client
|
|
):
|
|
"""A background create returns queued with no usage, so this poll's retrieval is the only
|
|
place the job's spend is ever seen. Without the origin stamp it is priced at zero like a
|
|
user-facing read (LIT-5602) and the job is never billed."""
|
|
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
|
from litellm.litellm_core_utils.internal_call_metadata import (
|
|
is_unbilled_non_inference_call,
|
|
)
|
|
|
|
mock_job = MagicMock()
|
|
mock_job.unified_object_id = "resp_test_billed"
|
|
mock_job.created_by = "test-user"
|
|
mock_job.id = "job-billed"
|
|
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_billed"}
|
|
|
|
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
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status = "completed"
|
|
|
|
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()
|
|
|
|
metadata = mock_aget.call_args[1]["litellm_metadata"]
|
|
foreground_read = {"background": False}
|
|
assert metadata[INTERNAL_CALL_ORIGIN_METADATA_KEY] == "background_response_cost_poll"
|
|
assert is_unbilled_non_inference_call("aget_responses", metadata, foreground_read) is False
|
|
assert is_unbilled_non_inference_call("aget_responses", None, foreground_read) is True
|