Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/pass_through_endpoints/streaming_handler.py: 20%
177 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import traceback
2from collections.abc import Coroutine, Sequence
3from dataclasses import dataclass
4from datetime import datetime, timezone
5from typing import Final, Protocol
7import httpx
9import litellm
10from litellm._logging import verbose_proxy_logger
11from litellm.litellm_core_utils.asyncify import asyncify
12from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
13from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
14from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
15from litellm.proxy._types import PassThroughEndpointLoggingResultValues
16from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
17from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
18from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
19from litellm.types.utils import StandardPassThroughResponseObject, Usage
21from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
22 AnthropicPassthroughLoggingHandler,
23)
24from .llm_provider_handlers.gemini_passthrough_logging_handler import (
25 GeminiPassthroughLoggingHandler,
26)
27from .llm_provider_handlers.openai_passthrough_logging_handler import (
28 OpenAIPassthroughLoggingHandler,
29)
30from .llm_provider_handlers.tinyfish_passthrough_logging_handler import (
31 TinyFishPassthroughLoggingHandler,
32 run_id_from_sse_frames,
33 sse_poller_spawned,
34)
35from .llm_provider_handlers.vertex_passthrough_logging_handler import (
36 VertexPassthroughLoggingHandler,
37)
38from .success_handler import PassThroughEndpointLogging
41class RouteStreamingLogging(Protocol):
42 def __call__( 42 ↛ exitline 42 didn't return from function '__call__' because
43 self,
44 *,
45 litellm_logging_obj: LiteLLMLoggingObj,
46 passthrough_success_handler_obj: PassThroughEndpointLogging,
47 url_route: str,
48 request_body: dict,
49 endpoint_type: EndpointType,
50 start_time: datetime,
51 raw_bytes: list[bytes],
52 end_time: datetime,
53 ) -> Coroutine[None, None, None]: ...
56@dataclass(frozen=True, slots=True)
57class PassThroughStreamContext:
58 passthrough_success_handler_obj: PassThroughEndpointLogging
59 url_route: str
60 start_time: datetime
63class PassThroughStreamingHandler:
64 @staticmethod
65 def _stamp_first_chunk_if_needed(litellm_logging_obj: LiteLLMLoggingObj) -> None:
66 if litellm_logging_obj.completion_start_time is None:
67 litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now())
69 @staticmethod
70 async def schedule_stream_failure_logging(
71 litellm_logging_obj: LiteLLMLoggingObj,
72 endpoint_type: EndpointType,
73 request_body: dict[str, object],
74 raw_bytes: Sequence[bytes],
75 exception: Exception,
76 stream_context: PassThroughStreamContext | None = None,
77 ) -> None:
78 # the tinyfish poller writes the one authoritative row; a failure row here would collide on its request_id
79 if endpoint_type == EndpointType.TINYFISH and sse_poller_spawned(litellm_logging_obj):
80 return
81 await asyncify(PassThroughStreamingHandler._record_partial_usage_for_failure)(
82 litellm_logging_obj=litellm_logging_obj,
83 endpoint_type=endpoint_type,
84 request_body=request_body,
85 raw_bytes=raw_bytes,
86 stream_context=stream_context,
87 )
88 try:
89 GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
90 async_coroutine=litellm_logging_obj.dispatch_failure_handlers(
91 exception, traceback.format_exc(), prefer_async_handlers=True
92 )
93 )
94 except Exception as e:
95 verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e)
97 @staticmethod
98 def _record_partial_usage_for_failure(
99 litellm_logging_obj: LiteLLMLoggingObj,
100 endpoint_type: EndpointType,
101 request_body: dict[str, object],
102 raw_bytes: Sequence[bytes],
103 stream_context: PassThroughStreamContext | None,
104 ) -> None:
105 if endpoint_type == EndpointType.ANTHROPIC:
106 AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure(
107 litellm_logging_obj=litellm_logging_obj, request_body=request_body, all_chunks=raw_bytes
108 )
109 return
110 if stream_context is None or not raw_bytes:
111 return
112 try:
113 partial_response, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result(
114 litellm_logging_obj=litellm_logging_obj,
115 passthrough_success_handler_obj=stream_context.passthrough_success_handler_obj,
116 url_route=stream_context.url_route,
117 request_body=request_body,
118 endpoint_type=endpoint_type,
119 start_time=stream_context.start_time,
120 raw_bytes=raw_bytes,
121 end_time=datetime.now(timezone.utc),
122 model=None,
123 )
124 except Exception as e:
125 verbose_proxy_logger.warning(
126 "Could not recover the partial usage of a failed %s pass-through stream: %s", endpoint_type.value, e
127 )
128 return
129 usage: Final = getattr(partial_response, "usage", None)
130 if not isinstance(usage, Usage):
131 return
132 response_cost: Final = kwargs.get("response_cost")
133 litellm_logging_obj.record_partial_usage_for_failure(
134 usage=usage,
135 response_cost=float(response_cost) if isinstance(response_cost, (int, float)) else 0.0,
136 )
138 @staticmethod
139 async def chunk_processor(
140 response: httpx.Response,
141 request_body: dict | None,
142 litellm_logging_obj: LiteLLMLoggingObj,
143 endpoint_type: EndpointType,
144 start_time: datetime,
145 passthrough_success_handler_obj: PassThroughEndpointLogging,
146 url_route: str,
147 route_streaming_logging: RouteStreamingLogging | None = None,
148 ):
149 resolved_route_streaming_logging: Final[RouteStreamingLogging] = (
150 route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
151 )
152 raw_bytes: Final[list[bytes]] = []
153 resolved_request_body: Final[dict[str, object]] = request_body or {}
155 def _build_logging_coroutine() -> Coroutine[None, None, None]:
156 return resolved_route_streaming_logging(
157 litellm_logging_obj=litellm_logging_obj,
158 passthrough_success_handler_obj=passthrough_success_handler_obj,
159 url_route=url_route,
160 request_body=resolved_request_body,
161 endpoint_type=endpoint_type,
162 start_time=start_time,
163 raw_bytes=raw_bytes,
164 end_time=datetime.now(),
165 )
167 logging_scheduled = False
168 model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
169 request_body=request_body,
170 url_route=url_route,
171 endpoint_type=endpoint_type,
172 litellm_logging_obj=litellm_logging_obj,
173 )
175 # Resolve once per stream rather than re-reading the global +
176 # re-branching on every chunk. ``include_cost_in_streaming_usage`` is
177 # set at config load and stable for the process, matching how the
178 # proxy-level streaming fast path resolves it.
179 cost_injection_active: Final = (
180 bool(getattr(litellm, "include_cost_in_streaming_usage", False))
181 and bool(model_name)
182 and (
183 endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI)
184 or (
185 endpoint_type == EndpointType.VERTEX_AI
186 and ("streamRawPredict" in url_route or "rawPredict" in url_route)
187 )
188 )
189 )
190 # TinyFish SSE bills via a detached poller spawned on the first run_id frame, so disconnects can't lose the charge
191 tinyfish_scan_active = endpoint_type == EndpointType.TINYFISH # rebind-ok: scan stops once the poller spawns
192 tinyfish_pending = b"" # rebind-ok: SSE frame reassembly buffer across transport chunks
193 try:
194 if not cost_injection_active:
195 # Hot path: just buffer for end-of-stream logging and forward.
196 async for chunk in response.aiter_bytes():
197 raw_bytes.append(chunk)
198 PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj)
199 if tinyfish_scan_active:
200 complete_frames, tinyfish_pending = split_complete_sse_frames(tinyfish_pending + chunk)
201 run_id = run_id_from_sse_frames(complete_frames) if b"run_id" in complete_frames else None
202 if run_id:
203 TinyFishPassthroughLoggingHandler.start_sse_run_billing(
204 run_id=run_id,
205 litellm_logging_obj=litellm_logging_obj,
206 start_time=start_time,
207 )
208 tinyfish_scan_active = False
209 yield chunk
210 else:
211 # ``cost_injection_active`` already requires ``model_name`` to
212 # be truthy; pin to a typed local so mypy narrows ``Optional[str]``
213 # -> ``str`` for the per-chunk call site.
214 assert model_name is not None
215 resolved_model_name: Final[str] = model_name
216 pending = b""
217 async for chunk in response.aiter_bytes():
218 raw_bytes.append(chunk)
219 PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj)
220 complete_frames, pending = split_complete_sse_frames(pending + chunk)
221 if complete_frames:
222 yield ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection(
223 complete_frames, resolved_model_name, litellm_logging_obj
224 )
225 if pending:
226 yield pending
227 # Stream completed cleanly. When the proxy armed deferred
228 # dispatch (post-call guardrails active), park the logging
229 # coroutine on logging_obj instead of enqueueing now, so
230 # ProxyLogging._fire_deferred_stream_logging fires it after
231 # guardrail end-of-stream blocks populate guardrail_information.
232 # Disconnect paths skip this and fall through to the immediate
233 # enqueue in ``finally`` to keep partial billing (LIT-2642);
234 # upstream exceptions log a failure instead (LIT-3798).
235 if (
236 getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None
237 and raw_bytes
238 and response.status_code < 400
239 ):
240 logging_scheduled = True
241 bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
242 litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
243 except Exception as e:
244 verbose_proxy_logger.error("Error in chunk_processor: %s", e)
245 if response.status_code < 400:
246 logging_scheduled = True
247 await PassThroughStreamingHandler.schedule_stream_failure_logging(
248 litellm_logging_obj=litellm_logging_obj,
249 endpoint_type=endpoint_type,
250 request_body=resolved_request_body,
251 raw_bytes=raw_bytes,
252 exception=e,
253 stream_context=PassThroughStreamContext(
254 passthrough_success_handler_obj=passthrough_success_handler_obj,
255 url_route=url_route,
256 start_time=start_time,
257 ),
258 )
259 raise
260 finally:
261 # GeneratorExit (raised on client disconnect) is not caught by
262 # `except Exception`; the finally block ensures partial usage
263 # still gets logged for spend tracking. See LIT-2642.
264 # Upstream 4xx/5xx responses are already logged as a failure by
265 # the caller before this generator starts (see
266 # _log_passthrough_upstream_failure); logging them again here as
267 # a success would double-log the same request.
268 if not logging_scheduled and raw_bytes and response.status_code < 400:
269 logging_scheduled = True
270 try:
271 GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine())
272 except Exception as e:
273 verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)
274 else:
275 bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
277 @staticmethod
278 async def _route_streaming_logging_to_handler(
279 litellm_logging_obj: LiteLLMLoggingObj,
280 passthrough_success_handler_obj: PassThroughEndpointLogging,
281 url_route: str,
282 request_body: dict,
283 endpoint_type: EndpointType,
284 start_time: datetime,
285 raw_bytes: Sequence[bytes],
286 end_time: datetime,
287 model: str | None = None,
288 ):
289 """
290 Route the logging for the collected chunks to the appropriate handler
292 Supported endpoint types:
293 - Anthropic
294 - Vertex AI
295 - OpenAI
296 """
297 from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
298 _is_message_stop_chunk, # pyright: ignore[reportPrivateUsage] # both native stream paths share terminal-event detection
299 _is_provider_error_chunk, # pyright: ignore[reportPrivateUsage] # provider errors must not become cache evidence
300 )
302 # Transport reads can split event names and JSON payloads. Recognize terminal
303 # events only after the shared SSE framer has reassembled the collected bytes.
304 complete_frames, incomplete_tail = split_complete_sse_frames(
305 b"".join(raw_bytes) if endpoint_type == EndpointType.ANTHROPIC else b""
306 )
307 litellm_logging_obj.model_call_details[ # rebind-ok: stamp evidence on the per-request state read by callbacks
308 "prompt_cache_response_complete"
309 ] = (
310 endpoint_type == EndpointType.ANTHROPIC
311 and not incomplete_tail.strip()
312 and _is_message_stop_chunk(complete_frames)
313 and not _is_provider_error_chunk(complete_frames)
314 )
315 try:
316 # TinyFish billing is owned by the detached poller; the $0 fallback below is only for streams with no run_id
317 if endpoint_type == EndpointType.TINYFISH:
318 if sse_poller_spawned(litellm_logging_obj):
319 return
320 late_run_id: Final = run_id_from_sse_frames(b"".join(raw_bytes))
321 if late_run_id:
322 # the run_id arrived in an unterminated frame; poll to terminal instead of mispricing a RUNNING run
323 TinyFishPassthroughLoggingHandler.start_sse_run_billing(
324 run_id=late_run_id,
325 litellm_logging_obj=litellm_logging_obj,
326 start_time=start_time,
327 )
328 return
329 tinyfish_payload: Final = (
330 await TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks(
331 litellm_logging_obj=litellm_logging_obj,
332 url_route=url_route,
333 start_time=start_time,
334 all_chunks=PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes),
335 end_time=end_time,
336 )
337 )
338 await litellm_logging_obj.dispatch_success_handlers(
339 result=tinyfish_payload["result"],
340 start_time=start_time,
341 end_time=end_time,
342 cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True,
343 prefer_async_handlers=True,
344 **tinyfish_payload["kwargs"],
345 )
346 return
347 (
348 standard_logging_response_object,
349 kwargs,
350 ) = await asyncify(PassThroughStreamingHandler._build_passthrough_logging_result)(
351 litellm_logging_obj=litellm_logging_obj,
352 passthrough_success_handler_obj=passthrough_success_handler_obj,
353 url_route=url_route,
354 request_body=request_body,
355 endpoint_type=endpoint_type,
356 start_time=start_time,
357 raw_bytes=raw_bytes,
358 end_time=end_time,
359 model=model,
360 )
361 # Always reached from an async context (anthropic_messages,
362 # google_genai, and proxy pass-through stream tasks). prefer_async_handlers
363 # keeps async-only loggers running even when call_type isn't pass_through
364 # and litellm_params lacks an async flag (e.g. aanthropic_messages).
365 await litellm_logging_obj.dispatch_success_handlers(
366 result=standard_logging_response_object,
367 start_time=start_time,
368 end_time=end_time,
369 cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True,
370 prefer_async_handlers=True,
371 **kwargs,
372 )
373 except Exception as e:
374 verbose_proxy_logger.error("Error in _route_streaming_logging_to_handler: %s", e)
376 @staticmethod
377 def _build_passthrough_logging_result(
378 litellm_logging_obj: LiteLLMLoggingObj,
379 passthrough_success_handler_obj: PassThroughEndpointLogging,
380 url_route: str,
381 request_body: dict,
382 endpoint_type: EndpointType,
383 start_time: datetime,
384 raw_bytes: Sequence[bytes],
385 end_time: datetime,
386 model: str | None,
387 ) -> tuple[PassThroughEndpointLoggingResultValues, dict]:
388 """
389 Synchronous, CPU-bound reconstruction of the standard logging payload
390 from collected raw SSE bytes. Extracted from
391 _route_streaming_logging_to_handler so the per-endpoint dispatch can
392 be unit-tested in isolation. The async callers run it in a worker
393 thread so the token counts inside stay off the event loop.
394 """
395 all_chunks: Final = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes)
396 standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None
397 kwargs: dict = {}
398 if endpoint_type == EndpointType.ANTHROPIC:
399 anthropic_passthrough_logging_handler_result: Final = (
400 AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
401 litellm_logging_obj=litellm_logging_obj,
402 passthrough_success_handler_obj=passthrough_success_handler_obj,
403 url_route=url_route,
404 request_body=request_body,
405 endpoint_type=endpoint_type,
406 start_time=start_time,
407 all_chunks=all_chunks,
408 end_time=end_time,
409 )
410 )
411 standard_logging_response_object = anthropic_passthrough_logging_handler_result["result"]
412 kwargs = anthropic_passthrough_logging_handler_result["kwargs"]
413 elif endpoint_type == EndpointType.VERTEX_AI:
414 vertex_passthrough_logging_handler_result: Final = (
415 VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks(
416 litellm_logging_obj=litellm_logging_obj,
417 passthrough_success_handler_obj=passthrough_success_handler_obj,
418 url_route=url_route,
419 request_body=request_body,
420 endpoint_type=endpoint_type,
421 start_time=start_time,
422 all_chunks=all_chunks,
423 end_time=end_time,
424 model=model,
425 )
426 )
427 standard_logging_response_object = vertex_passthrough_logging_handler_result["result"]
428 kwargs = vertex_passthrough_logging_handler_result["kwargs"]
429 elif endpoint_type == EndpointType.GEMINI:
430 gemini_passthrough_logging_handler_result: Final = (
431 GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( # pyright: ignore[reportPrivateUsage] # mirrors sibling handler dispatch
432 litellm_logging_obj=litellm_logging_obj,
433 passthrough_success_handler_obj=passthrough_success_handler_obj,
434 url_route=url_route,
435 request_body=request_body,
436 endpoint_type=endpoint_type,
437 start_time=start_time,
438 all_chunks=all_chunks,
439 end_time=end_time,
440 model=model,
441 )
442 )
443 standard_logging_response_object = ( # rebind-ok: branch bind in shared if/elif dispatch
444 gemini_passthrough_logging_handler_result["result"]
445 )
446 kwargs = ( # rebind-ok: branch bind in shared if/elif dispatch
447 gemini_passthrough_logging_handler_result["kwargs"]
448 )
449 elif endpoint_type == EndpointType.OPENAI:
450 openai_passthrough_logging_handler_result: Final = (
451 OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks(
452 litellm_logging_obj=litellm_logging_obj,
453 passthrough_success_handler_obj=passthrough_success_handler_obj,
454 url_route=url_route,
455 request_body=request_body,
456 endpoint_type=endpoint_type,
457 start_time=start_time,
458 all_chunks=all_chunks,
459 end_time=end_time,
460 )
461 )
462 standard_logging_response_object = openai_passthrough_logging_handler_result["result"]
463 kwargs = openai_passthrough_logging_handler_result["kwargs"]
465 if standard_logging_response_object is None:
466 standard_logging_response_object = StandardPassThroughResponseObject(
467 response=f"cannot parse chunks to standard response object. Chunks={all_chunks}"
468 )
469 return standard_logging_response_object, kwargs
471 @staticmethod
472 def _extract_model_for_cost_injection(
473 request_body: dict | None,
474 url_route: str,
475 endpoint_type: EndpointType,
476 litellm_logging_obj: LiteLLMLoggingObj,
477 ) -> str | None:
478 """
479 Extract model name for cost injection from various sources.
480 """
481 # Try to get model from request body
482 if request_body:
483 model = request_body.get("model")
484 if model:
485 return model
487 # Try to get model from logging object
488 if hasattr(litellm_logging_obj, "model_call_details"):
489 model = litellm_logging_obj.model_call_details.get("model")
490 if model:
491 return model
493 # For Vertex AI, try to extract from URL
494 if endpoint_type == EndpointType.VERTEX_AI:
495 model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
496 if model and model != "unknown":
497 return model
499 return None
501 @staticmethod
502 def _convert_raw_bytes_to_str_lines(raw_bytes: Sequence[bytes]) -> list[str]:
503 """
504 Converts a list of raw bytes into a list of string lines, similar to aiter_lines()
506 Args:
507 raw_bytes: List of bytes chunks from aiter.bytes()
509 Returns:
510 List of string lines, with each line being a complete data: {} chunk
511 """
512 # errors="replace" so a stream cut mid-multibyte-sequence (client disconnect)
513 # still decodes and logs the usage events already received, instead of raising
514 # and dropping the whole request from SpendLogs
515 combined_str: Final = b"".join(raw_bytes).decode("utf-8", errors="replace")
517 # Split by newlines and filter out empty lines
518 lines: Final = [line.strip() for line in combined_str.split("\n") if line.strip()]
520 return lines