Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_attribution.py: 27%

1"""Spend attribution for batches created through a passthrough endpoint. 

2 

3The creating key and its tags are read off the passthrough request's metadata and 

4persisted on the managed object row, because the batch cost lands hours later in a 

5background poll that has no request to read them from. 

6""" 

7 

8import asyncio 

9from collections.abc import Mapping, Sequence 

10from typing import Final 

11 

12from litellm._logging import verbose_proxy_logger 

13from litellm.litellm_core_utils.safe_json_dumps import strip_null_bytes 

14 

15 

16def optional_str(value: object) -> str | None: 

17 return value if isinstance(value, str) else None 

18 

19 

20def _sanitized_str_tuple(value: object) -> tuple[str, ...] | None: 

21 if not isinstance(value, list): 

22 return None 

23 items: Final[Sequence[object]] = value 

24 return tuple(strip_null_bytes(tag) for tag in items if isinstance(tag, str)) 

25 

26 

27def is_collection_route(url_route: str, collection_suffix: str) -> bool: 

28 """Whether the route addresses the batch collection itself rather than one batch. 

29 A POST to the collection is the create; every id-scoped route is a retrieve, 

30 results or cancel. 

31 """ 

32 return url_route.split("?")[0].rstrip("/").endswith(collection_suffix) 

33 

34 

35def request_tags_from_metadata(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None: 

36 """Tags for the batch-cost spend row: the request's own tags when it sent any, 

37 otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a 

38 tagged key does not put its tags in the top-level metadata "tags" on the 

39 passthrough path) 

40 """ 

41 tags: Final = _sanitized_str_tuple(request_metadata.get("tags")) 

42 if tags: 

43 return tags 

44 key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata") 

45 if isinstance(key_auth_metadata, dict): 

46 return _sanitized_str_tuple(key_auth_metadata.get("tags")) 

47 return None 

48 

49 

50def log_batch_registration_result( 

51 finished: asyncio.Task[None], 

52 provider: str, 

53 unified_object_id: str, 

54 model_object_id: str, 

55 is_batch_create: bool, 

56) -> None: 

57 """Report the outcome of the fire-and-forget managed object write. A create that 

58 fails is not retried by a later poll, so its cost is never tracked at all. 

59 """ 

60 error: Final = finished.exception() if not finished.cancelled() else None 

61 if finished.cancelled() or error is not None: 

62 consequence: Final = ( 

63 "its cost will not be tracked" if is_batch_create else "its status and output file may be stale" 

64 ) 

65 verbose_proxy_logger.error( 

66 "Failed to store %s batch managed object with unified_object_id=%s, batch_id=%s; %s: %s", 

67 provider, 

68 unified_object_id, 

69 model_object_id, 

70 consequence, 

71 error, 

72 ) 

73 return 

74 verbose_proxy_logger.info( 

75 "Stored %s batch managed object with unified_object_id=%s, batch_id=%s", 

76 provider, 

77 unified_object_id, 

78 model_object_id, 

79 )