From 220ae6600d361469fe88d14f73da3307f68afe24 Mon Sep 17 00:00:00 2001 From: yuwk <1729065730@qq.com> Date: Sat, 29 Aug 2026 22:16:34 +0800 Subject: [PATCH] =?UTF-8?q?fix(anthropic):=20=E9=99=90=E5=88=B6=E6=B5=81?= =?UTF-8?q?=E5=BC=8F=E5=88=AB=E5=90=8D=E9=87=8D=E5=86=99=E7=BC=93=E5=86=B2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 将 Anthropic SSE 重写逻辑下沉到提供商模块。 - 首帧缓冲超过上限时切换为原样透传。 - 重写 message_start 后停止缓冲并补充边界测试。 --- .../messages/model_alias_rewriter.py | 106 ++++++++++++++++++ litellm/proxy/common_request_processing.py | 81 +------------ .../proxy/test_response_model_sanitization.py | 43 ++++++- 3 files changed, 148 insertions(+), 82 deletions(-) create mode 100644 litellm/llms/anthropic/experimental_pass_through/messages/model_alias_rewriter.py diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/model_alias_rewriter.py b/litellm/llms/anthropic/experimental_pass_through/messages/model_alias_rewriter.py new file mode 100644 index 00000000000..b34126a9ef9 --- /dev/null +++ b/litellm/llms/anthropic/experimental_pass_through/messages/model_alias_rewriter.py @@ -0,0 +1,106 @@ +"""Bounded Anthropic SSE rewriting for the public model alias.""" + +import json +from collections.abc import AsyncGenerator, AsyncIterator +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames + +MAX_MESSAGE_START_BUFFER_BYTES: Final = 1024 * 1024 +_SSE_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def _restamp_message_start_data_line(line: bytes, requested_model: str) -> tuple[bytes, bool]: + """Replace the model in one Anthropic message_start data line.""" + body: Final = line.rstrip(b"\r\n") + line_ending: Final = line[len(body) :] + if not body.startswith(b"data:"): + return line, False + try: + payload: Final = _SSE_OBJECT_ADAPTER.validate_json(body[len(b"data:") :].strip()) + except ValidationError: + return line, False + if payload.get("type") != "message_start": + return line, False + try: + message: Final = _SSE_OBJECT_ADAPTER.validate_python(payload.get("message")) + except ValidationError: + return line, False + if "model" not in message: + return line, False + restamped_payload: Final = {**payload, "message": {**message, "model": requested_model}} + rewritten: Final = b"data: " + json.dumps(restamped_payload, separators=(",", ":")).encode() + line_ending + return rewritten, True + + +def _restamp_message_start_frames(frames: bytes, requested_model: str) -> tuple[bytes, bool]: + """Restore the alias in complete frames and report when message_start was found.""" + output: list[bytes] = [] + restamped = False + for line in frames.splitlines(keepends=True): + rewritten_line, line_restamped = _restamp_message_start_data_line(line, requested_model) + output.append(rewritten_line) + restamped = restamped or line_restamped + return b"".join(output), restamped + + +class AnthropicMessageStartModelRewriter: + """Reassemble only the first Anthropic frames with a bounded pending buffer.""" + + __slots__ = ("_finished", "_pending", "_requested_model") + + def __init__(self, requested_model: str) -> None: + self._requested_model: Final = requested_model + self._pending = b"" + self._finished = False + + def feed(self, chunk: bytes) -> bytes: + """Return rewritten frames, or bounded pass-through for a malformed stream.""" + if self._finished: + return chunk + + buffered: Final = self._pending + chunk + if len(buffered) > MAX_MESSAGE_START_BUFFER_BYTES: + self._pending = b"" + self._finished = True + return buffered + + complete_frames, self._pending = split_complete_sse_frames(buffered) + rewritten_frames, restamped = _restamp_message_start_frames( + complete_frames, + self._requested_model, + ) + if not restamped: + return rewritten_frames + + tail: Final = self._pending + self._pending = b"" + self._finished = True + return rewritten_frames + tail + + def flush(self) -> bytes: + """Return the final unterminated tail without dropping provider bytes.""" + tail: Final = self._pending + self._pending = b"" + return tail + + +async def restamp_anthropic_streaming_response_model( + response: AsyncIterator[object], + requested_model: str, +) -> AsyncGenerator[object, None]: + """Restamp message_start while preserving arbitrary transport chunk boundaries.""" + rewriter: Final = AnthropicMessageStartModelRewriter(requested_model) + async for chunk in response: + if isinstance(chunk, (bytes, bytearray)): + if rewritten_frames := rewriter.feed(bytes(chunk)): + yield rewritten_frames + continue + if buffered_tail := rewriter.flush(): + yield buffered_tail + yield chunk + final_tail: Final = rewriter.flush() + if final_tail: + yield final_tail diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8d738999bda..ca1c35c5086 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,7 +4,7 @@ import json import logging import math import traceback -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -15,7 +15,6 @@ import httpx import orjson from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse -from pydantic import TypeAdapter, ValidationError from starlette.types import Receive, Scope, Send import litellm @@ -49,6 +48,9 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) +from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import ( + restamp_anthropic_streaming_response_model, +) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model from litellm.proxy.auth.auth_utils import check_response_size_is_safe @@ -60,7 +62,6 @@ from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, coerce_keepalive_interval, resolve_ttft_keepalive_interval, - split_complete_sse_frames, wrap_sse_stream_with_keepalive_pings, ) from litellm.proxy.dd_span_tagger import DDSpanTagger @@ -1268,78 +1269,6 @@ def _override_openai_response_model( ) -_ANTHROPIC_SSE_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object]) - - -def _restamp_anthropic_message_start_data_line(line: bytes, requested_model: str) -> bytes: - """Replace the model in one Anthropic message_start data line.""" - body: Final = line.rstrip(b"\r\n") - line_ending: Final = line[len(body) :] - if not body.startswith(b"data:"): - return line - try: - payload: Final = _ANTHROPIC_SSE_OBJECT_ADAPTER.validate_json(body[len(b"data:") :].strip()) - except ValidationError: - return line - if payload.get("type") != "message_start": - return line - try: - message: Final = _ANTHROPIC_SSE_OBJECT_ADAPTER.validate_python(payload.get("message")) - except ValidationError: - return line - if "model" not in message: - return line - restamped_payload: Final = {**payload, "message": {**message, "model": requested_model}} - return b"data: " + json.dumps(restamped_payload, separators=(",", ":")).encode() + line_ending - - -def _restamp_anthropic_message_start_frames(frames: bytes, requested_model: str) -> bytes: - """Restore the client-requested model in complete Anthropic SSE frames.""" - return b"".join( - _restamp_anthropic_message_start_data_line(line, requested_model) for line in frames.splitlines(keepends=True) - ) - - -class _AnthropicMessageStartModelRewriter: - """Reassemble transport chunks before rewriting Anthropic SSE frames.""" - - __slots__ = ("_pending", "_requested_model") - - def __init__(self, requested_model: str) -> None: - self._requested_model: Final = requested_model - self._pending = b"" - - def feed(self, chunk: bytes) -> bytes: - """Buffer a transport chunk and return rewritten complete frames.""" - complete_frames, self._pending = split_complete_sse_frames(self._pending + chunk) - return _restamp_anthropic_message_start_frames(complete_frames, self._requested_model) - - def flush(self) -> bytes: - """Return the final unterminated tail without dropping provider bytes.""" - tail: Final = self._pending - self._pending = b"" - return tail - - -async def _restamp_anthropic_streaming_response_model( - response: AsyncIterator[object], - requested_model: str, -) -> AsyncGenerator[object, None]: - """Restamp message_start while preserving arbitrary transport chunk boundaries.""" - rewriter: Final = _AnthropicMessageStartModelRewriter(requested_model) - async for chunk in response: - if isinstance(chunk, (bytes, bytearray)): - if rewritten_frames := rewriter.feed(bytes(chunk)): - yield rewritten_frames - continue - if buffered_tail := rewriter.flush(): - yield buffered_tail - yield chunk - final_tail: Final = rewriter.flush() - if final_tail: - yield final_tail - - def _anthropic_client_requested_model(request_data: Mapping[str, object]) -> str | None: """Return the public model name only for native Anthropic Messages streams.""" logging_obj: Final = request_data.get("litellm_logging_obj") @@ -3493,7 +3422,7 @@ class ProxyBaseLLMRequestProcessing: delivered_chunk = False anthropic_requested_model: Final = _anthropic_client_requested_model(request_data) stream_response: Final = ( - _restamp_anthropic_streaming_response_model(response, anthropic_requested_model) + restamp_anthropic_streaming_response_model(response, anthropic_requested_model) if anthropic_requested_model is not None else response ) diff --git a/tests/test_litellm/proxy/test_response_model_sanitization.py b/tests/test_litellm/proxy/test_response_model_sanitization.py index c63cc8881f2..eac7739dc1f 100644 --- a/tests/test_litellm/proxy/test_response_model_sanitization.py +++ b/tests/test_litellm/proxy/test_response_model_sanitization.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterator from unittest.mock import AsyncMock, MagicMock import pytest @@ -338,22 +338,22 @@ async def test_proxy_streaming_chunks_use_client_requested_model_before_alias_ma @pytest.mark.asyncio async def test_anthropic_streaming_message_start_uses_client_requested_model_across_chunks( - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ): """Keep the public model alias when an Anthropic message_start spans transport chunks.""" from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - async def upstream_stream(): + async def upstream_stream() -> AsyncGenerator[bytes, None]: yield b'event: message_start\ndata: {"type":"message_start","message":{"model":"backend' yield b'-model","id":"msg_1"}}\n\nevent: message_stop\ndata: {"type":"message_stop"}\n\n' async def passthrough_iterator_hook( user_api_key_dict: UserAPIKeyAuth, - response: AsyncGenerator, - request_data: dict, - ): + response: AsyncIterator[object], + request_data: dict[str, object], + ) -> AsyncGenerator[object, None]: async for chunk in response: yield chunk @@ -385,6 +385,37 @@ async def test_anthropic_streaming_message_start_uses_client_requested_model_acr assert b'event: message_stop\ndata: {"type":"message_stop"}\n\n' in streamed +def test_anthropic_message_start_rewriter_bounds_unterminated_streams(): + """Fall back to byte-for-byte pass-through when the first SSE frame is oversized.""" + from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import ( + MAX_MESSAGE_START_BUFFER_BYTES, + AnthropicMessageStartModelRewriter, + ) + + rewriter = AnthropicMessageStartModelRewriter("my-alias") + pending = b"x" * MAX_MESSAGE_START_BUFFER_BYTES + + assert rewriter.feed(pending) == b"" + assert rewriter.feed(b"y") == pending + b"y" + assert rewriter.feed(b"z") == b"z" + + +def test_anthropic_message_start_rewriter_stops_buffering_after_rewrite(): + """Pass later transport chunks through directly after the alias has been restored.""" + from litellm.llms.anthropic.experimental_pass_through.messages.model_alias_rewriter import ( + MAX_MESSAGE_START_BUFFER_BYTES, + AnthropicMessageStartModelRewriter, + ) + + rewriter = AnthropicMessageStartModelRewriter("my-alias") + message_start = b'event: message_start\ndata: {"type":"message_start","message":{"model":"backend-model"}}\n\n' + rewritten = rewriter.feed(message_start) + later_chunk = b"z" * (MAX_MESSAGE_START_BUFFER_BYTES + 1) + + assert b'"model":"my-alias"' in rewritten + assert rewriter.feed(later_chunk) == later_chunk + + @pytest.mark.asyncio async def test_proxy_streaming_azure_model_router_preserves_actual_model(monkeypatch): """