Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/spend_tracking/spend_event.py: 63%
235 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""Compact, typed success event handed from an inference worker to the collector.
3``build_spend_event`` runs on the inference worker right after ``Logging.async_success_handler``
4has built the ``standard_logging_object`` (so the cost is already known). It validates the success
5callback's ``kwargs`` into the projection ``_PROXY_track_cost_callback`` and
6``DBSpendUpdateWriter.update_database`` actually read: identities and metadata, timings, usage, the
7standard logging payload without its prompt/response bodies, and the tool names. The request
8messages, the raw ``proxy_server_request`` body and the full response travel only when spend logs
9are configured to store prompts and responses. The cache key is the preset key the caching layer
10already computed, never a fresh hash over the request body.
12``spend_event_callback_args`` rebuilds the ``(kwargs, response_obj, start_time, end_time)`` tuple
13the existing cost pipeline consumes, so the sidecar runs the unchanged pipeline against the event.
14"""
16from collections.abc import Mapping, Sequence
17from dataclasses import dataclass
18from datetime import datetime
19from types import MappingProxyType
20from typing import Final, Literal, TypeAlias
22from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
23from typing_extensions import NotRequired, ReadOnly, TypedDict
25import litellm
26from litellm.proxy._types import UserAPIKeyAuth
27from litellm.proxy.db.spend_log_tool_index import response_tool_call_names
28from litellm.types.interactions import InteractionsAPIResponse
29from litellm.types.utils import LiteLLMBatch, Usage
31SPEND_EVENT_VERSION: Final = 1
32CACHE_OFF_KEY: Final = "Cache OFF"
34ObjectMapping: TypeAlias = Mapping[str, object]
36_UNSERIALIZABLE_METADATA_KEYS: Final = frozenset({"user_api_key_auth", "litellm_parent_otel_span"})
37_STANDARD_LOGGING_BODY_KEYS: Final = frozenset({"messages", "response"})
38_STANDARD_LOGGING_DROPPED_KEYS: Final = frozenset({"model_parameters"})
39_NOT_OFFLOADED_RESPONSE_TYPES: Final = (LiteLLMBatch, InteractionsAPIResponse)
42class _LitellmParams(TypedDict, total=False):
43 api_base: ReadOnly[str | None]
44 custom_llm_provider: ReadOnly[str | None]
45 litellm_call_id: ReadOnly[str | None]
46 user_api_key_end_user_id: ReadOnly[str | None]
47 metadata: ReadOnly[ObjectMapping | None]
48 litellm_metadata: ReadOnly[ObjectMapping | None]
49 proxy_server_request: ReadOnly[ObjectMapping | None]
50 preset_cache_key: ReadOnly[str | None]
53class _DynamicParams(TypedDict, total=False):
54 turn_off_message_logging: ReadOnly[bool | None]
57class _RequestBody(TypedDict, total=False):
58 tools: ReadOnly[Sequence[ObjectMapping] | None]
61class _PassthroughPayload(TypedDict, total=False):
62 request_body: ReadOnly[_RequestBody | None]
65class _ToolCallFunction(TypedDict):
66 name: ReadOnly[str]
67 arguments: ReadOnly[str]
70class _ToolCall(TypedDict):
71 id: ReadOnly[str | None]
72 type: ReadOnly[Literal["function"]]
73 function: ReadOnly[_ToolCallFunction]
76class _ToolCallMessage(TypedDict):
77 role: ReadOnly[Literal["assistant"]]
78 content: ReadOnly[None]
79 tool_calls: ReadOnly[Sequence[_ToolCall]]
82class _ToolCallChoice(TypedDict):
83 index: ReadOnly[int]
84 finish_reason: ReadOnly[Literal["tool_calls"]]
85 message: ReadOnly[_ToolCallMessage]
88class CompactResponse(TypedDict, total=False):
89 """What the spend pipeline reads off a response: its id, usage and which tools it called."""
91 id: ReadOnly[object]
92 model: ReadOnly[object]
93 usage: ReadOnly[object]
94 usage_info: ReadOnly[object]
95 status: ReadOnly[object]
96 background: ReadOnly[object]
97 choices: ReadOnly[Sequence[_ToolCallChoice]]
100class _SuccessKwargs(TypedDict, total=False):
101 """The success callback's ``kwargs`` (``Logging.model_call_details``), validated and projected."""
103 litellm_call_id: ReadOnly[str | None]
104 call_type: ReadOnly[str | None]
105 model: ReadOnly[str | None]
106 custom_llm_provider: ReadOnly[str | None]
107 stream: ReadOnly[bool | None]
108 complete_streaming_response: ReadOnly[object]
109 cache_hit: ReadOnly[bool | None]
110 response_cost: ReadOnly[float | None]
111 completion_start_time: ReadOnly[datetime | None]
112 agent_id: ReadOnly[str | None]
113 litellm_trace_id: ReadOnly[str | None]
114 litellm_params: ReadOnly[_LitellmParams]
115 standard_logging_object: ReadOnly[ObjectMapping | None]
116 standard_callback_dynamic_params: ReadOnly[_DynamicParams | None]
117 combined_usage_object: ReadOnly[Usage | None]
118 realtime_tools: ReadOnly[Sequence[object] | None]
119 realtime_tool_calls: ReadOnly[Sequence[object] | None]
120 tools: ReadOnly[Sequence[ObjectMapping] | None]
121 passthrough_logging_payload: ReadOnly[_PassthroughPayload | None]
124class _FunctionToolFunction(TypedDict):
125 name: ReadOnly[str]
128class _FunctionTool(TypedDict):
129 type: ReadOnly[Literal["function"]]
130 function: ReadOnly[_FunctionToolFunction]
133class SpendCallbackKwargs(TypedDict):
134 """The ``kwargs`` handed to ``_PROXY_track_cost_callback`` on the sidecar."""
136 litellm_call_id: ReadOnly[str | None]
137 call_type: ReadOnly[str | None]
138 model: ReadOnly[str | None]
139 custom_llm_provider: ReadOnly[str | None]
140 stream: ReadOnly[bool | None]
141 cache_hit: ReadOnly[bool | None]
142 response_cost: ReadOnly[float | None]
143 completion_start_time: ReadOnly[datetime | None]
144 agent_id: ReadOnly[str | None]
145 litellm_trace_id: ReadOnly[str | None]
146 litellm_params: ReadOnly[_LitellmParams]
147 standard_logging_object: ReadOnly[ObjectMapping | None]
148 standard_callback_dynamic_params: ReadOnly[_DynamicParams | None]
149 combined_usage_object: ReadOnly[Usage | None]
150 realtime_tools: ReadOnly[Sequence[object] | None]
151 realtime_tool_calls: ReadOnly[Sequence[object] | None]
152 tools: ReadOnly[Sequence[_FunctionTool] | None]
153 complete_streaming_response: NotRequired[ReadOnly[CompactResponse | None]]
156_NO_LITELLM_PARAMS: Final[_LitellmParams] = {}
157_SUCCESS_KWARGS: Final = TypeAdapter(_SuccessKwargs)
158_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping)
159_COMPACT_RESPONSE: Final = TypeAdapter(CompactResponse)
162class SpendEvent(BaseModel):
163 model_config = ConfigDict(frozen=True, extra="forbid")
165 version: Literal[1]
166 litellm_call_id: str | None
167 call_type: str | None
168 model: str | None
169 custom_llm_provider: str | None
170 stream: bool | None
171 complete_streaming_response: bool
172 cache_hit: bool | None
173 response_cost: float | None
174 start_time: datetime
175 end_time: datetime
176 completion_start_time: datetime | None
177 agent_id: str | None
178 litellm_trace_id: str | None
179 litellm_params: _LitellmParams
180 standard_logging_object: ObjectMapping | None
181 standard_callback_dynamic_params: _DynamicParams | None
182 response: CompactResponse | None
183 combined_usage: ObjectMapping | None
184 realtime_tools: Sequence[object] | None
185 realtime_tool_calls: Sequence[object] | None
186 request_tool_names: tuple[str, ...]
189@dataclass(frozen=True, slots=True)
190class SpendEventCallbackArgs:
191 kwargs: SpendCallbackKwargs
192 response_obj: CompactResponse | None
193 start_time: datetime
194 end_time: datetime
197@dataclass(frozen=True, slots=True)
198class SpendEventBuildError:
199 reason: str
202@dataclass(frozen=True, slots=True)
203class SpendEventDecodeError:
204 reason: str
207def is_offloadable_success(response_obj: object) -> bool:
208 """Batch retrieves and interaction polls branch on the concrete response class, so they stay in-process."""
209 return not isinstance(response_obj, _NOT_OFFLOADED_RESPONSE_TYPES)
212def _json_fallback(value: object) -> str:
213 return str(value)
216def _mapping_or_none(value: object) -> ObjectMapping | None:
217 try:
218 return _OBJECT_MAPPING.validate_python(value)
219 except ValidationError:
220 return None
223def _drop_keys(mapping: ObjectMapping, keys: frozenset[str]) -> ObjectMapping:
224 return MappingProxyType({key: value for key, value in mapping.items() if key not in keys})
227def _budget_reservation(metadata: ObjectMapping) -> ObjectMapping | None:
228 """The admission-time reservation, wherever the request setup left it, so the sidecar can reconcile it."""
229 direct: Final = _mapping_or_none(metadata.get("user_api_key_budget_reservation"))
230 if direct is not None:
231 return direct
232 auth: Final = metadata.get("user_api_key_auth")
233 if isinstance(auth, UserAPIKeyAuth):
234 return auth.budget_reservation
235 auth_mapping: Final = _mapping_or_none(auth)
236 return _mapping_or_none(auth_mapping.get("budget_reservation")) if auth_mapping is not None else None
239def _metadata_for_event(
240 metadata: ObjectMapping | None, budget_reservation: ObjectMapping | None
241) -> ObjectMapping | None:
242 if metadata is None:
243 return None
244 kept: Final = _drop_keys(metadata, _UNSERIALIZABLE_METADATA_KEYS)
245 if budget_reservation is None:
246 return kept
247 return MappingProxyType({**kept, "user_api_key_budget_reservation": budget_reservation})
250def _litellm_params_for_event(
251 litellm_params: _LitellmParams, cache_key: str | None, store_bodies: bool
252) -> _LitellmParams:
253 metadata: Final = litellm_params.get("metadata")
254 litellm_metadata: Final = litellm_params.get("litellm_metadata")
255 budget_reservation: Final = next(
256 (
257 reservation
258 for source in (litellm_metadata, metadata)
259 if source is not None and (reservation := _budget_reservation(source)) is not None
260 ),
261 None,
262 )
263 projected: Final[_LitellmParams] = {
264 "api_base": litellm_params.get("api_base"),
265 "custom_llm_provider": litellm_params.get("custom_llm_provider"),
266 "litellm_call_id": litellm_params.get("litellm_call_id"),
267 "user_api_key_end_user_id": litellm_params.get("user_api_key_end_user_id"),
268 "metadata": _metadata_for_event(metadata, budget_reservation),
269 "litellm_metadata": _metadata_for_event(litellm_metadata, budget_reservation),
270 "proxy_server_request": litellm_params.get("proxy_server_request") if store_bodies else None,
271 "preset_cache_key": cache_key,
272 }
273 return projected
276def _standard_logging_for_event(sl_object: ObjectMapping | None, store_bodies: bool) -> ObjectMapping | None:
277 if sl_object is None:
278 return None
279 dropped: Final = (
280 _STANDARD_LOGGING_DROPPED_KEYS if store_bodies else _STANDARD_LOGGING_DROPPED_KEYS | _STANDARD_LOGGING_BODY_KEYS
281 )
282 return _drop_keys(sl_object, dropped)
285def _tool_call(name: str) -> _ToolCall:
286 tool_call: Final[_ToolCall] = {"id": None, "type": "function", "function": {"name": name, "arguments": "{}"}}
287 return tool_call
290def _tool_call_choice(names: Sequence[str]) -> _ToolCallChoice:
291 choice: Final[_ToolCallChoice] = {
292 "index": 0,
293 "finish_reason": "tool_calls",
294 "message": {"role": "assistant", "content": None, "tool_calls": tuple(_tool_call(name) for name in names)},
295 }
296 return choice
299def _compact_response(response_obj: object) -> CompactResponse | None:
300 """Usage, identity and tool calls of the response, in chat-completions shape, without the content."""
301 dumped: Final = response_obj.model_dump() if isinstance(response_obj, BaseModel) else _mapping_or_none(response_obj)
302 if dumped is None:
303 return None
304 scalars: Final = _COMPACT_RESPONSE.validate_python(_drop_keys(dumped, frozenset({"choices"})))
305 tool_call_names: Final = response_tool_call_names(response_obj)
306 if not tool_call_names:
307 return scalars
308 with_tool_calls: Final[CompactResponse] = {**scalars, "choices": (_tool_call_choice(tool_call_names),)}
309 return with_tool_calls
312def _tool_name(tool: ObjectMapping) -> str | None:
313 """Chat tools nest the name under ``function``; Anthropic and Responses API tools keep it at the top."""
314 function: Final = _mapping_or_none(tool.get("function"))
315 name: Final = function.get("name") if function is not None else tool.get("name")
316 return name.strip() if isinstance(name, str) and name.strip() else None
319def _request_tool_names(kwargs: _SuccessKwargs) -> tuple[str, ...]:
320 passthrough: Final = kwargs.get("passthrough_logging_payload")
321 request_body: Final = passthrough.get("request_body") if passthrough is not None else None
322 passthrough_tools: Final = request_body.get("tools") if request_body is not None else None
323 return tuple(
324 name
325 for source in (kwargs.get("tools"), passthrough_tools)
326 if source is not None
327 for tool in source
328 if (name := _tool_name(tool)) is not None
329 )
332def preset_spend_log_cache_key(litellm_params: _LitellmParams) -> str | None:
333 """The key the caching layer already stored in ``litellm_params``, or ``Cache OFF``; never hashes the body."""
334 if litellm.cache is None:
335 return CACHE_OFF_KEY
336 return litellm_params.get("preset_cache_key")
339def _function_tool(name: str) -> _FunctionTool:
340 tool: Final[_FunctionTool] = {"type": "function", "function": {"name": name}}
341 return tool
344def build_spend_event(
345 raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool
346) -> bytes | SpendEventBuildError:
347 """Validate the success callback's kwargs and serialize the event once, as a single JSON line."""
348 try:
349 kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs)
350 except ValidationError as error:
351 return SpendEventBuildError(reason=str(error))
352 litellm_params: Final = kwargs.get("litellm_params", _NO_LITELLM_PARAMS)
353 sl_object: Final = kwargs.get("standard_logging_object")
354 cache_key: Final = preset_spend_log_cache_key(litellm_params)
355 response_cost: Final = sl_object.get("response_cost") if sl_object is not None else kwargs.get("response_cost")
356 combined_usage: Final = kwargs.get("combined_usage_object")
357 event: Final = SpendEvent(
358 version=SPEND_EVENT_VERSION,
359 litellm_call_id=kwargs.get("litellm_call_id"),
360 call_type=kwargs.get("call_type"),
361 model=kwargs.get("model"),
362 custom_llm_provider=kwargs.get("custom_llm_provider"),
363 stream=kwargs.get("stream"),
364 complete_streaming_response="complete_streaming_response" in kwargs,
365 cache_hit=kwargs.get("cache_hit"),
366 response_cost=response_cost if isinstance(response_cost, (int, float)) else None,
367 start_time=start_time,
368 end_time=end_time,
369 completion_start_time=kwargs.get("completion_start_time"),
370 agent_id=kwargs.get("agent_id"),
371 litellm_trace_id=kwargs.get("litellm_trace_id"),
372 litellm_params=_litellm_params_for_event(litellm_params, cache_key, store_bodies),
373 standard_logging_object=_standard_logging_for_event(sl_object, store_bodies),
374 standard_callback_dynamic_params=kwargs.get("standard_callback_dynamic_params"),
375 response=_compact_response(response_obj),
376 combined_usage=combined_usage.model_dump() if combined_usage is not None else None,
377 realtime_tools=kwargs.get("realtime_tools"),
378 realtime_tool_calls=kwargs.get("realtime_tool_calls"),
379 request_tool_names=_request_tool_names(kwargs),
380 )
381 return event.model_dump_json(fallback=_json_fallback).encode() + b"\n"
384def decode_spend_event(line: bytes) -> SpendEvent | SpendEventDecodeError:
385 try:
386 return SpendEvent.model_validate_json(line)
387 except ValidationError as error:
388 return SpendEventDecodeError(reason=str(error))
391def spend_event_callback_args(event: SpendEvent) -> SpendEventCallbackArgs:
392 """The ``(kwargs, response_obj, start_time, end_time)`` the in-process cost callback receives."""
393 tools: Final = tuple(_function_tool(name) for name in event.request_tool_names)
394 kwargs: Final[SpendCallbackKwargs] = {
395 "litellm_call_id": event.litellm_call_id,
396 "call_type": event.call_type,
397 "model": event.model,
398 "custom_llm_provider": event.custom_llm_provider,
399 "stream": event.stream,
400 "cache_hit": event.cache_hit,
401 "response_cost": event.response_cost,
402 "completion_start_time": event.completion_start_time,
403 "agent_id": event.agent_id,
404 "litellm_trace_id": event.litellm_trace_id,
405 "litellm_params": event.litellm_params,
406 "standard_logging_object": event.standard_logging_object,
407 "standard_callback_dynamic_params": event.standard_callback_dynamic_params,
408 "combined_usage_object": Usage.model_validate(event.combined_usage)
409 if event.combined_usage is not None
410 else None,
411 "realtime_tools": event.realtime_tools,
412 "realtime_tool_calls": event.realtime_tool_calls,
413 "tools": tools or None,
414 }
415 if not event.complete_streaming_response:
416 return SpendEventCallbackArgs(kwargs, event.response, event.start_time, event.end_time)
417 streaming_kwargs: Final[SpendCallbackKwargs] = {**kwargs, "complete_streaming_response": event.response}
418 return SpendEventCallbackArgs(streaming_kwargs, event.response, event.start_time, event.end_time)