diff --git a/litellm/proxy/guardrails/anthropic_sse.py b/litellm/proxy/guardrails/anthropic_sse.py index 3c7580a14a4..9e00e0398f1 100644 --- a/litellm/proxy/guardrails/anthropic_sse.py +++ b/litellm/proxy/guardrails/anthropic_sse.py @@ -47,11 +47,15 @@ def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: ) -def assemble_anthropic_sse_stream(all_chunks: Sequence[object]) -> ModelResponse | None: - """Assemble raw Anthropic SSE frames into a ModelResponse, restoring the upstream id and model. +def assemble_anthropic_sse_stream( + all_chunks: Sequence[object], *, restore_identity: bool = False +) -> ModelResponse | None: + """Assemble raw Anthropic SSE frames into a ModelResponse. - The assembler does not carry either through, and they are written onto the freshly built object - rather than copied into a new one because it is unreachable from caller state until returned. + ``restore_identity`` stamps the upstream message id and model onto the result, which the + assembler does not carry through. It is off by default so callers that re-emit the assembled + 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, @@ -63,7 +67,7 @@ def assemble_anthropic_sse_stream(all_chunks: Sequence[object]) -> ModelResponse message_start: Final = _anthropic_message_start(sse_stream) if message_start is None: return None - model: Final = message_start.get("model") + 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 all_chunks=(sse_stream,), @@ -74,6 +78,8 @@ def assemble_anthropic_sse_stream(all_chunks: Sequence[object]) -> ModelResponse return None if not isinstance(assembled, ModelResponse): return None + if not restore_identity: + return assembled message_id: Final = message_start.get("id") if isinstance(message_id, str): assembled.id = message_id diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index abb5a185c1e..24119a37eeb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -2593,7 +2593,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # /v1/messages arrives as SSE frames, which stream_chunk_builder cannot assemble raw_sse: Final = is_raw_sse_stream(all_chunks) assembled_model_response: ModelResponse | TextCompletionResponse | None = ( - assemble_anthropic_sse_stream(all_chunks) if raw_sse else stream_chunk_builder(chunks=all_chunks) + assemble_anthropic_sse_stream(all_chunks, restore_identity=True) + if raw_sse + else stream_chunk_builder(chunks=all_chunks) ) if isinstance(assembled_model_response, ModelResponse): pre_guardrail_text: Final = model_response_text(assembled_model_response) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 6cfd0dde2f8..60ba197317a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -1215,6 +1215,25 @@ class TestToolPermissionGuardrailAnthropicMessages: ) assert '"stop_reason": "tool_use"' not in body + @pytest.mark.asyncio + async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self): + """The SSE helpers moved to litellm/proxy/guardrails/anthropic_sse.py must be a pure move. + + That module can stamp the upstream message id and model onto the assembled response for + callers that ask for it; this path never did, and a client reads those bytes. + """ + with patch.object(self.rewriting, "should_run_guardrail", return_value=True): + out = await self._drain(self.rewriting, self._sse_chunks("Read")) + + body = b"".join(c if isinstance(c, bytes) else str(c).encode() for c in out).decode() + message_start = next( + json.loads(line[6:]) + for line in body.splitlines() + if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_start" + )["message"] + assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" + assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model" + def _resplit(self, chunks, size=7): joined = b"".join(chunks) return [joined[i : i + size] for i in range(0, len(joined), size)]