Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/middleware/billable_request_metrics_middleware.py: 91%
127 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Counts HTTP requests to LLM inference, MCP, and A2A endpoints.
4Feeds two independent sinks off one classification:
6- ``GatewayRequestSink`` receives every classified request with its status and
7 is the source of truth for SGR (successful gateway requests) on the admin UI.
8 Not license-gated (see litellm.proxy.db.gateway_request_tracking). It is not
9 told which deployment served the request: it persists its counts, so every
10 dimension it takes has to be one the proxy chooses.
11- ``BillingRecorder`` receives 2xx requests only and exports them for
12 enterprise metering (see litellm.proxy.enterprise_billing.billing_metrics).
14Both are injected. When neither is present the middleware is a transparent
15pass-through.
16"""
18import re
19import threading
20from collections.abc import Callable, Sequence
21from enum import Enum
22from typing import Final, Protocol, runtime_checkable
24from starlette.types import ASGIApp, Message, Receive, Scope, Send
26from litellm._logging import verbose_proxy_logger
27from litellm.proxy._types import LiteLLMRoutes
30class BillableCategory(str, Enum):
31 LLM = "llm"
32 MCP = "mcp"
33 A2A = "a2a"
36@runtime_checkable
37class BillingRecorder(Protocol):
38 def record(self, *, category: BillableCategory, route: str, status_code: int, model_id: str | None) -> None: ... 38 ↛ exitline 38 didn't return from function 'record' because
41@runtime_checkable
42class GatewayRequestSink(Protocol):
43 """
44 Records every classified request, 2xx or not, for the SGR dashboard.
46 Distinct from BillingRecorder on three counts: this is not license-gated,
47 it is not restricted to 2xx, and it takes no model id. The deployment that
48 served a request is deliberately not part of what it records, because the
49 dashboard aggregates by route and a per-deployment dimension would only
50 multiply the rows it has to sum back together.
51 """
53 def record(self, *, category: BillableCategory, route: str, status_code: int) -> None: ... 53 ↛ exitline 53 didn't return from function 'record' because
56_MODEL_ID_HEADER: Final = b"x-litellm-model-id"
58# Ordered: a longer suffix that shares an ending with a shorter one must come
59# first, e.g. "/chat/completions" before "/completions". This is the POST
60# inference surface that writes a SpendLogs row on success, so the exported
61# count lines up with the admin UI usage page for inference traffic. Billing
62# is a deliberate lower bound on SpendLogs rows: management writes that also
63# log (batch/file/fine-tuning creation, interaction cancel) and non-POST calls
64# that log (passthrough reads) never bill, so drift only ever undercounts.
65_LLM_ROUTE_SUFFIXES: Final[tuple[str, ...]] = (
66 "/chat/completions",
67 "/completions",
68 "/embeddings",
69 "/responses",
70 "/rerank",
71 "/moderations",
72 "/images/generations",
73 "/images/edits",
74 "/images/variations",
75 "/audio/transcriptions",
76 "/audio/translations",
77 "/audio/speech",
78 "/videos", # create; GET list is excluded by the POST gate
79 "/remix", # /v1/videos/{id}/remix
80 "/ocr",
81 "/search", # /v1/search and /v1/vector_stores/{id}/search
82 "/rag/query",
83 "/rag/ingest",
84 ":generateContent", # Gemini-native /v1beta/models/{model}:generateContent
85 ":streamGenerateContent",
86)
88# Exact paths only: a suffix match would also catch non-inference resources that
89# share the ending, e.g. the OpenAI Assistants route /v1/threads/{id}/messages
90# writes no SpendLogs row and must not bill, unlike Anthropic /v1/messages.
91_LLM_ROUTE_EXACT: Final[tuple[str, ...]] = (
92 "/v1/messages",
93 "/interactions", # Google Interactions create; /{id} reads and /cancel do not match
94 "/v1beta/interactions",
95 "/comprehendmedical", # AWS-SDK-shaped passthrough: the operation rides in the X-Amz-Target header
96 "/transcribe",
97)
99# Provider passthrough prefixes (e.g. /bedrock/..., /vertex-ai/...) carry real
100# inference calls that write SpendLogs rows, so they bill. Anchored to the
101# routes enum so new providers are picked up without touching this module.
102# /langfuse forwards observability traffic, not inference: it writes no
103# SpendLogs row and must not bill.
104_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse"})
105_PASSTHROUGH_PREFIXES: Final[tuple[str, ...]] = tuple(
106 prefix
107 for prefix in LiteLLMRoutes.mapped_pass_through_routes.value
108 if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES
109)
112def _classify_llm_route(path: str) -> str | None:
113 exact_match: Final = next((route for route in _LLM_ROUTE_EXACT if path == route), None)
114 if exact_match is not None:
115 return exact_match
116 suffix_match = next((suffix for suffix in _LLM_ROUTE_SUFFIXES if path == suffix or path.endswith(suffix)), None)
117 if suffix_match is not None:
118 return suffix_match
119 # Deep passthrough paths only: the bare prefix itself is not an inference call.
120 return next((prefix for prefix in _PASSTHROUGH_PREFIXES if path.startswith(f"{prefix}/")), None)
123_MCP_MANAGEMENT_PREFIX: Final = "/v1/mcp"
124_MCP_DYNAMIC_TRANSPORT: Final = re.compile(r"/(?:toolset/)?[^/]+/mcp")
125# The REST wrapper's tool-call endpoint executes a tool and fires the same MCP
126# spend logging as the /mcp transport; its list/test siblings do not bill.
127_MCP_REST_TOOL_CALL: Final = "/mcp-rest/tools/call"
129_A2A_INVOKE_SUFFIX: Final = "/message/send"
130_A2A_TRANSPORT_PREFIXES: Final[tuple[str, ...]] = ("/v1/a2a/", "/a2a/")
131# Bare POST /a2a/{agent_id} carries the JSON-RPC method in the body, not the
132# path. Only message/send and message/stream write a SpendLogs row there; the
133# task RPCs (tasks/get, tasks/cancel, tasks/pushNotificationConfig/*, ...) are
134# forwarded upstream and write none. A path-only classifier cannot separate
135# them, so the bare route does not bill: counting a task RPC would overcount,
136# while missing a bare-path message/send only undercounts, and undercounting is
137# the sole direction this metric is allowed to drift. The /mcp transport is
138# method-agnostic by contrast because its list path logs a SpendLogs row too.
141def _classify_mcp_route(path: str) -> str | None:
142 if path == _MCP_MANAGEMENT_PREFIX or path.startswith(f"{_MCP_MANAGEMENT_PREFIX}/"):
143 return None
144 if path == "/mcp" or path.startswith("/mcp/"):
145 return "/mcp"
146 if path == _MCP_REST_TOOL_CALL:
147 return "/mcp"
148 if _MCP_DYNAMIC_TRANSPORT.fullmatch(path) is not None:
149 return "/mcp"
150 return None
153def _classify_a2a_route(path: str) -> str | None:
154 if path.endswith(_A2A_INVOKE_SUFFIX) and any(path.startswith(prefix) for prefix in _A2A_TRANSPORT_PREFIXES):
155 return "/a2a"
156 return None
159def classify_billable_request(path: str, method: str = "POST") -> tuple[BillableCategory, str] | None:
160 """Map a request path to its (category, normalized route), or None if not billable."""
161 normalized: Final = path.rstrip("/") or "/"
163 mcp_route: Final = _classify_mcp_route(normalized)
164 if mcp_route is not None:
165 return (BillableCategory.MCP, mcp_route)
167 a2a_route: Final = _classify_a2a_route(normalized)
168 if a2a_route is not None:
169 return (BillableCategory.A2A, a2a_route)
171 # POST-only is a conservative gate: non-POST calls can still write a
172 # SpendLogs row (passthrough reads, resource GETs) but must not bill, so
173 # any classifier-vs-dashboard mismatch is an undercount, never an overcount.
174 if method.upper() != "POST":
175 return None
177 llm_route: Final = _classify_llm_route(normalized)
178 if llm_route is not None:
179 return (BillableCategory.LLM, llm_route)
180 return None
183def _extract_model_id(headers: Sequence[tuple[bytes, bytes]]) -> str | None:
184 return next(
185 (value.decode("latin-1") for name, value in headers if name.lower() == _MODEL_ID_HEADER and value),
186 None,
187 )
190class BillableRequestMetricsMiddleware:
191 """
192 Pure ASGI middleware that classifies each request once and fans the result
193 out to the SGR sink (any status) and the billing recorder (2xx only).
194 Modeled on InFlightRequestsMiddleware: it wraps `send`, reads the final
195 status and the x-litellm-model-id header off the `http.response.start`
196 message, and never blocks or fails the request path.
197 """
199 def __init__(
200 self,
201 app: ASGIApp,
202 recorder: BillingRecorder | None = None,
203 recorder_factory: Callable[[], BillingRecorder | None] | None = None,
204 sink: GatewayRequestSink | None = None,
205 sink_factory: Callable[[], GatewayRequestSink | None] | None = None,
206 ) -> None:
207 self.app = app
208 self.recorder = recorder
209 # The factory defers recorder construction to the first request, AFTER the
210 # startup event has loaded the YAML config's environment_variables (license
211 # and cert env vars). Building at import time captured recorder=None for
212 # deployments configured that way. Resolved exactly once; the result
213 # (including None) is cached.
214 self._recorder_factory = recorder_factory
215 self._resolved = recorder_factory is None
216 self._resolve_lock = threading.Lock()
217 # Resolved on the same schedule and for the same reason: the DB is not
218 # connected at import time, so the sink cannot be built there either.
219 self.sink = sink
220 self._sink_factory = sink_factory
221 self._sink_resolved = sink_factory is None
222 self._sink_resolve_lock = threading.Lock()
224 def _resolve_recorder(self) -> BillingRecorder | None:
225 if self._resolved:
226 return self.recorder
227 # The lock keeps concurrent first requests from each building their own
228 # MeterProvider (and leaking its background exporter thread).
229 with self._resolve_lock:
230 if not self._resolved: 230 ↛ 234line 230 didn't jump to line 234
231 factory: Final = self._recorder_factory
232 self.recorder = factory() if factory is not None else self.recorder
233 self._resolved = True
234 return self.recorder
236 def _resolve_sink(self) -> GatewayRequestSink | None:
237 if self._sink_resolved:
238 return self.sink
239 with self._sink_resolve_lock:
240 if not self._sink_resolved: 240 ↛ 244line 240 didn't jump to line 244
241 factory: Final = self._sink_factory
242 self.sink = factory() if factory is not None else self.sink
243 self._sink_resolved = True
244 return self.sink
246 async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
247 if scope["type"] != "http":
248 await self.app(scope, receive, send)
249 return
251 recorder: Final = self._resolve_recorder()
252 sink: Final = self._resolve_sink()
253 if recorder is None and sink is None: 253 ↛ 254line 253 didn't jump to line 254 because the condition on line 253 was never true
254 await self.app(scope, receive, send)
255 return
257 classification: Final = classify_billable_request(scope.get("path", ""), scope.get("method", "POST"))
258 if classification is None:
259 await self.app(scope, receive, send)
260 return
262 category, route = classification
263 status_code = 0
264 model_id: str | None = None
266 async def send_wrapper(message: Message) -> None:
267 nonlocal status_code, model_id
268 if message["type"] == "http.response.start":
269 status_code = message["status"]
270 model_id = _extract_model_id(message.get("headers", []))
271 await send(message)
273 await self.app(scope, receive, send_wrapper)
275 if sink is not None: 275 ↛ 281line 275 didn't jump to line 281 because the condition on line 275 was always true
276 try:
277 sink.record(category=category, route=route, status_code=status_code)
278 except Exception: # noqa: BLE001 -- metering must never fail a request that was already served
279 verbose_proxy_logger.warning("gateway request metering failed for %s", route, exc_info=True)
281 if recorder is not None and 200 <= status_code < 300: 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 try:
283 recorder.record(category=category, route=route, status_code=status_code, model_id=model_id)
284 except Exception: # noqa: BLE001 -- metering must never fail a request that was already served
285 verbose_proxy_logger.warning("billable request metering failed for %s", route, exc_info=True)