Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/prompt_caching_requests.py: 89%

1from collections.abc import Callable, Mapping 

2from datetime import datetime, timezone 

3from types import MappingProxyType 

4from typing import TYPE_CHECKING, Annotated, Final 

5 

6from fastapi import APIRouter, Depends, HTTPException, Query 

7from pydantic import BaseModel, Json, TypeAdapter 

8 

9from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth, user_api_key_has_admin_view 

10from litellm.proxy.auth.user_api_key_auth import user_api_key_auth 

11from litellm.proxy.spend_tracking.savings import ( 

12 extract_cache_creation_tokens, 

13 extract_cache_read_tokens, 

14 marks_gateway_injection, 

15 prompt_caching_savings_for_request, 

16) 

17from litellm.proxy.spend_tracking.spend_tracking_utils import ( 

18 _query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below 

19) 

20from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY 

21from litellm.types.management_endpoints.prompt_caching_requests import ( 

22 PromptCachingRequest, 

23 PromptCachingRequestCursor, 

24 PromptCachingRequestFilter, 

25 PromptCachingRequestsResponse, 

26) 

27 

28if TYPE_CHECKING: 28 ↛ 29line 28 didn't jump to line 29 because the condition on line 28 was never true

29 from litellm.router import Router 

30 

31router: Final = APIRouter() 

32 

33 

34def _numeric_token_sql(path: str) -> str: 

35 value: Final = f"metadata #> '{{usage_object,{path}}}'" 

36 return ( 

37 f"CASE WHEN jsonb_typeof({value}) = 'number' THEN ({value} #>> '{{}}')::numeric " 

38 f"WHEN {value} = 'true'::jsonb THEN 1 WHEN {value} = 'false'::jsonb THEN 0 END" 

39 ) 

40 

41 

42def _cache_tokens_sql(*paths: str) -> str: 

43 candidates: Final = ", ".join(f"NULLIF(({_numeric_token_sql(path)}), 0)" for path in paths) 

44 return f"TRUNC(COALESCE({candidates}, 0))" 

45 

46 

47_CACHE_READ_SQL: Final = _cache_tokens_sql("cache_read_input_tokens", "prompt_tokens_details,cached_tokens") 

48_CACHE_CREATION_SQL: Final = _cache_tokens_sql( 

49 "cache_creation_input_tokens", 

50 "prompt_tokens_details,cache_write_tokens", 

51 "prompt_tokens_details,cache_creation_tokens", 

52) 

53_GATEWAY_INJECTED_SQL: Final = ( 

54 f"(jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' " 

55 f"AND (metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = '' " 

56 f"OR metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' = model_id))" 

57) 

58_FILTER_SQL: Final = MappingProxyType( 

59 { 

60 "all": f"({_GATEWAY_INJECTED_SQL} OR {_CACHE_READ_SQL} > 0 OR {_CACHE_CREATION_SQL} > 0)", 

61 "injected": _GATEWAY_INJECTED_SQL, 

62 "hits": f"{_CACHE_READ_SQL} > 0", 

63 } 

64) 

65 

66 

67def prompt_caching_requests_sql(filter: PromptCachingRequestFilter) -> str: 

68 return f""" 

69 SELECT request_id, "startTime" AS start_time, "endTime" AS end_time, 

70 model, model_id, custom_llm_provider, spend, 

71 CASE WHEN jsonb_typeof(metadata->'usage_object') = 'object' 

72 THEN metadata->'usage_object' END AS usage_object, 

73 CASE WHEN jsonb_typeof(metadata->'cost_breakdown') = 'object' 

74 THEN metadata->'cost_breakdown' END AS cost_breakdown, 

75 CASE WHEN jsonb_typeof(metadata->'{GATEWAY_INJECTED_CACHE_METADATA_KEY}') = 'string' 

76 THEN metadata->>'{GATEWAY_INJECTED_CACHE_METADATA_KEY}' END AS gateway_marker 

77 FROM "LiteLLM_SpendLogs" 

78 WHERE "startTime" >= ($1::text::timestamptz AT TIME ZONE 'UTC') 

79 AND "startTime" <= ($2::text::timestamptz AT TIME ZONE 'UTC') 

80 AND COALESCE(LOWER(cache_hit), 'false') != 'true' 

81 AND {_FILTER_SQL[filter]} 

82 AND ($4::text::timestamptz IS NULL OR 

83 ("startTime", request_id) < (($4::text::timestamptz AT TIME ZONE 'UTC'), $5::text)) 

84 ORDER BY "startTime" DESC, request_id DESC 

85 LIMIT $3::integer 

86 """ 

87 

88 

89class _PromptCachingRow(BaseModel): 

90 request_id: str 

91 start_time: datetime 

92 end_time: datetime 

93 model: str 

94 model_id: str | None 

95 custom_llm_provider: str | None 

96 spend: float 

97 usage_object: Json[Mapping[str, object]] | Mapping[str, object] | None 

98 cost_breakdown: Json[Mapping[str, object]] | Mapping[str, object] | None 

99 gateway_marker: str | None 

100 

101 

102_REQUEST_ROWS: Final = TypeAdapter(tuple[_PromptCachingRow, ...]) 

103 

104 

105def _request_result(row: _PromptCachingRow, llm_router: "Callable[[], Router | None]") -> PromptCachingRequest: 

106 return PromptCachingRequest( 

107 request_id=row.request_id, 

108 start_time=row.start_time.replace(tzinfo=timezone.utc) if row.start_time.tzinfo is None else row.start_time, 

109 model=row.model, 

110 gateway_injected=marks_gateway_injection( 

111 MappingProxyType({GATEWAY_INJECTED_CACHE_METADATA_KEY: row.gateway_marker}), row.model_id 

112 ), 

113 cache_read_tokens=extract_cache_read_tokens(row.usage_object), 

114 cache_creation_tokens=extract_cache_creation_tokens(row.usage_object), 

115 spend=row.spend, 

116 net_savings=prompt_caching_savings_for_request( 

117 model=row.model, 

118 custom_llm_provider=row.custom_llm_provider, 

119 usage_object=row.usage_object, 

120 model_id=row.model_id, 

121 llm_router=llm_router, 

122 cost_breakdown=row.cost_breakdown, 

123 billed_at=row.end_time, 

124 ), 

125 ) 

126 

127 

128@router.get( 

129 "/cost_optimization/prompt_caching/requests", 

130 tags=["Cost Optimization"], # mutable-ok: FastAPI's route API requires a list 

131 response_model=PromptCachingRequestsResponse, 

132) 

133async def get_prompt_caching_requests( 

134 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], 

135 start_date: datetime, 

136 end_date: datetime, 

137 page_size: Annotated[int, Query(ge=1, le=100)] = 50, 

138 filter: PromptCachingRequestFilter = "all", 

139 cursor_start_time: datetime | None = None, 

140 cursor_request_id: Annotated[str | None, Query(min_length=1)] = None, 

141) -> PromptCachingRequestsResponse: 

142 from litellm.proxy.proxy_server import llm_router, prisma_client 

143 

144 if not user_api_key_has_admin_view(user_api_key_dict): 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 raise HTTPException(status_code=403, detail="Only proxy admin roles can view prompt caching requests") 

146 if (cursor_start_time is None) != (cursor_request_id is None): 

147 raise HTTPException(status_code=400, detail="cursor_start_time and cursor_request_id must be provided together") 

148 if prisma_client is None: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true

149 raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) 

150 start: Final = start_date.replace(tzinfo=timezone.utc) if start_date.tzinfo is None else start_date 

151 end: Final = end_date.replace(tzinfo=timezone.utc) if end_date.tzinfo is None else end_date 

152 if end < start: 

153 raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date") 

154 cursor_time: Final = ( 

155 cursor_start_time.replace(tzinfo=timezone.utc) 

156 if cursor_start_time is not None and cursor_start_time.tzinfo is None 

157 else cursor_start_time 

158 ) 

159 rows: Final = _REQUEST_ROWS.validate_python( 

160 await _query_raw_rows( 

161 prisma_client, 

162 prompt_caching_requests_sql(filter), 

163 start.isoformat(), 

164 end.isoformat(), 

165 page_size + 1, 

166 cursor_time.isoformat() if cursor_time is not None else None, 

167 cursor_request_id, 

168 ) 

169 or () 

170 ) 

171 

172 def current_router() -> "Router | None": 

173 return llm_router 

174 

175 requests: Final = tuple(_request_result(row, current_router) for row in rows[:page_size]) 

176 has_more: Final = len(rows) > page_size 

177 return PromptCachingRequestsResponse( 

178 requests=requests, 

179 page_size=page_size, 

180 has_more=has_more, 

181 next_cursor=PromptCachingRequestCursor(start_time=requests[-1].start_time, request_id=requests[-1].request_id) 

182 if has_more 

183 else None, 

184 )