Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/anthropic_sse.py: 20%

1"""Anthropic SSE <-> ModelResponse conversion for guardrail streaming hooks. 

2 

3`/v1/messages` streams reach a guardrail's `async_post_call_streaming_iterator_hook` as raw SSE 

4frames rather than chunk objects, which `stream_chunk_builder` cannot assemble. These helpers let a 

5hook scan such a stream, and re-emit it when the guardrail rewrote the response. 

6""" 

7 

8from __future__ import annotations 

9 

10import codecs 

11import json 

12from collections.abc import Mapping, Sequence 

13from typing import Final 

14 

15from litellm.types.utils import Choices, ModelResponse 

16 

17_ANTHROPIC_EVENT_TYPES: Final = frozenset( 

18 { 

19 "message_start", 

20 "message_delta", 

21 "message_stop", 

22 "content_block_start", 

23 "content_block_delta", 

24 "content_block_stop", 

25 "ping", 

26 "error", 

27 } 

28) 

29 

30 

31def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool: 

32 return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks) 

33 

34 

35def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None: 

36 raw: Final = b"".join( 

37 chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") 

38 for chunk in all_chunks 

39 if isinstance(chunk, (str, bytes)) 

40 ) 

41 try: 

42 return codecs.getincrementaldecoder("utf-8")().decode(raw, final=False) 

43 except UnicodeDecodeError: 

44 return None 

45 

46 

47def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]: 

48 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( 

49 AnthropicPassthroughLoggingHandler, 

50 ) 

51 

52 return tuple( 

53 event_data 

54 for event in AnthropicPassthroughLoggingHandler._split_sse_chunk_into_events(sse_stream) # pyright: ignore[reportPrivateUsage] # same parser the assembler uses 

55 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 

56 ) 

57 

58 

59def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None: 

60 return next( 

61 ( 

62 message 

63 for event_data in _parsed_sse_events(sse_stream) 

64 if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict) 

65 ), 

66 None, 

67 ) 

68 

69 

70def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool: 

71 """Whether raw SSE frames are Anthropic Messages events. 

72 

73 ``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the 

74 only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own 

75 stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format 

76 its client cannot parse, so the surface is decided on the event types actually present. 

77 """ 

78 sse_stream: Final = _joined_sse_stream(all_chunks) 

79 if sse_stream is None: 

80 return False 

81 return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream)) 

82 

83 

84def assemble_anthropic_sse_stream( 

85 all_chunks: Sequence[object], *, restore_identity: bool = False 

86) -> ModelResponse | None: 

87 """Assemble raw Anthropic SSE frames into a ModelResponse. 

88 

89 ``restore_identity`` stamps the upstream message id and model onto the result, which the 

90 assembler does not carry through. It is off by default so callers that re-emit the assembled 

91 response keep the wire shape they had before this helper was shared. The writes land on a 

92 freshly built object that is unreachable from caller state until returned. 

93 """ 

94 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( 

95 AnthropicPassthroughLoggingHandler, 

96 ) 

97 

98 sse_stream: Final = _joined_sse_stream(all_chunks) 

99 if sse_stream is None: 

100 return None 

101 message_start: Final = _anthropic_message_start(sse_stream) 

102 if message_start is None: 

103 return None 

104 model: Final = message_start.get("model") if restore_identity else None 

105 try: 

106 assembled: Final = AnthropicPassthroughLoggingHandler._build_complete_streaming_response( # pyright: ignore[reportPrivateUsage] # the only SSE-to-ModelResponse assembler; reimplementing it here would fork the parser 

107 all_chunks=(sse_stream,), 

108 litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None 

109 model=model if isinstance(model, str) else "", 

110 ) 

111 except Exception: # noqa: BLE001 # stream_chunk_builder re-raises every assembly failure as litellm.APIError 

112 return None 

113 if not isinstance(assembled, ModelResponse): 

114 return None 

115 if not restore_identity: 

116 return assembled 

117 message_id: Final = message_start.get("id") 

118 if isinstance(message_id, str): 

119 assembled.id = message_id 

120 if isinstance(model, str) and model: 

121 assembled.model = model 

122 return assembled 

123 

124 

125def model_response_text(response: ModelResponse) -> str: 

126 """Assistant text of a response, used to detect whether a guardrail rewrote it.""" 

127 return "".join( 

128 choice.message.content 

129 for choice in response.choices 

130 if isinstance(choice, Choices) # pyright: ignore[reportUnnecessaryIsInstance] # runtime choices can be StreamingChoices 

131 and isinstance(choice.message.content, str) 

132 ) 

133 

134 

135def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]: 

136 """Anthropic error event, for a failure discovered after the response headers were flushed. 

137 

138 Once a keepalive ping has been sent a raise cannot reach the client, so the failure has to 

139 travel as a frame. 

140 """ 

141 body: Final = json.dumps(message) 

142 return ( 

143 f'event: error\ndata: {{"type": "error", "error": {{"type": "guardrail_error", ' 

144 f'"message": {body}}}}}\n\n'.encode(), 

145 ) 

146 

147 

148def is_sse_error_stream(all_chunks: Sequence[object]) -> bool: 

149 """Whether the buffered stream carries nothing but error frames. 

150 

151 post_call guardrails run in a chain, so a hook can be handed the terminal error frames an 

152 earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing 

153 them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the 

154 Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload. 

155 """ 

156 if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks): 

157 # A stream mixing typed chunks with an error frame still carries content to scan, and the 

158 # frames-only join below would drop exactly the part that has to be scanned 

159 return False 

160 sse_stream: Final = _joined_sse_stream(all_chunks) 

161 if sse_stream is None: 

162 return False 

163 events: Final = _parsed_sse_events(sse_stream) 

164 return len(events) > 0 and all( 

165 event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events 

166 ) 

167 

168 

169def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]: 

170 from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( 

171 LiteLLMAnthropicMessagesAdapter, 

172 ) 

173 from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( 

174 FakeAnthropicMessagesStreamIterator, 

175 ) 

176 

177 anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( 

178 response=assembled 

179 ) 

180 return tuple(FakeAnthropicMessagesStreamIterator(response=anthropic_response).chunks)