From 7f5abfffb4fb3fca21118ea99287c3e3de9df6e5 Mon Sep 17 00:00:00 2001 From: nate-berri Date: Sat, 10 Oct 2026 12:05:37 -0700 Subject: [PATCH] refactor(anthropic): move passthrough SSE stream assembly into litellm/llms (#45744) Part of LIT-9423 --- .../chat/guardrail_translation/handler.py | 11 +- .../anthropic/pass_through/stream_assembly.py | 397 +++++++++++++++++ litellm/proxy/guardrails/anthropic_sse.py | 14 +- .../anthropic_passthrough_logging_handler.py | 398 +----------------- .../test_anthropic_guardrail_handler.py | 8 +- .../test_bedrock_guardrails.py | 4 +- ...t_anthropic_passthrough_logging_handler.py | 5 +- 7 files changed, 426 insertions(+), 411 deletions(-) create mode 100644 litellm/llms/anthropic/pass_through/stream_assembly.py diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e3442bca1e7..7c1bee14b5b 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -29,6 +29,10 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import ( LiteLLMAnthropicMessagesAdapter, is_provider_native_tool_dict, ) +from litellm.llms.anthropic.pass_through.stream_assembly import ( + build_complete_streaming_response, + build_usage_only_response_from_chunks, +) from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -47,9 +51,6 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( stream_item_fingerprint, unappliable_request_rewrite, ) -from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( - AnthropicPassthroughLoggingHandler, -) from litellm.types.llms.anthropic import ( AllAnthropicToolsValues, AnthropicMessagesRequest, @@ -322,7 +323,7 @@ class AnthropicMessagesHandler(BaseTranslation): if not chunks: return None try: - return AnthropicPassthroughLoggingHandler.build_usage_only_response_from_chunks( + return build_usage_only_response_from_chunks( all_chunks=chunks, model=str((request_data or {}).get("model") or ""), ) @@ -1226,7 +1227,7 @@ class AnthropicMessagesHandler(BaseTranslation): has_ended: Final = self._check_streaming_has_ended(responses_so_far) if has_ended: # build the model response from the responses_so_far - built_response: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response( + built_response: Final = build_complete_streaming_response( all_chunks=responses_so_far, litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj), model="", diff --git a/litellm/llms/anthropic/pass_through/stream_assembly.py b/litellm/llms/anthropic/pass_through/stream_assembly.py new file mode 100644 index 00000000000..76d1896ac32 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/stream_assembly.py @@ -0,0 +1,397 @@ +import json +from collections.abc import Sequence +from typing import Final + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.core_helpers import map_finish_reason +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.chat.handler import ( + ModelResponseIterator as AnthropicModelResponseIterator, +) +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.types.utils import ( + Choices, + Message, + ModelResponse, + TextCompletionResponse, +) + + +def split_sse_chunk_into_events(chunk: str | bytes) -> list[str]: + """ + Split a chunk that may contain multiple SSE events into individual events. + + SSE format: "event: type\ndata: {...}\n\n" + Multiple events in a single chunk are separated by double newlines. + + Args: + chunk: Raw chunk string that may contain multiple SSE events + + Returns: + List of individual SSE event strings (each containing "event: X\ndata: {...}") + """ + # Handle bytes input + if isinstance(chunk, bytes): + chunk = chunk.decode("utf-8") + + # Split on double newlines to separate SSE events + # Filter out empty strings + events: Final = [event.strip() for event in chunk.split("\n\n") if event.strip()] + + return events + + +def build_complete_streaming_response( + all_chunks: Sequence[str | bytes], + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + speed: str | None = None, +) -> ModelResponse | TextCompletionResponse | None: + """ + Builds complete response from raw Anthropic chunks. + + Fast path: for the dominant case of a pure-text streaming response + (no tool_use / thinking / non-text content blocks), the long run of + ``content_block_delta`` text deltas is collapsed into a single + equivalent SSE event before conversion. ``chunk_parser`` and + ``stream_chunk_builder`` remain the single source of truth for chunk + shape, usage math and finish-reason mapping, so the rebuilt response + (and therefore the logged/billed payload) is identical -- this is + asserted by a parity test. Anything non-trivial falls back to the + unchanged legacy reconstruction. + + Per-event Pydantic ``ModelResponseStream`` construction dominated + event-loop CPU under concurrent streaming; collapsing the homogeneous + text run removes O(num_output_tokens) of it. + """ + collapsed: Final = collapse_pure_text_chunks(all_chunks) + if collapsed is not None: + return build_complete_streaming_response_legacy( + all_chunks=collapsed, + litellm_logging_obj=litellm_logging_obj, + model=model, + speed=speed, + ) + return build_complete_streaming_response_legacy( + all_chunks=all_chunks, + litellm_logging_obj=litellm_logging_obj, + model=model, + speed=speed, + ) + + +# Anthropic SSE block/delta types that the fast path is NOT allowed to +# collapse -- their presence forces the unchanged legacy path so tool +# calls, thinking, citations, etc. keep byte-identical reconstruction. +_FAST_PATH_DISALLOWED_DELTA_TYPES: Final = frozenset( + { + "input_json_delta", + "thinking_delta", + "signature_delta", + "citations_delta", + } +) + + +def collapse_pure_text_chunks( + all_chunks: Sequence[str | bytes], +) -> list[str] | None: + """ + Return a new chunk list with the contiguous run of text-only + ``content_block_delta`` events replaced by a single equivalent event, + or ``None`` if the stream is not a pure single-text-block response + (in which case the caller uses the legacy path unchanged). + + Only ``message_start`` / ``content_block_start(text)`` / + ``content_block_delta(text_delta)`` / ``content_block_stop`` / + ``message_delta`` / ``message_stop`` / ``ping`` events are accepted. + Any other content-block type or delta type returns ``None``. + """ + normalized: Final[list[str]] = [] + for raw in all_chunks: + line = raw.decode("utf-8") if isinstance(raw, bytes) else raw + for ev in line.split("\n\n"): + ev = ev.strip() + if ev: + normalized.append(ev) + + text_block_indexes: Final[set] = set() + out: Final[list[str]] = [] + pending_text: list[str] = [] + pending_index: int | None = None + saw_any_text_delta = False + + def flush() -> None: + nonlocal pending_text, pending_index + if pending_text: + merged: Final = { + "type": "content_block_delta", + "index": pending_index if pending_index is not None else 0, + "delta": {"type": "text_delta", "text": "".join(pending_text)}, + } + out.append("data: " + json.dumps(merged)) + pending_text = [] + pending_index = None + + for ev in normalized: + idx = ev.find("data:") + if idx == -1: + # Bare "event: " line. The legacy converter turns this + # into an empty ModelResponseStream that contributes nothing + # to stream_chunk_builder. Drop the high-frequency interior + # markers (content_block_delta / ping); keep every other + # bare event line verbatim so chunk ordering and the + # load-bearing chunks[0] (event: message_start) are retained. + name = ev[len("event:") :].strip() if ev.startswith("event:") else "" + if name in ("content_block_delta", "ping"): + continue + flush() + out.append(ev) + continue + + json_str = ev[idx + len("data:") :].strip() + try: + data = json.loads(json_str) + except (json.JSONDecodeError, ValueError): + return None + + etype = data.get("type") + if etype == "content_block_start": + block = data.get("content_block") or {} + if block.get("type") != "text": + return None + text_block_indexes.add(data.get("index")) + flush() + out.append(ev) + elif etype == "content_block_delta": + delta = data.get("delta") or {} + dtype = delta.get("type") + if dtype in _FAST_PATH_DISALLOWED_DELTA_TYPES: + return None + if dtype != "text_delta": + return None + cur_index = data.get("index") + if cur_index not in text_block_indexes: + return None + # Defensive: Anthropic sends blocks strictly sequentially + # (start/deltas/stop, then next block), so pending_text from + # block N must be flushed by content_block_stop before block + # N+1's deltas arrive. If we ever see a delta whose index + # disagrees with the current pending buffer, the stream is + # interleaved -- fall back to legacy rather than risk merging + # text from different blocks under a single index. + if pending_text and pending_index is not None and cur_index != pending_index: + return None + saw_any_text_delta = True + pending_index = cur_index + pending_text.append(delta.get("text") or "") + elif etype == "ping": + # Interior no-op; legacy maps it to an empty chunk. + continue + else: + # message_start / content_block_stop / message_delta / + # message_stop / error: pass through unchanged. + flush() + out.append(ev) + + flush() + + if not saw_any_text_delta: + return None + return out + + +def build_complete_streaming_response_legacy( + all_chunks: Sequence[str | bytes], + litellm_logging_obj: LiteLLMLoggingObj, + model: str, + speed: str | None = None, +) -> ModelResponse | TextCompletionResponse | None: + """ + Original reconstruction: convert every SSE event to a generic chunk + and assemble via stream_chunk_builder. Kept verbatim as the fallback + / source of truth for the fast path's parity test. + + - Splits multi-event chunks into individual SSE events + - Converts str chunks to generic chunks + - Converts generic chunks to litellm chunks (OpenAI format) + - Builds complete response from litellm chunks + """ + verbose_proxy_logger.debug("Building complete streaming response from %d chunks", len(all_chunks)) + anthropic_model_response_iterator: Final = AnthropicModelResponseIterator( + streaming_response=None, + sync_stream=False, + speed=speed, + ) + all_openai_chunks: Final = [] + + # Process each chunk - a chunk may contain multiple SSE events + for _chunk_str in all_chunks: + # Split chunk into individual SSE events + individual_events = split_sse_chunk_into_events(_chunk_str) + + # Process each individual event + for event_str in individual_events: + try: + # Skip OpenAI-style [DONE] sentinels some Anthropic-compatible + # providers emit. Match the whole SSE line so a valid chunk whose + # text payload happens to contain "[DONE]" is not dropped. + if any(line.strip() == "data: [DONE]" for line in event_str.split("\n")): + continue + transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( + chunk=event_str + ) + if transformed_openai_chunk is not None: + all_openai_chunks.append(transformed_openai_chunk) + + except (StopIteration, StopAsyncIteration): + break + except json.JSONDecodeError: + # Some upstreams emit non-JSON SSE lines; skip them so the + # logging pipeline is not broken by a single bad frame. + verbose_proxy_logger.debug( + "Skipping non-JSON SSE event: %s", + event_str[:200], + ) + continue + + complete_streaming_response: Final = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, + ) + verbose_proxy_logger.debug("Complete streaming response built: %s", complete_streaming_response) + return complete_streaming_response + + +def extract_sse_data(event_str: str) -> dict | None: + """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" + for line in event_str.splitlines(): + stripped = line.strip() + if stripped.startswith("data:"): + payload = stripped[len("data:") :].strip() + if not payload or payload == "[DONE]": + return None + try: + return json.loads(payload) + except (ValueError, TypeError): + return None + return None + + +def build_usage_only_response_from_chunks( + all_chunks: Sequence[str | bytes], + model: str, + speed: str | None = None, +) -> ModelResponse | None: + """ + Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for + cost tracking when stream_chunk_builder cannot reassemble the stream. + + Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an + initial output_tokens) and the final ``message_delta`` (cumulative output_tokens) + regardless of the content/tool shape, so cost is recoverable even when full + content assembly fails. Returns ``None`` if no usage event is found. + """ + input_tokens = 0 + cache_read = 0 + cache_creation = 0 + cache_creation_5m: int | None = None + cache_creation_1h: int | None = None + output_tokens = 0 + web_search_requests: int | None = None + tool_search_requests: int | None = None + inference_geo: str | None = None + speed_from_stream: str | None = None + stop_reason: str | None = None + found_usage = False + resolved_model = model + for _chunk_str in all_chunks: + for event_str in split_sse_chunk_into_events(_chunk_str): + data = extract_sse_data(event_str) + if not data: + continue + event_type = data.get("type") + if event_type == "message_start": + message = data.get("message") or {} + if not resolved_model or resolved_model == "unknown": + resolved_model = message.get("model") or resolved_model + usage = message.get("usage") or {} + input_tokens = usage.get("input_tokens") or input_tokens + cache_read = usage.get("cache_read_input_tokens") or cache_read + cache_creation = usage.get("cache_creation_input_tokens") or cache_creation + _cc = usage.get("cache_creation") + if isinstance(_cc, dict): + cache_creation_5m = _cc.get("ephemeral_5m_input_tokens") + cache_creation_1h = _cc.get("ephemeral_1h_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + if isinstance(usage.get("speed"), str): + speed_from_stream = usage.get("speed") + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + found_usage = True + elif event_type == "message_delta": + _delta_stop = (data.get("delta") or {}).get("stop_reason") + if _delta_stop: + stop_reason = _delta_stop + usage = data.get("usage") or {} + if usage.get("output_tokens") is not None: + output_tokens = usage.get("output_tokens") + _stu = usage.get("server_tool_use") + if isinstance(_stu, dict): + if _stu.get("web_search_requests") is not None: + web_search_requests = _stu.get("web_search_requests") + if _stu.get("tool_search_requests") is not None: + tool_search_requests = _stu.get("tool_search_requests") + if usage.get("cache_read_input_tokens") is not None: + cache_read = usage.get("cache_read_input_tokens") + if usage.get("inference_geo") is not None: + inference_geo = usage.get("inference_geo") + if isinstance(usage.get("speed"), str): + speed_from_stream = usage.get("speed") + found_usage = True + if not found_usage: + return None + # If only the 5m/1h split was provided, derive the cache_creation total from it. + if not cache_creation and (cache_creation_5m or cache_creation_1h): + cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) + # build usage via the same AnthropicConfig.calculate_usage path the success + # cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use / + # inference_geo tokens are priced instead of left at $0 + usage_object: Final[dict] = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + if cache_read: + usage_object["cache_read_input_tokens"] = cache_read + if cache_creation: + usage_object["cache_creation_input_tokens"] = cache_creation + if cache_creation_5m is not None or cache_creation_1h is not None: + usage_object["cache_creation"] = { + "ephemeral_5m_input_tokens": cache_creation_5m or 0, + "ephemeral_1h_input_tokens": cache_creation_1h or 0, + } + if web_search_requests is not None or tool_search_requests is not None: + _server_tool_use: Final[dict] = {} + if web_search_requests is not None: + _server_tool_use["web_search_requests"] = web_search_requests + if tool_search_requests is not None: + _server_tool_use["tool_search_requests"] = tool_search_requests + usage_object["server_tool_use"] = _server_tool_use + if inference_geo is not None: + usage_object["inference_geo"] = inference_geo + if speed_from_stream is not None: + usage_object["speed"] = speed_from_stream + usage_obj: Final = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None, speed=speed) + return ModelResponse( + model=resolved_model, + choices=[ + Choices( + finish_reason=(map_finish_reason(stop_reason) if stop_reason else "stop"), + index=0, + message=Message(role="assistant", content=""), + ) + ], + usage=usage_obj, + ) diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index c446cfaa815..07a5d511117 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -45,14 +45,12 @@ def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None: def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]: - from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( - AnthropicPassthroughLoggingHandler, - ) + from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data, split_sse_chunk_into_events return tuple( event_data - for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses - if (event_data := AnthropicPassthroughLoggingHandler._extract_sse_data(event)) is not None # pyright: ignore[reportPrivateUsage] # same parser the assembler uses; a private import beats forking SSE parsing + for event in split_sse_chunk_into_events(sse_stream) + if (event_data := extract_sse_data(event)) is not None ) @@ -91,9 +89,7 @@ def assemble_anthropic_sse_stream( response keep the wire shape they had before this helper was shared. The writes land on a freshly built object that is unreachable from caller state until returned. """ - from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( - AnthropicPassthroughLoggingHandler, - ) + from litellm.llms.anthropic.pass_through.stream_assembly import build_complete_streaming_response sse_stream: Final = _joined_sse_stream(all_chunks) if sse_stream is None: @@ -103,7 +99,7 @@ def assemble_anthropic_sse_stream( return None model: Final = message_start.get("model") if restore_identity else None try: - assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser + assembled: Final = build_complete_streaming_response( all_chunks=(sse_stream,), litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None model=model if isinstance(model, str) else "", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 73d40369d30..d787c16a534 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -10,17 +10,19 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE -from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) from litellm.llms.anthropic import get_anthropic_config -from litellm.llms.anthropic.chat.handler import ( - ModelResponseIterator as AnthropicModelResponseIterator, +from litellm.llms.anthropic.pass_through.stream_assembly import ( + build_complete_streaming_response, + build_complete_streaming_response_legacy, + build_usage_only_response_from_chunks, + collapse_pure_text_chunks, + split_sse_chunk_into_events, ) -from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -33,9 +35,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( PassthroughStandardLoggingPayload, ) from litellm.types.utils import ( - Choices, LiteLLMBatch, - Message, ModelResponse, TextCompletionResponse, Usage, @@ -516,388 +516,12 @@ class AnthropicPassthroughLoggingHandler: _handle_logging_anthropic_collected_chunks = handle_logging_anthropic_collected_chunks - @staticmethod - def _split_sse_chunk_into_events(chunk: str | bytes) -> list[str]: - """ - Split a chunk that may contain multiple SSE events into individual events. - - SSE format: "event: type\ndata: {...}\n\n" - Multiple events in a single chunk are separated by double newlines. - - Args: - chunk: Raw chunk string that may contain multiple SSE events - - Returns: - List of individual SSE event strings (each containing "event: X\ndata: {...}") - """ - # Handle bytes input - if isinstance(chunk, bytes): - chunk = chunk.decode("utf-8") - - # Split on double newlines to separate SSE events - # Filter out empty strings - events: Final = [event.strip() for event in chunk.split("\n\n") if event.strip()] - - return events - - @staticmethod - def build_complete_streaming_response( - all_chunks: Sequence[str | bytes], - litellm_logging_obj: LiteLLMLoggingObj, - model: str, - speed: str | None = None, - ) -> ModelResponse | TextCompletionResponse | None: - """ - Builds complete response from raw Anthropic chunks. - - Fast path: for the dominant case of a pure-text streaming response - (no tool_use / thinking / non-text content blocks), the long run of - ``content_block_delta`` text deltas is collapsed into a single - equivalent SSE event before conversion. ``chunk_parser`` and - ``stream_chunk_builder`` remain the single source of truth for chunk - shape, usage math and finish-reason mapping, so the rebuilt response - (and therefore the logged/billed payload) is identical -- this is - asserted by a parity test. Anything non-trivial falls back to the - unchanged legacy reconstruction. - - Per-event Pydantic ``ModelResponseStream`` construction dominated - event-loop CPU under concurrent streaming; collapsing the homogeneous - text run removes O(num_output_tokens) of it. - """ - collapsed: Final = AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks) - if collapsed is not None: - return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy( - all_chunks=collapsed, - litellm_logging_obj=litellm_logging_obj, - model=model, - speed=speed, - ) - return AnthropicPassthroughLoggingHandler._build_complete_streaming_response_legacy( - all_chunks=all_chunks, - litellm_logging_obj=litellm_logging_obj, - model=model, - speed=speed, - ) - + _split_sse_chunk_into_events = staticmethod(split_sse_chunk_into_events) + build_complete_streaming_response = staticmethod(build_complete_streaming_response) _build_complete_streaming_response = build_complete_streaming_response - - # Anthropic SSE block/delta types that the fast path is NOT allowed to - # collapse -- their presence forces the unchanged legacy path so tool - # calls, thinking, citations, etc. keep byte-identical reconstruction. - _FAST_PATH_DISALLOWED_DELTA_TYPES = frozenset( - { - "input_json_delta", - "thinking_delta", - "signature_delta", - "citations_delta", - } - ) - - @staticmethod - def _collapse_pure_text_chunks( - all_chunks: Sequence[str | bytes], - ) -> list[str] | None: - """ - Return a new chunk list with the contiguous run of text-only - ``content_block_delta`` events replaced by a single equivalent event, - or ``None`` if the stream is not a pure single-text-block response - (in which case the caller uses the legacy path unchanged). - - Only ``message_start`` / ``content_block_start(text)`` / - ``content_block_delta(text_delta)`` / ``content_block_stop`` / - ``message_delta`` / ``message_stop`` / ``ping`` events are accepted. - Any other content-block type or delta type returns ``None``. - """ - normalized: Final[list[str]] = [] - for raw in all_chunks: - line = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for ev in line.split("\n\n"): - ev = ev.strip() - if ev: - normalized.append(ev) - - text_block_indexes: Final[set] = set() - out: Final[list[str]] = [] - pending_text: list[str] = [] - pending_index: int | None = None - saw_any_text_delta = False - - def flush() -> None: - nonlocal pending_text, pending_index - if pending_text: - merged: Final = { - "type": "content_block_delta", - "index": pending_index if pending_index is not None else 0, - "delta": {"type": "text_delta", "text": "".join(pending_text)}, - } - out.append("data: " + json.dumps(merged)) - pending_text = [] - pending_index = None - - for ev in normalized: - idx = ev.find("data:") - if idx == -1: - # Bare "event: " line. The legacy converter turns this - # into an empty ModelResponseStream that contributes nothing - # to stream_chunk_builder. Drop the high-frequency interior - # markers (content_block_delta / ping); keep every other - # bare event line verbatim so chunk ordering and the - # load-bearing chunks[0] (event: message_start) are retained. - name = ev[len("event:") :].strip() if ev.startswith("event:") else "" - if name in ("content_block_delta", "ping"): - continue - flush() - out.append(ev) - continue - - json_str = ev[idx + len("data:") :].strip() - try: - data = json.loads(json_str) - except (json.JSONDecodeError, ValueError): - return None - - etype = data.get("type") - if etype == "content_block_start": - block = data.get("content_block") or {} - if block.get("type") != "text": - return None - text_block_indexes.add(data.get("index")) - flush() - out.append(ev) - elif etype == "content_block_delta": - delta = data.get("delta") or {} - dtype = delta.get("type") - if dtype in AnthropicPassthroughLoggingHandler._FAST_PATH_DISALLOWED_DELTA_TYPES: - return None - if dtype != "text_delta": - return None - cur_index = data.get("index") - if cur_index not in text_block_indexes: - return None - # Defensive: Anthropic sends blocks strictly sequentially - # (start/deltas/stop, then next block), so pending_text from - # block N must be flushed by content_block_stop before block - # N+1's deltas arrive. If we ever see a delta whose index - # disagrees with the current pending buffer, the stream is - # interleaved -- fall back to legacy rather than risk merging - # text from different blocks under a single index. - if pending_text and pending_index is not None and cur_index != pending_index: - return None - saw_any_text_delta = True - pending_index = cur_index - pending_text.append(delta.get("text") or "") - elif etype == "ping": - # Interior no-op; legacy maps it to an empty chunk. - continue - else: - # message_start / content_block_stop / message_delta / - # message_stop / error: pass through unchanged. - flush() - out.append(ev) - - flush() - - if not saw_any_text_delta: - return None - return out - - @staticmethod - def _build_complete_streaming_response_legacy( - all_chunks: Sequence[str | bytes], - litellm_logging_obj: LiteLLMLoggingObj, - model: str, - speed: str | None = None, - ) -> ModelResponse | TextCompletionResponse | None: - """ - Original reconstruction: convert every SSE event to a generic chunk - and assemble via stream_chunk_builder. Kept verbatim as the fallback - / source of truth for the fast path's parity test. - - - Splits multi-event chunks into individual SSE events - - Converts str chunks to generic chunks - - Converts generic chunks to litellm chunks (OpenAI format) - - Builds complete response from litellm chunks - """ - verbose_proxy_logger.debug("Building complete streaming response from %d chunks", len(all_chunks)) - anthropic_model_response_iterator: Final = AnthropicModelResponseIterator( - streaming_response=None, - sync_stream=False, - speed=speed, - ) - all_openai_chunks: Final = [] - - # Process each chunk - a chunk may contain multiple SSE events - for _chunk_str in all_chunks: - # Split chunk into individual SSE events - individual_events = AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(_chunk_str) - - # Process each individual event - for event_str in individual_events: - try: - # Skip OpenAI-style [DONE] sentinels some Anthropic-compatible - # providers emit. Match the whole SSE line so a valid chunk whose - # text payload happens to contain "[DONE]" is not dropped. - if any(line.strip() == "data: [DONE]" for line in event_str.split("\n")): - continue - transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk( - chunk=event_str - ) - if transformed_openai_chunk is not None: - all_openai_chunks.append(transformed_openai_chunk) - - except (StopIteration, StopAsyncIteration): - break - except json.JSONDecodeError: - # Some upstreams emit non-JSON SSE lines; skip them so the - # logging pipeline is not broken by a single bad frame. - verbose_proxy_logger.debug( - "Skipping non-JSON SSE event: %s", - event_str[:200], - ) - continue - - complete_streaming_response: Final = litellm.stream_chunk_builder( - chunks=all_openai_chunks, - logging_obj=litellm_logging_obj, - ) - verbose_proxy_logger.debug("Complete streaming response built: %s", complete_streaming_response) - return complete_streaming_response - - @staticmethod - def _extract_sse_data(event_str: str) -> dict | None: - """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" - for line in event_str.splitlines(): - stripped = line.strip() - if stripped.startswith("data:"): - payload = stripped[len("data:") :].strip() - if not payload or payload == "[DONE]": - return None - try: - return cast(dict, json.loads(payload)) - except (ValueError, TypeError): - return None - return None - - @staticmethod - def build_usage_only_response_from_chunks( - all_chunks: Sequence[str | bytes], - model: str, - speed: str | None = None, - ) -> ModelResponse | None: - """ - Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for - cost tracking when stream_chunk_builder cannot reassemble the stream. - - Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an - initial output_tokens) and the final ``message_delta`` (cumulative output_tokens) - regardless of the content/tool shape, so cost is recoverable even when full - content assembly fails. Returns ``None`` if no usage event is found. - """ - input_tokens = 0 - cache_read = 0 - cache_creation = 0 - cache_creation_5m: int | None = None - cache_creation_1h: int | None = None - output_tokens = 0 - web_search_requests: int | None = None - tool_search_requests: int | None = None - inference_geo: str | None = None - speed_from_stream: str | None = None - stop_reason: str | None = None - found_usage = False - resolved_model = model - for _chunk_str in all_chunks: - for event_str in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(_chunk_str): - data = AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) - if not data: - continue - event_type = data.get("type") - if event_type == "message_start": - message = data.get("message") or {} - if not resolved_model or resolved_model == "unknown": - resolved_model = message.get("model") or resolved_model - usage = message.get("usage") or {} - input_tokens = usage.get("input_tokens") or input_tokens - cache_read = usage.get("cache_read_input_tokens") or cache_read - cache_creation = usage.get("cache_creation_input_tokens") or cache_creation - _cc = usage.get("cache_creation") - if isinstance(_cc, dict): - cache_creation_5m = _cc.get("ephemeral_5m_input_tokens") - cache_creation_1h = _cc.get("ephemeral_1h_input_tokens") - if usage.get("inference_geo") is not None: - inference_geo = usage.get("inference_geo") - if isinstance(usage.get("speed"), str): - speed_from_stream = usage.get("speed") - if usage.get("output_tokens") is not None: - output_tokens = usage.get("output_tokens") - found_usage = True - elif event_type == "message_delta": - _delta_stop = (data.get("delta") or {}).get("stop_reason") - if _delta_stop: - stop_reason = _delta_stop - usage = data.get("usage") or {} - if usage.get("output_tokens") is not None: - output_tokens = usage.get("output_tokens") - _stu = usage.get("server_tool_use") - if isinstance(_stu, dict): - if _stu.get("web_search_requests") is not None: - web_search_requests = _stu.get("web_search_requests") - if _stu.get("tool_search_requests") is not None: - tool_search_requests = _stu.get("tool_search_requests") - if usage.get("cache_read_input_tokens") is not None: - cache_read = usage.get("cache_read_input_tokens") - if usage.get("inference_geo") is not None: - inference_geo = usage.get("inference_geo") - if isinstance(usage.get("speed"), str): - speed_from_stream = usage.get("speed") - found_usage = True - if not found_usage: - return None - # If only the 5m/1h split was provided, derive the cache_creation total from it. - if not cache_creation and (cache_creation_5m or cache_creation_1h): - cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0) - # build usage via the same AnthropicConfig.calculate_usage path the success - # cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use / - # inference_geo tokens are priced instead of left at $0 - usage_object: Final[dict] = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - } - if cache_read: - usage_object["cache_read_input_tokens"] = cache_read - if cache_creation: - usage_object["cache_creation_input_tokens"] = cache_creation - if cache_creation_5m is not None or cache_creation_1h is not None: - usage_object["cache_creation"] = { - "ephemeral_5m_input_tokens": cache_creation_5m or 0, - "ephemeral_1h_input_tokens": cache_creation_1h or 0, - } - if web_search_requests is not None or tool_search_requests is not None: - _server_tool_use: Final[dict] = {} - if web_search_requests is not None: - _server_tool_use["web_search_requests"] = web_search_requests - if tool_search_requests is not None: - _server_tool_use["tool_search_requests"] = tool_search_requests - usage_object["server_tool_use"] = _server_tool_use - if inference_geo is not None: - usage_object["inference_geo"] = inference_geo - if speed_from_stream is not None: - usage_object["speed"] = speed_from_stream - usage_obj: Final = AnthropicConfig().calculate_usage( - usage_object=usage_object, reasoning_content=None, speed=speed - ) - return ModelResponse( - model=resolved_model, - choices=[ - Choices( - finish_reason=(map_finish_reason(stop_reason) if stop_reason else "stop"), - index=0, - message=Message(role="assistant", content=""), - ) - ], - usage=usage_obj, - ) - + _collapse_pure_text_chunks = staticmethod(collapse_pure_text_chunks) + _build_complete_streaming_response_legacy = staticmethod(build_complete_streaming_response_legacy) + build_usage_only_response_from_chunks = staticmethod(build_usage_only_response_from_chunks) _build_usage_only_response_from_chunks = build_usage_only_response_from_chunks @staticmethod diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 9c0fb5c6a98..4f27d9637af 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -182,7 +182,7 @@ class TestAnthropicMessagesHandlerStreamingRequestData: with ( patch.object(handler, "_check_streaming_has_ended", return_value=True), patch( - "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler.build_complete_streaming_response", + "litellm.llms.anthropic.chat.guardrail_translation.handler.build_complete_streaming_response", return_value=mock_response, ), ): @@ -241,7 +241,7 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: with ( patch.object(handler, "_check_streaming_has_ended", return_value=True), patch( - "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler.build_complete_streaming_response", + "litellm.llms.anthropic.chat.guardrail_translation.handler.build_complete_streaming_response", return_value=None, ), ): @@ -1353,7 +1353,7 @@ class TestAnthropicMessagesHandlerInputProcessing: with ( patch.object(handler, "_check_streaming_has_ended", return_value=True), patch( - "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler.build_complete_streaming_response", + "litellm.llms.anthropic.chat.guardrail_translation.handler.build_complete_streaming_response", return_value=mock_response, ), ): @@ -1400,7 +1400,7 @@ class TestAnthropicMessagesHandlerInputProcessing: with ( patch.object(handler, "_check_streaming_has_ended", return_value=True), patch( - "litellm.llms.anthropic.chat.guardrail_translation.handler.AnthropicPassthroughLoggingHandler.build_complete_streaming_response", + "litellm.llms.anthropic.chat.guardrail_translation.handler.build_complete_streaming_response", return_value=mock_response, ), ): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index f9795197f90..ccf3852c2c5 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -3255,9 +3255,7 @@ async def test_streaming_hook_fails_closed_when_assembler_raises_api_error(): guardrail = _sse_guardrail() with patch( - "litellm.proxy.pass_through_endpoints.llm_provider_handlers." - "anthropic_passthrough_logging_handler.AnthropicPassthroughLoggingHandler." - "_build_complete_streaming_response", + "litellm.llms.anthropic.pass_through.stream_assembly.build_complete_streaming_response", side_effect=litellm.APIError( status_code=500, message="Error building chunks for logging/streaming usage calculation", diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 93afea1b106..24fb60e5389 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -11,6 +11,7 @@ import pytest from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) @@ -2288,9 +2289,7 @@ class TestAnthropicUsageOnlyFallback: def test_extract_sse_data_handles_malformed_and_sentinel_lines( self, event_str, expected ): - assert ( - AnthropicPassthroughLoggingHandler._extract_sse_data(event_str) == expected - ) + assert extract_sse_data(event_str) == expected def _real_logging_obj(self): from litellm.litellm_core_utils.litellm_logging import Logging as RealLoggingObj