mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(enterprise): expire a background response only on a 404 naming it and skip a bad deployment entry per job
This commit is contained in:
mateo-berri
parent
bc49a02530
commit
76f7e7d49e
2 changed files with 124 additions and 41 deletions
|
|
@ -47,8 +47,8 @@ def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_Manag
|
|||
return ManagedObjectRepository(prisma_client).table
|
||||
|
||||
|
||||
def _is_provider_not_found(error: Exception) -> bool:
|
||||
return getattr(error, "status_code", None) == 404
|
||||
def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -> bool:
|
||||
return getattr(error, "status_code", None) == 404 and provider_response_id in str(error)
|
||||
|
||||
|
||||
class CheckResponsesCost:
|
||||
|
|
@ -196,13 +196,16 @@ class CheckResponsesCost:
|
|||
if model_name:
|
||||
litellm_metadata["model"] = model_name
|
||||
litellm_metadata["model_group"] = model_name # Use same value for model_group
|
||||
via_router = self._resolve_deployment(responses_id_security)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
via_router = self._resolve_deployment(responses_id_security)
|
||||
provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id(
|
||||
responses_id_security
|
||||
).get("response_id", responses_id_security)
|
||||
try:
|
||||
response = await self._get_response(
|
||||
response_id=responses_id_security,
|
||||
|
|
@ -215,7 +218,7 @@ class CheckResponsesCost:
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
if via_router and _is_provider_not_found(e):
|
||||
if via_router and _is_response_gone_at_provider(e, provider_response_id):
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ Unit tests for CheckResponsesCost class
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -12,6 +14,19 @@ from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
|||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
|
||||
|
||||
class _RecordingManagedObjectTable:
|
||||
def __init__(self, rows: tuple[object, ...]) -> None:
|
||||
self.rows: Final = rows
|
||||
self.updates: tuple[tuple[Mapping[str, object], Mapping[str, object]], ...] = ()
|
||||
|
||||
async def find_many(self, **query: object) -> tuple[object, ...]:
|
||||
return self.rows
|
||||
|
||||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int:
|
||||
self.updates = (*self.updates, (where, data))
|
||||
return len(self.rows)
|
||||
|
||||
|
||||
class TestCheckResponsesCost:
|
||||
"""Test suite for CheckResponsesCost class"""
|
||||
|
||||
|
|
@ -387,16 +402,12 @@ class TestCheckResponsesCost:
|
|||
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404"}
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
side_effect=litellm.NotFoundError(
|
||||
message="Response not found", model="gpt-5", llm_provider="openai"
|
||||
message="Response with id 'resp_upstream_404' not found.", model="gpt-5", llm_provider="openai"
|
||||
)
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
table = _RecordingManagedObjectTable(rows=(mock_job,))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
|
||||
await check_responses_cost_instance.check_responses_cost()
|
||||
|
|
@ -404,10 +415,7 @@ class TestCheckResponsesCost:
|
|||
mock_sdk_aget.assert_not_called()
|
||||
mock_llm_router.get_deployment.assert_called_once_with(model_id="deployment-404")
|
||||
assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == encoded_response_id
|
||||
update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many
|
||||
update_many.assert_awaited_once()
|
||||
assert update_many.call_args.kwargs["where"] == {"id": {"in": ["job-404"]}}
|
||||
assert update_many.call_args.kwargs["data"] == {"status": "stale_expired"}
|
||||
assert table.updates == (({"id": {"in": ["job-404"]}}, {"status": "stale_expired"}),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_marks_mapped_provider_404_stale_expired(
|
||||
|
|
@ -456,21 +464,14 @@ class TestCheckResponsesCost:
|
|||
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404-mapped"}
|
||||
mock_llm_router.aget_responses = AsyncMock(side_effect=mapped.value)
|
||||
|
||||
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
|
||||
)
|
||||
table = _RecordingManagedObjectTable(rows=(mock_job,))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
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_called()
|
||||
update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many
|
||||
update_many.assert_awaited_once()
|
||||
assert update_many.call_args.kwargs["where"] == {"id": {"in": ["job-404-mapped"]}}
|
||||
assert update_many.call_args.kwargs["data"] == {"status": "stale_expired"}
|
||||
assert table.updates == (({"id": {"in": ["job-404-mapped"]}}, {"status": "stale_expired"}),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_404_without_router_deployment_keeps_row_for_retry(
|
||||
|
|
@ -494,12 +495,8 @@ class TestCheckResponsesCost:
|
|||
|
||||
mock_llm_router.get_deployment.return_value = 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=0
|
||||
)
|
||||
table = _RecordingManagedObjectTable(rows=(mock_job,))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
with patch(
|
||||
"litellm.aget_responses",
|
||||
|
|
@ -512,8 +509,96 @@ class TestCheckResponsesCost:
|
|||
|
||||
mock_sdk_aget.assert_awaited_once()
|
||||
mock_llm_router.aget_responses.assert_not_called()
|
||||
update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many
|
||||
assert update_many.call_args_list == []
|
||||
assert table.updates == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_404_not_naming_the_response_keeps_row_for_retry(
|
||||
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""A 404 that does not name the response (a gateway or a reconfigured deployment) is retried, not expired."""
|
||||
import litellm
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="azure",
|
||||
model_id="deployment-renamed",
|
||||
response_id="resp_upstream_still_there",
|
||||
)
|
||||
|
||||
mock_job = MagicMock()
|
||||
mock_job.unified_object_id = encoded_response_id
|
||||
mock_job.created_by = "test-user"
|
||||
mock_job.id = "job-404-other"
|
||||
mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id}
|
||||
|
||||
mock_llm_router.get_deployment.return_value = {"model_id": "deployment-renamed"}
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
side_effect=litellm.NotFoundError(
|
||||
message="Resource not found", model="gpt-5", llm_provider="azure"
|
||||
)
|
||||
)
|
||||
|
||||
table = _RecordingManagedObjectTable(rows=(mock_job,))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
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_called()
|
||||
assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == encoded_response_id
|
||||
assert table.updates == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_deployment_lookup_error_skips_only_that_job(
|
||||
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
|
||||
):
|
||||
"""A deployment lookup that raises skips its own job and the cycle still records the next job."""
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
broken_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id="deployment-broken", response_id="resp_upstream_broken"
|
||||
)
|
||||
healthy_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id="deployment-healthy", response_id="resp_upstream_healthy"
|
||||
)
|
||||
|
||||
broken_job = MagicMock()
|
||||
broken_job.unified_object_id = broken_response_id
|
||||
broken_job.created_by = "test-user"
|
||||
broken_job.id = "job-broken"
|
||||
broken_job.file_object = {"model": "gpt-5.5", "id": broken_response_id}
|
||||
|
||||
healthy_job = MagicMock()
|
||||
healthy_job.unified_object_id = healthy_response_id
|
||||
healthy_job.created_by = "test-user"
|
||||
healthy_job.id = "job-healthy"
|
||||
healthy_job.file_object = {"model": "gpt-5.5", "id": healthy_response_id}
|
||||
|
||||
mock_llm_router.get_deployment.side_effect = [
|
||||
Exception("Model invalid format - <class 'str'>"),
|
||||
{"model_id": "deployment-healthy"},
|
||||
]
|
||||
mock_llm_router.aget_responses = AsyncMock(
|
||||
return_value=ResponsesAPIResponse(
|
||||
id=healthy_response_id,
|
||||
object="response",
|
||||
status="completed",
|
||||
created_at=int(datetime.now().timestamp()),
|
||||
output=[],
|
||||
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
|
||||
)
|
||||
)
|
||||
|
||||
table = _RecordingManagedObjectTable(rows=(broken_job, healthy_job))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
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_called()
|
||||
mock_llm_router.aget_responses.assert_awaited_once()
|
||||
assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == healthy_response_id
|
||||
assert table.updates == (({"id": {"in": ["job-healthy"]}}, {"status": "completed"}),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_non_404_error_keeps_row_for_retry(
|
||||
|
|
@ -542,20 +627,15 @@ class TestCheckResponsesCost:
|
|||
)
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
table = _RecordingManagedObjectTable(rows=(mock_job,))
|
||||
mock_prisma_client.db.litellm_managedobjecttable = table
|
||||
|
||||
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_called()
|
||||
mock_llm_router.aget_responses.assert_awaited_once()
|
||||
update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many
|
||||
assert update_many.call_args_list == []
|
||||
assert table.updates == ()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_responses_cost_multiple_jobs(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue