fix(batches): attribute Anthropic passthrough batch cost to the creating key, team and tags (#36468)

The Anthropic batch create never persisted the creating key's hashed token or its
request tags on the managed object, so when CheckBatchCost billed the batch hours
later there was nothing to attribute it to. Key spend, key budgets and tag spend
never moved for batch usage.

Persist both from the create, the way the Vertex passthrough already does, and
register the batch only from the collection route. An id-scoped route cannot
rebuild the unified object id, because it embeds the model and the model comes
from the create's request body, so it could only claim a row it did not create or
fail the model_object_id unique constraint.

The shared metadata helpers, the route predicate and the registration-result
logging now live in batch_attribution instead of being copied per provider. The
Anthropic write previously logged success unconditionally, before the
fire-and-forget task had run.

Resolves LIT-5288
This commit is contained in:
yucheng-berri 2026-08-11 20:52:43 -07:00 committed by GitHub
parent 7a55ca811b
commit 8bfb7772e4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 403 additions and 79 deletions

View file

@ -472,6 +472,8 @@ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float(
### ANTHROPIC CONSTANTS ###
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01")
ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02"
ANTHROPIC_BATCHES_ROUTE: Final = "/v1/messages/batches"
VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs"
ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = {
"low": 1,
"medium": 5,

View file

@ -39,11 +39,10 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
CallTypes.pass_through.value,
CallTypes.llm_passthrough_route.value,
CallTypes.allm_passthrough_route.value,
# CheckBatchCost's synthetic logging_obj for a completed managed batch only ever
# carries user_api_key_user_id (from LiteLLM_ManagedObjectTable.created_by) and
# user_api_key_team_id (from .team_id) -- both are None for batches created with
# the master key or a team-less key, since the table never stores the raw key
# hash. The batch already incurred real provider cost, so track it regardless.
# CheckBatchCost's synthetic logging_obj for a completed managed batch carries
# whatever LiteLLM_ManagedObjectTable stored at create time, and all of it is
# None for a batch created before those columns were persisted, or by the master
# key. The batch already incurred real provider cost, so track it regardless.
CallTypes.aretrieve_batch.value,
}
)

View file

@ -1,3 +1,4 @@
import asyncio
import json
from collections.abc import Sequence
from datetime import datetime
@ -7,6 +8,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import ANTHROPIC_BATCHES_ROUTE
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model
@ -20,6 +22,12 @@ from litellm.llms.anthropic.chat.handler import (
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
is_collection_route,
log_batch_registration_result,
optional_str,
request_tags_from_metadata,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
@ -833,13 +841,14 @@ class AnthropicPassthroughLoggingHandler:
# Store the managed object for cost tracking
# This will be picked up by check_batch_cost polling mechanism
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=litellm_batch_response,
model_object_id=batch_id,
logging_obj=logging_obj,
**kwargs,
)
if is_collection_route(url_route, ANTHROPIC_BATCHES_ROUTE):
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=litellm_batch_response,
model_object_id=batch_id,
logging_obj=logging_obj,
**kwargs,
)
# Create a batch job response for logging
litellm_model_response = ModelResponse()
@ -964,8 +973,12 @@ class AnthropicPassthroughLoggingHandler:
**kwargs,
) -> None:
"""
Store batch managed object for cost tracking.
Register a newly created batch for cost tracking.
This will be picked up by the check_batch_cost polling mechanism.
Only the create reaches here, so the row records the creating key and its tags.
An id-scoped route cannot rebuild the unified object id anyway: the model comes
from the create's request body, which a retrieve does not have.
"""
try:
# Get the managed files hook from the logging object
@ -981,7 +994,7 @@ class AnthropicPassthroughLoggingHandler:
user_api_key_dict: Final = UserAPIKeyAuth(
user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
api_key="",
api_key=optional_str(_request_metadata.get("user_api_key")),
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
@ -1003,9 +1016,7 @@ class AnthropicPassthroughLoggingHandler:
)
# Store the unified object for batch cost tracking
import asyncio
asyncio.create_task(
task: Final = asyncio.create_task(
managed_files_hook.store_unified_object_id(
unified_object_id=unified_object_id,
file_object=batch_object,
@ -1013,13 +1024,14 @@ class AnthropicPassthroughLoggingHandler:
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=request_tags_from_metadata(_request_metadata),
persist_attribution=True,
)
)
verbose_proxy_logger.info(
"Stored Anthropic batch managed object with unified_object_id=%s, batch_id=%s",
unified_object_id,
model_object_id,
task.add_done_callback(
lambda finished: log_batch_registration_result(
finished, "Anthropic", unified_object_id, model_object_id, is_batch_create=True
)
)
else:
verbose_proxy_logger.warning(

View file

@ -0,0 +1,78 @@
"""Spend attribution for batches created through a passthrough endpoint.
The creating key and its tags are read off the passthrough request's metadata and
persisted on the managed object row, because the batch cost lands hours later in a
background poll that has no request to read them from.
"""
import asyncio
from collections.abc import Mapping, Sequence
from typing import Final
from litellm._logging import verbose_proxy_logger
def optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def _optional_str_tuple(value: object) -> tuple[str, ...] | None:
if not isinstance(value, list):
return None
items: Final[Sequence[object]] = value
return tuple(tag for tag in items if isinstance(tag, str))
def is_collection_route(url_route: str, collection_suffix: str) -> bool:
"""Whether the route addresses the batch collection itself rather than one batch.
A POST to the collection is the create; every id-scoped route is a retrieve,
results or cancel.
"""
return url_route.split("?")[0].rstrip("/").endswith(collection_suffix)
def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None:
"""Tags for the batch-cost spend row: the request's own tags when it sent any,
otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a
tagged key does not put its tags in the top-level metadata "tags" on the
passthrough path)
"""
tags: Final = _optional_str_tuple(request_metadata.get("tags"))
if tags:
return tags
key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata")
if isinstance(key_auth_metadata, dict):
return _optional_str_tuple(key_auth_metadata.get("tags"))
return None
def log_batch_registration_result(
finished: asyncio.Task[None],
provider: str,
unified_object_id: str,
model_object_id: str,
is_batch_create: bool,
) -> None:
"""Report the outcome of the fire-and-forget managed object write. A create that
fails is not retried by a later poll, so its cost is never tracked at all.
"""
error: Final = finished.exception() if not finished.cancelled() else None
if finished.cancelled() or error is not None:
consequence: Final = (
"its cost will not be tracked" if is_batch_create else "its status and output file may be stale"
)
verbose_proxy_logger.error(
"Failed to store %s batch managed object with unified_object_id=%s, batch_id=%s; %s: %s",
provider,
unified_object_id,
model_object_id,
consequence,
error,
)
return
verbose_proxy_logger.info(
"Stored %s batch managed object with unified_object_id=%s, batch_id=%s",
provider,
unified_object_id,
model_object_id,
)

View file

@ -1,6 +1,5 @@
import asyncio
import re
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
@ -9,6 +8,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as VertexModelResponseIterator,
@ -18,6 +18,12 @@ from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
)
from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
is_collection_route,
log_batch_registration_result,
optional_str,
request_tags_from_metadata,
)
from litellm.types.utils import (
Choices,
EmbeddingResponse,
@ -41,32 +47,6 @@ else:
EndpointType = Any
def _optional_str(value: object) -> str | None:
return value if isinstance(value, str) else None
def _optional_str_tuple(value: object) -> tuple[str, ...] | None:
if not isinstance(value, list):
return None
items: Final = cast(list[object], value) # cast-ok: isinstance-narrowed; element type unknown
return tuple(tag for tag in items if isinstance(tag, str))
def _request_tags(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None:
"""Tags for the batch-cost spend row: the request's own tags when it sent any,
otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a
tagged key does not put its tags in the top-level metadata "tags" on the
passthrough path)
"""
tags: Final = _optional_str_tuple(request_metadata.get("tags"))
if tags:
return tags
key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata")
if isinstance(key_auth_metadata, dict):
return _optional_str_tuple(key_auth_metadata.get("tags"))
return None
class VertexPassthroughLoggingHandler:
@staticmethod
def vertex_passthrough_handler(
@ -685,7 +665,7 @@ class VertexPassthroughLoggingHandler:
# Store the managed object for cost tracking
# This will be picked up by check_batch_cost polling mechanism
is_batch_create: Final = url_route.split("?")[0].rstrip("/").endswith("batchPredictionJobs")
is_batch_create: Final = is_collection_route(url_route, VERTEX_BATCH_PREDICTION_JOBS_ROUTE)
VertexPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=litellm_batch_response,
@ -809,29 +789,6 @@ class VertexPassthroughLoggingHandler:
"kwargs": kwargs,
}
@staticmethod
def _log_batch_registration_result(
finished: asyncio.Task, unified_object_id: str, model_object_id: str, is_batch_create: bool
) -> None:
error: Final = finished.exception() if not finished.cancelled() else None
if finished.cancelled() or error is not None:
consequence: Final = (
"its cost will not be tracked" if is_batch_create else "its status and output file may be stale"
)
verbose_proxy_logger.error(
"Failed to store batch managed object with unified_object_id=%s, batch_id=%s; %s: %s",
unified_object_id,
model_object_id,
consequence,
error,
)
return
verbose_proxy_logger.info(
"Stored batch managed object with unified_object_id=%s, batch_id=%s",
unified_object_id,
model_object_id,
)
@staticmethod
def _store_batch_managed_object(
unified_object_id: str,
@ -863,7 +820,7 @@ class VertexPassthroughLoggingHandler:
user_api_key_dict: Final = UserAPIKeyAuth(
user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
api_key=_optional_str(_request_metadata.get("user_api_key")),
api_key=optional_str(_request_metadata.get("user_api_key")),
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
@ -893,14 +850,14 @@ class VertexPassthroughLoggingHandler:
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=_request_tags(_request_metadata),
request_tags=request_tags_from_metadata(_request_metadata),
persist_attribution=is_batch_create,
create_if_missing=is_batch_create,
)
)
task.add_done_callback(
lambda finished: VertexPassthroughLoggingHandler._log_batch_registration_result(
finished, unified_object_id, model_object_id, is_batch_create
lambda finished: log_batch_registration_result(
finished, "Vertex AI", unified_object_id, model_object_id, is_batch_create
)
)
else:

View file

@ -604,8 +604,8 @@ async def test_create_still_upserts_and_claims_attribution():
@pytest.mark.asyncio
async def test_default_callers_still_create_their_rows():
"""create_if_missing defaults to True, so the fine-tune, Responses and Anthropic
callers, none of which pass it, keep upserting exactly as before."""
"""create_if_missing defaults to True, so the fine-tune, Responses and managed
/v1/batches callers, none of which passes it, keep upserting exactly as before."""
managed_files, mock_prisma = _make_object_store_instance()
await managed_files.store_unified_object_id(

View file

@ -1,3 +1,4 @@
import asyncio
import json
import os
import sys
@ -17,6 +18,13 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passth
)
async def _drain_tasks():
"""Await the fire-and-forget managed object write and let its done callback run."""
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
await asyncio.gather(*pending, return_exceptions=True)
await asyncio.sleep(0)
class TestAnthropicLoggingHandlerModelFallback:
"""Test the model fallback logic in the anthropic passthrough logging handler."""
@ -925,6 +933,114 @@ class TestAnthropicBatchPassthroughCostTracking:
assert call_kwargs["user_api_key_dict"].user_id == expected_user_id
assert call_kwargs["user_api_key_dict"].team_id == expected_team_id
async def _store_with_metadata(self, mock_logging_obj, metadata):
mock_managed_files_hook = MagicMock()
mock_managed_files_hook.store_unified_object_id = AsyncMock()
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl,
patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
),
):
mock_pl.get_proxy_hook.return_value = mock_managed_files_hook
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id="uoi",
batch_object={"id": "b1", "object": "batch", "status": "validating"},
model_object_id="b1",
logging_obj=mock_logging_obj,
litellm_params={"metadata": metadata},
)
await _drain_tasks()
mock_managed_files_hook.store_unified_object_id.assert_awaited_once()
return mock_managed_files_hook.store_unified_object_id.call_args[1]
@pytest.mark.asyncio
async def test_create_persists_key_hash_and_tags(self, mock_logging_obj):
"""Regression (LIT-5288): the batch create must persist the creating key's hashed
token and its tags so CheckBatchCost can attribute the batch-cost spend row to the
key, team and tags. Before this fix the stored api_key was always "" and no tags
were stored, so key/team/tag spend and budgets never moved for batch usage."""
call_kwargs = await self._store_with_metadata(
mock_logging_obj,
{
"user_api_key": "hashed-key-a",
"user_api_key_user_id": "alice",
"user_api_key_team_id": "team-alpha",
"user_api_key_auth_metadata": {"tags": ["env:prod", 7, "team:ml"]},
},
)
assert call_kwargs["user_api_key_dict"].api_key == "hashed-key-a"
assert call_kwargs["request_tags"] == ("env:prod", "team:ml")
assert call_kwargs["persist_attribution"] is True
@pytest.mark.asyncio
async def test_failed_create_write_is_reported_not_swallowed(self, mock_logging_obj):
"""The managed object write is fire-and-forget, and only the create writes the row,
so a failed create is never back-filled by a later retrieve and that batch's cost
is never tracked. The failure has to reach the log instead of being reported as a
success."""
mock_managed_files_hook = MagicMock()
mock_managed_files_hook.store_unified_object_id = AsyncMock(
side_effect=RuntimeError("db down")
)
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl,
patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
) as mock_logger,
):
mock_pl.get_proxy_hook.return_value = mock_managed_files_hook
AnthropicPassthroughLoggingHandler._store_batch_managed_object(
unified_object_id="uoi",
batch_object={"id": "b1", "object": "batch", "status": "validating"},
model_object_id="b1",
logging_obj=mock_logging_obj,
litellm_params={"metadata": {"user_api_key": "hashed-key-a"}},
)
await _drain_tasks()
mock_logger.info.assert_not_called()
mock_logger.error.assert_called_once()
assert "its cost will not be tracked" in mock_logger.error.call_args[0]
assert "Anthropic" in mock_logger.error.call_args[0]
@pytest.mark.parametrize(
"url_route, registers",
[
("https://api.anthropic.com/v1/messages/batches", True),
("https://api.anthropic.com/v1/messages/batches/", True),
("https://api.anthropic.com/v1/messages/batches?limit=20", True),
("https://api.anthropic.com/v1/messages/batches/msgbatch_123", False),
("https://api.anthropic.com/v1/messages/batches/msgbatch_123/results", False),
("https://api.anthropic.com/v1/messages/batches/msgbatch_123/cancel", False),
],
)
def test_batch_is_registered_from_the_create_route_only(
self, mock_logging_obj, mock_httpx_response, mock_request_body, url_route, registers
):
"""Only a POST to the collection route registers the batch. Every id-scoped route
is a retrieve, results or cancel, and none of them can rebuild the unified object
id anyway: it embeds the model, which comes from the create's request body. Before
this gate an id-scoped route reached the store with a mismatched id, where it could
only either claim a row it did not create or fail the model_object_id unique
constraint."""
with patch.object(
AnthropicPassthroughLoggingHandler, "_store_batch_managed_object"
) as mock_store:
AnthropicPassthroughLoggingHandler.batch_creation_handler(
httpx_response=mock_httpx_response,
logging_obj=mock_logging_obj,
url_route=url_route,
result="success",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
request_body=mock_request_body,
)
assert mock_store.call_count == (1 if registers else 0)
def test_batch_creation_handler_failure_status_code(
self, mock_logging_obj, mock_request_body
):
@ -978,6 +1094,7 @@ class TestAnthropicBatchPassthroughCostTracking:
batch_object=batch_object,
model_object_id="msgbatch_123",
logging_obj=mock_logging_obj,
is_batch_create=True,
user_id="test-user",
)

View file

@ -0,0 +1,159 @@
import asyncio
from unittest.mock import patch
import pytest
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
is_collection_route,
log_batch_registration_result,
optional_str,
request_tags_from_metadata,
)
@pytest.mark.parametrize(
"value, expected",
[("a", "a"), ("", ""), (None, None), (7, None), (["a"], None)],
)
def test_optional_str(value, expected):
assert optional_str(value) == expected
class TestRequestTagsFromMetadata:
"""Tags for the batch-cost spend row. These feed LiteLLM_ManagedObjectTable.request_tags,
which is the only record of the creating request's tags by the time CheckBatchCost bills
the batch hours later."""
@pytest.mark.parametrize(
"metadata, expected",
[
# a request that sent its own tags (x-litellm-tags header or body metadata)
({"tags": ["req:a", "req:b"]}, ("req:a", "req:b")),
# request tags win over the key's own tags
(
{"tags": ["req:a"], "user_api_key_auth_metadata": {"tags": ["key:b"]}},
("req:a",),
),
# no request tags: fall back to the tags the key itself carries, because a
# tagged key does not put its tags in the top-level metadata on this path
({"user_api_key_auth_metadata": {"tags": ["key:b"]}}, ("key:b",)),
# an empty request tag list is not a selection, so the key's tags still apply
(
{"tags": [], "user_api_key_auth_metadata": {"tags": ["key:b"]}},
("key:b",),
),
# neither: no tags on the spend row
({}, None),
# order is preserved, so the spend row is reproducible
({"tags": ["z", "a", "m"]}, ("z", "a", "m")),
],
)
def test_precedence(self, metadata, expected):
assert request_tags_from_metadata(metadata) == expected
@pytest.mark.parametrize(
"raw, expected",
[
# non-string entries are dropped rather than crashing the create
(["env:prod", 7, None, "team:ml"], ("env:prod", "team:ml")),
# nothing usable survives, so this is treated as no request tags at all
([7, None], None),
# a non-list is not a tag list
("env:prod", None),
({"env": "prod"}, None),
(None, None),
],
)
def test_malformed_tags_are_dropped(self, raw, expected):
assert request_tags_from_metadata({"tags": raw}) == expected
def test_malformed_key_auth_metadata_is_ignored(self):
assert request_tags_from_metadata({"user_api_key_auth_metadata": "nope"}) is None
@pytest.mark.parametrize(
"url_route, suffix, expected",
[
("https://api.anthropic.com/v1/messages/batches", "/v1/messages/batches", True),
("https://api.anthropic.com/v1/messages/batches/", "/v1/messages/batches", True),
("https://api.anthropic.com/v1/messages/batches?limit=20", "/v1/messages/batches", True),
("https://api.anthropic.com/v1/messages/batches/msgbatch_1", "/v1/messages/batches", False),
# a proxied base with a path prefix still resolves, because this is a suffix match
("https://gateway.internal/anthropic/v1/messages/batches", "/v1/messages/batches", True),
("https://aiplatform.googleapis.com/v1/projects/p/locations/l/batchPredictionJobs", "batchPredictionJobs", True),
("https://aiplatform.googleapis.com/v1/projects/p/locations/l/batchPredictionJobs/9", "batchPredictionJobs", False),
],
)
def test_is_collection_route(url_route, suffix, expected):
assert is_collection_route(url_route, suffix) is expected
class TestLogBatchRegistrationResult:
"""The managed object write is fire and forget, so its outcome only ever reaches an
operator through this log line."""
@staticmethod
async def _finished_task(coro):
task = asyncio.ensure_future(coro)
await asyncio.gather(task, return_exceptions=True)
return task
@pytest.mark.asyncio
async def test_success_names_the_provider(self):
async def ok():
return None
task = await self._finished_task(ok())
with patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
) as logger:
log_batch_registration_result(task, "Anthropic", "uoi", "b1", is_batch_create=True)
logger.error.assert_not_called()
logger.info.assert_called_once()
assert "Anthropic" in logger.info.call_args[0]
@pytest.mark.asyncio
async def test_a_failed_create_says_the_cost_is_lost(self):
async def boom():
raise RuntimeError("db down")
task = await self._finished_task(boom())
with patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
) as logger:
log_batch_registration_result(task, "Vertex AI", "uoi", "b1", is_batch_create=True)
logger.info.assert_not_called()
assert "its cost will not be tracked" in logger.error.call_args[0]
@pytest.mark.asyncio
async def test_a_failed_refresh_says_the_row_is_stale(self):
async def boom():
raise RuntimeError("db down")
task = await self._finished_task(boom())
with patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
) as logger:
log_batch_registration_result(task, "Vertex AI", "uoi", "b1", is_batch_create=False)
logger.info.assert_not_called()
assert "its status and output file may be stale" in logger.error.call_args[0]
@pytest.mark.asyncio
async def test_a_cancelled_write_is_reported_not_reraised(self):
async def slow():
await asyncio.sleep(60)
task = asyncio.ensure_future(slow())
await asyncio.sleep(0)
task.cancel()
await asyncio.gather(task, return_exceptions=True)
with patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution.verbose_proxy_logger"
) as logger:
log_batch_registration_result(task, "Anthropic", "uoi", "b1", is_batch_create=True)
logger.info.assert_not_called()
logger.error.assert_called_once()