From 8bfb7772e4effe3f699f5ac327320cc9a6b20f9b Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 11 Aug 2026 20:52:43 -0700 Subject: [PATCH] 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 --- litellm/constants.py | 2 + .../proxy/hooks/proxy_track_cost_callback.py | 9 +- .../anthropic_passthrough_logging_handler.py | 46 +++-- .../batch_attribution.py | 78 +++++++++ .../vertex_passthrough_logging_handler.py | 67 ++------ .../proxy/test_managed_files_hook.py | 4 +- ...t_anthropic_passthrough_logging_handler.py | 117 +++++++++++++ .../test_batch_attribution.py | 159 ++++++++++++++++++ 8 files changed, 403 insertions(+), 79 deletions(-) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py diff --git a/litellm/constants.py b/litellm/constants.py index c9d9ff155ff..8c3541067a5 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0e22b5324c1..4551680e1b4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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, } ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 9fb967e570f..d0ca61e1fd1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py new file mode 100644 index 00000000000..e94145b3efe --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py @@ -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, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 7dee0e4a364..621b3ff9c83 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -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: diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 34cd0cabc2c..d260f79a09a 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -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( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 947a7a64beb..1c98e3ce535 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -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", ) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py new file mode 100644 index 00000000000..5109cbb5991 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_batch_attribution.py @@ -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()