mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge pull request #35816 from BerriAI/litellm_anthropic_stream_model_alias
fix(proxy): report requested model on Anthropic streaming message_start
This commit is contained in:
Mateo Wang
GitHub
commit
3dac3f7a36
3 changed files with 495 additions and 3 deletions
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
174
litellm/proxy/anthropic_endpoints/streaming_model_restamp.py
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
"""
|
||||
Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only
|
||||
stream event carrying a model, so streamed responses report the requested model like
|
||||
non-streaming ones do.
|
||||
|
||||
Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the
|
||||
provider passthrough path) or as event dicts (fake-stream and agentic paths).
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
_MESSAGE_START_EVENT: Final = "message_start"
|
||||
_MESSAGE_START_MARKER: Final = b"message_start"
|
||||
_SSE_DATA_FIELD: Final = "data:"
|
||||
_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n")
|
||||
_MAX_HELD_BYTES: Final = 65536
|
||||
_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"')
|
||||
|
||||
_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None:
|
||||
message: Final = event.get("message")
|
||||
if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict):
|
||||
return None
|
||||
if message.get("model") == requested_model:
|
||||
return None
|
||||
return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is
|
||||
|
||||
|
||||
def _restamped_data_line(line: str, requested_model: str) -> str | None:
|
||||
stripped: Final = line.strip()
|
||||
if not stripped.startswith(_SSE_DATA_FIELD):
|
||||
return None
|
||||
payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
return None
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_json(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
restamped: Final = _restamped_event(event, requested_model)
|
||||
if restamped is None:
|
||||
return None
|
||||
terminator: Final = line[len(line.rstrip("\r\n")) :]
|
||||
return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}"
|
||||
|
||||
|
||||
def _restamped_frame(frame: str, requested_model: str) -> str | None:
|
||||
lines: Final = frame.splitlines(keepends=True)
|
||||
restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines)
|
||||
if all(line is None for line in restamped):
|
||||
return None
|
||||
return "".join(new if new is not None else old for new, old in zip(restamped, lines))
|
||||
|
||||
|
||||
def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object:
|
||||
"""
|
||||
Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``.
|
||||
|
||||
Chunks that carry no model are returned unchanged.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
try:
|
||||
event: Final = _EVENT_ADAPTER.validate_python(chunk)
|
||||
except ValidationError:
|
||||
return chunk
|
||||
return _restamped_event(event, requested_model) or chunk
|
||||
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
if _MESSAGE_START_EVENT.encode() not in chunk:
|
||||
return chunk
|
||||
restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model)
|
||||
return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8")
|
||||
|
||||
if isinstance(chunk, str):
|
||||
if _MESSAGE_START_EVENT not in chunk:
|
||||
return chunk
|
||||
restamped_text: Final = _restamped_frame(chunk, requested_model)
|
||||
return chunk if restamped_text is None else restamped_text
|
||||
|
||||
return chunk
|
||||
|
||||
|
||||
def _is_ping_frame(frame: bytes) -> bool:
|
||||
return any(marker in frame for marker in _PING_MARKERS)
|
||||
|
||||
|
||||
class AnthropicStreamModelRestamper:
|
||||
"""
|
||||
Per-stream restamper for the encoded passthrough path, where chunks are raw
|
||||
transport reads: the ``message_start`` SSE frame can arrive split across
|
||||
chunks or coalesced with later frames. Complete frames (``\\n\\n``,
|
||||
``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator
|
||||
closes them and an incomplete tail is held until it completes, so the
|
||||
restamp never misses a torn frame; ``flush`` returns whatever is still held
|
||||
when the stream ends so no bytes are swallowed. Once ``message_start`` has
|
||||
been handled, or the first real event proves the stream carries none, every
|
||||
later chunk passes through untouched.
|
||||
"""
|
||||
|
||||
def __init__(self, requested_model: str) -> None:
|
||||
self._requested_model: Final = requested_model
|
||||
self._held = b""
|
||||
self._armed = True
|
||||
|
||||
def process(self, chunk: object) -> object:
|
||||
if not self._armed:
|
||||
return chunk
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
return self._process_encoded(bytes(chunk))
|
||||
if isinstance(chunk, str):
|
||||
return self._process_encoded(chunk.encode("utf-8"))
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model)
|
||||
if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"):
|
||||
self._armed = False
|
||||
return restamped
|
||||
|
||||
def flush(self) -> bytes:
|
||||
held: Final = self._held
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
if not held:
|
||||
return b""
|
||||
restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model)
|
||||
return restamped if isinstance(restamped, bytes) else held
|
||||
|
||||
def _process_encoded(self, data: bytes) -> bytes:
|
||||
combined: Final = self._held + data
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined))
|
||||
if not boundaries:
|
||||
if len(combined) > _MAX_HELD_BYTES:
|
||||
self._held = b""
|
||||
self._armed = False
|
||||
return combined
|
||||
self._held = combined
|
||||
return b""
|
||||
emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]])
|
||||
tail: Final = combined[boundaries[-1] :]
|
||||
if not self._armed:
|
||||
self._held = b""
|
||||
return emitted + tail
|
||||
self._held = tail
|
||||
return emitted
|
||||
|
||||
def _restamped_closed_block(self, closed: bytes) -> bytes:
|
||||
boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed))
|
||||
frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries))
|
||||
decider: Final = next(
|
||||
(
|
||||
index
|
||||
for index, frame in enumerate(frames)
|
||||
if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame))
|
||||
),
|
||||
None,
|
||||
)
|
||||
if decider is None:
|
||||
return closed
|
||||
self._armed = False
|
||||
if _MESSAGE_START_MARKER not in frames[decider]:
|
||||
return closed
|
||||
restamped_text: Final = _restamped_frame(
|
||||
frames[decider].decode("utf-8", errors="ignore"), self._requested_model
|
||||
)
|
||||
if restamped_text is None:
|
||||
return closed
|
||||
return b"".join(
|
||||
restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames)
|
||||
)
|
||||
|
|
@ -176,6 +176,9 @@ if TYPE_CHECKING:
|
|||
ProxyConfig = _ProxyConfig
|
||||
else:
|
||||
ProxyConfig = Any
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
add_litellm_data_to_request,
|
||||
refresh_proxy_server_request_body_snapshot,
|
||||
|
|
@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data=self.data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
request=request,
|
||||
restamp_model=(
|
||||
None if _should_return_raw_model_name(self.data) else requested_model_from_client
|
||||
),
|
||||
)
|
||||
return await create_response(
|
||||
generator=wrap_sse_stream_with_keepalive_pings(
|
||||
|
|
@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
else:
|
||||
return chunk
|
||||
|
||||
@staticmethod
|
||||
def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer:
|
||||
if restamper is None:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk
|
||||
|
||||
def serialize(chunk: object) -> str:
|
||||
return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk))
|
||||
|
||||
return serialize
|
||||
|
||||
@staticmethod
|
||||
async def _finalize_streaming_generator_cleanup(
|
||||
request: Request | None,
|
||||
|
|
@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing:
|
|||
serialize_chunk: StreamChunkSerializer,
|
||||
serialize_error: StreamErrorSerializer,
|
||||
request: Request | None = None,
|
||||
flush_tail: Callable[[], bytes] | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Shared streaming data generator: runs proxy iterator hook, per-chunk hook,
|
||||
cost injection, then yields chunks via serialize_chunk; on exception runs
|
||||
failure hook and yields via serialize_error. Use for SSE or NDJSON.
|
||||
|
||||
``flush_tail`` runs once after the upstream iterator completes cleanly and
|
||||
its non-empty result is yielded, so a serializer that buffers bytes across
|
||||
chunks can emit anything still held at end of stream.
|
||||
"""
|
||||
verbose_proxy_logger.debug("inside generator")
|
||||
# Resolve per-stream (not per-chunk) whether the heavy per-chunk path
|
||||
|
|
@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# so it must not suppress that refund.
|
||||
delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
yield serialize_chunk(chunk)
|
||||
held_tail: Final = flush_tail() if flush_tail is not None else b""
|
||||
if held_tail:
|
||||
yield serialize_chunk(held_tail)
|
||||
stream_completed = True
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
# Client disconnected mid-stream. CancelledError / GeneratorExit
|
||||
|
|
@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# billing and release exactly once. This is the outermost generator
|
||||
# Starlette closes on disconnect, so the nested iterator hook (which
|
||||
# only sees GeneratorExit on GC) cannot own the refund.
|
||||
if not stream_completed:
|
||||
client_disconnected = True
|
||||
client_disconnected = not stream_completed
|
||||
if not delivered_chunk and not _withheld_provider_output(response):
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_budget_reservation_on_cancel,
|
||||
|
|
@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request_data: dict,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
request: Request | None = None,
|
||||
restamp_model: str | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""
|
||||
Anthropic /messages and Google /generateContent streaming data generator require SSE events.
|
||||
|
|
@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing:
|
|||
SSE serializers directly (rather than re-wrapping it in another
|
||||
``async for: yield`` trampoline), so a streamed chunk traverses one
|
||||
fewer async-generator layer / coroutine resume on the hot path.
|
||||
|
||||
``restamp_model`` publishes that name on the Anthropic ``message_start``
|
||||
event in place of the provider's model, matching what the non-streaming
|
||||
response reports.
|
||||
"""
|
||||
restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None
|
||||
return ProxyBaseLLMRequestProcessing.async_streaming_data_generator(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=request_data,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk,
|
||||
serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper),
|
||||
serialize_error=lambda proxy_exc: (
|
||||
f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n"
|
||||
),
|
||||
request=request,
|
||||
flush_tail=None if restamper is None else restamper.flush,
|
||||
)
|
||||
|
||||
@overload
|
||||
|
|
|
|||
|
|
@ -0,0 +1,288 @@
|
|||
"""
|
||||
Tests for restamping the public model on Anthropic Messages streaming chunks.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.anthropic_endpoints.streaming_model_restamp import (
|
||||
AnthropicStreamModelRestamper,
|
||||
restamp_anthropic_stream_chunk_model,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
|
||||
def _message_start_frame(model: str, line_end: str = "\n") -> bytes:
|
||||
payload = {
|
||||
"type": "message_start",
|
||||
"message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []},
|
||||
}
|
||||
return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode()
|
||||
|
||||
|
||||
def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock:
|
||||
async def _iterator_hook(**_kwargs):
|
||||
for frame in frames:
|
||||
yield frame
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook
|
||||
proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"])
|
||||
return proxy_logging_obj
|
||||
|
||||
|
||||
def _model_from_frame(frame: bytes | str) -> str:
|
||||
text = frame.decode("utf-8") if isinstance(frame, bytes) else frame
|
||||
data_line = next(line for line in text.split("\n") if line.startswith("data:"))
|
||||
return json.loads(data_line[len("data:") :])["message"]["model"]
|
||||
|
||||
|
||||
def test_restamps_sse_bytes_frame():
|
||||
restamped = restamp_anthropic_stream_chunk_model(
|
||||
_message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1"
|
||||
)
|
||||
|
||||
assert isinstance(restamped, bytes)
|
||||
assert _model_from_frame(restamped) == "claude-auto-1"
|
||||
assert b"event: message_start" in restamped
|
||||
|
||||
|
||||
def test_restamps_event_dict():
|
||||
chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}
|
||||
|
||||
restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2")
|
||||
|
||||
assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}}
|
||||
assert chunk["message"]["model"] == "claude-sonnet-4-6"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk",
|
||||
[
|
||||
b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n',
|
||||
{"type": "content_block_delta", "delta": {"text": "hi"}},
|
||||
{"type": "message_start", "message": "not-a-dict"},
|
||||
b"event: message_start\ndata: not-json\n\n",
|
||||
b"data: [DONE]\n\n",
|
||||
],
|
||||
)
|
||||
def test_leaves_chunks_without_a_model_untouched(chunk):
|
||||
assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_publishes_requested_model_on_message_start():
|
||||
"""The message_start event reports the requested model, not the provider's."""
|
||||
delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-auto-1"
|
||||
assert chunks[1] == delta_frame
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_keeps_provider_model_when_restamping_is_off():
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001"
|
||||
|
||||
|
||||
def test_restamps_message_start_split_across_transport_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
held = restamper.process(frame[:25])
|
||||
emitted = restamper.process(frame[25:])
|
||||
|
||||
assert held == b""
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_emits_coalesced_frames_with_only_message_start_rewritten():
|
||||
delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
combined = _message_start_frame("claude-haiku-4-5-20251001") + delta
|
||||
|
||||
emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined)
|
||||
|
||||
assert isinstance(restamper_output, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
assert emitted.endswith(delta)
|
||||
|
||||
|
||||
def test_ping_frames_keep_the_restamper_armed():
|
||||
ping = b'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(ping) == ping
|
||||
reassembled = restamper.process(frame[:10])
|
||||
reassembled += restamper.process(frame[10:])
|
||||
|
||||
assert _model_from_frame(reassembled) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_first_non_ping_event_disarms_the_restamper():
|
||||
delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n'
|
||||
late_message_start = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(delta) == delta
|
||||
assert restamper.process(late_message_start) == late_message_start
|
||||
|
||||
|
||||
def test_oversized_unterminated_chunk_flushes_unmodified():
|
||||
blob = b"data: " + b"x" * 70000
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(blob) == blob
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
assert restamper.process(frame) == frame
|
||||
|
||||
|
||||
def test_dict_message_start_disarms_after_restamp():
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}})
|
||||
second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}}
|
||||
|
||||
assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}}
|
||||
assert restamper.process(second) == second
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_restamps_message_start_split_across_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks)
|
||||
assert _model_from_frame(joined) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_restamps_crlf_terminated_message_start_frame():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n'
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
emitted = restamper.process(frame)
|
||||
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
assert emitted.endswith(b"\r\n\r\n")
|
||||
assert restamper.process(delta) == delta
|
||||
|
||||
|
||||
def test_restamps_cr_terminated_message_start_frame():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
emitted = restamper.process(frame)
|
||||
|
||||
assert isinstance(emitted, bytes)
|
||||
assert b'"model":"claude-auto-1"' in emitted
|
||||
assert emitted.endswith(b"\r\r")
|
||||
|
||||
|
||||
def test_restamps_crlf_message_start_split_across_transport_chunks():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
held = restamper.process(frame[:25])
|
||||
emitted = restamper.process(frame[25:])
|
||||
|
||||
assert held == b""
|
||||
assert isinstance(emitted, bytes)
|
||||
assert _model_from_frame(emitted) == "claude-auto-1"
|
||||
|
||||
|
||||
def test_flush_returns_restamped_held_tail():
|
||||
unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2]
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
|
||||
assert restamper.process(unterminated) == b""
|
||||
flushed = restamper.flush()
|
||||
|
||||
assert b'"model":"claude-auto-1"' in flushed
|
||||
assert restamper.flush() == b""
|
||||
|
||||
|
||||
def test_flush_disarms_the_restamper():
|
||||
restamper = AnthropicStreamModelRestamper("claude-auto-1")
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001")
|
||||
|
||||
assert restamper.flush() == b""
|
||||
assert restamper.process(frame) == frame
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_flushes_held_tail_at_end_of_stream():
|
||||
unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2]
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([unterminated])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks)
|
||||
assert b'"model":"claude-auto-1"' in joined
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sse_generator_restamps_crlf_stream():
|
||||
frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n")
|
||||
delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n'
|
||||
proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta])
|
||||
|
||||
chunks = [
|
||||
chunk
|
||||
async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator(
|
||||
response=MagicMock(),
|
||||
user_api_key_dict=MagicMock(),
|
||||
request_data={"model": "claude-auto-1"},
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
restamp_model="claude-auto-1",
|
||||
)
|
||||
]
|
||||
|
||||
assert _model_from_frame(chunks[0]) == "claude-auto-1"
|
||||
assert chunks[1] == delta
|
||||
Loading…
Add table
Reference in a new issue