1use aws_sdk_bedrockruntime::error::SdkError;
2use aws_sdk_bedrockruntime::primitives::event_stream::EventReceiver;
3use aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError;
4use aws_sdk_bedrockruntime::types::{
5 ContentBlockDelta, ContentBlockStart, ConverseStreamOutput, ReasoningContentBlockDelta,
6 StopReason as BedrockStopReason, TokenUsage as BedrockTokenUsage,
7};
8use aws_smithy_types::event_stream::RawMessage;
9use futures::{Stream, StreamExt, stream};
10use tracing::{error, warn};
11
12use crate::providers::stream_assembler::{StreamAssembler, assemble};
13use crate::{LlmError, LlmResponse, ProviderError, StopReason, TokenUsage, Tokens};
14
15pub fn process_bedrock_stream(
16 events: impl Stream<Item = crate::Result<ConverseStreamOutput>> + Send,
17) -> impl Stream<Item = crate::Result<LlmResponse>> + Send {
18 stream::iter([Ok(LlmResponse::Start)]).chain(assemble(events, |event, turn| Ok(decode_event(event, turn))))
19}
20
21pub(crate) fn converse_events(
22 receiver: EventReceiver<ConverseStreamOutput, ConverseStreamOutputError>,
23) -> impl Stream<Item = crate::Result<ConverseStreamOutput>> + Send {
24 stream::unfold(receiver, |mut receiver| async move {
25 let event = receiver.recv().await.map_err(|e| {
26 error!("Bedrock stream recv error: {e}");
27 LlmError::from(e)
28 });
29 event.transpose().map(|event| (event, receiver))
30 })
31}
32
33impl From<&BedrockTokenUsage> for TokenUsage {
34 fn from(usage: &BedrockTokenUsage) -> Self {
35 let cache_read = usage.cache_read_input_tokens().and_then(|v| u32::try_from(v).ok()).map(Tokens::from);
36 let cache_creation = usage.cache_write_input_tokens().and_then(|v| u32::try_from(v).ok()).map(Tokens::from);
37 let cached = cache_read.unwrap_or_default() + cache_creation.unwrap_or_default();
39 TokenUsage {
40 input_tokens: Tokens::from(u32::try_from(usage.input_tokens).unwrap_or(0)) + cached,
41 output_tokens: u32::try_from(usage.output_tokens).unwrap_or(0).into(),
42 cache_read_tokens: cache_read,
43 cache_creation_tokens: cache_creation,
44 ..TokenUsage::default()
45 }
46 }
47}
48
49impl From<SdkError<ConverseStreamOutputError, RawMessage>> for LlmError {
50 fn from(e: SdkError<ConverseStreamOutputError, RawMessage>) -> Self {
51 let message = format!("Bedrock stream error: {e}");
52 let provider = match e {
53 SdkError::ServiceError(svc) => {
54 let inner = svc.err();
55 if inner.is_throttling_exception() {
56 ProviderError::rate_limit(message)
57 } else if inner.is_service_unavailable_exception()
58 || inner.is_internal_server_exception()
59 || inner.is_model_stream_error_exception()
60 {
61 ProviderError::stream_interrupted(message)
62 } else {
63 ProviderError::api(message)
64 }
65 }
66 _ => ProviderError::stream_interrupted(message),
67 };
68 Self::from(provider)
69 }
70}
71
72fn decode_event(event: ConverseStreamOutput, turn: &mut StreamAssembler<i32>) -> Vec<LlmResponse> {
73 let response = match event {
74 ConverseStreamOutput::ContentBlockStart(event) => match event.start {
75 Some(ContentBlockStart::ToolUse(tool)) => {
76 Some(turn.start_tool(event.content_block_index, tool.tool_use_id, tool.name))
77 }
78 _ => None,
79 },
80 ConverseStreamOutput::ContentBlockDelta(event) => match event.delta {
81 Some(ContentBlockDelta::Text(text)) if !text.is_empty() => Some(LlmResponse::Text { chunk: text }),
82 Some(ContentBlockDelta::ToolUse(delta)) => turn.append_tool_args(&event.content_block_index, delta.input),
83 Some(ContentBlockDelta::ReasoningContent(ReasoningContentBlockDelta::Text(text))) if !text.is_empty() => {
84 Some(LlmResponse::Reasoning { chunk: text })
85 }
86 _ => None,
87 },
88 ConverseStreamOutput::ContentBlockStop(event) => turn.complete_tool(&event.content_block_index),
89 ConverseStreamOutput::MessageStop(event) => {
90 turn.stop(map_bedrock_stop_reason(&event.stop_reason));
91 turn.allow_eof();
92 None
93 }
94 ConverseStreamOutput::Metadata(event) => {
95 event.usage.as_ref().map(|usage| LlmResponse::Usage { tokens: usage.into() })
96 }
97 ConverseStreamOutput::MessageStart(_) => None,
98 other => {
99 warn!("Unhandled Bedrock stream event: {other:?}");
100 None
101 }
102 };
103
104 response.into_iter().collect()
105}
106
107fn map_bedrock_stop_reason(reason: &BedrockStopReason) -> StopReason {
108 match reason {
109 BedrockStopReason::EndTurn | BedrockStopReason::StopSequence => StopReason::EndTurn,
110 BedrockStopReason::ToolUse => StopReason::ToolCalls,
111 BedrockStopReason::MaxTokens | BedrockStopReason::ModelContextWindowExceeded => StopReason::Length,
112 BedrockStopReason::ContentFiltered | BedrockStopReason::GuardrailIntervened => StopReason::ContentFilter,
113 other => StopReason::Unknown(format!("{other:?}")),
114 }
115}
116
117#[cfg(test)]
118mod tests {
119 use super::*;
120 use crate::ProviderErrorKind;
121 use crate::testing::llm_response;
122 use aws_sdk_bedrockruntime::types::{
123 ContentBlockDeltaEvent, ContentBlockStartEvent, ContentBlockStopEvent, ConversationRole,
124 ConverseStreamMetadataEvent, MessageStartEvent, MessageStopEvent, ToolUseBlockDelta, ToolUseBlockStart,
125 };
126
127 #[tokio::test]
128 async fn test_text_stream() {
129 let responses = collect_responses(
130 bedrock_stream().text(0, &["Hello", "", " world"]).message_stop(BedrockStopReason::EndTurn).build(),
131 )
132 .await;
133
134 assert_eq!(responses, llm_response().text(&["Hello", " world"]).build_with_stop_reason(StopReason::EndTurn));
135 }
136
137 #[tokio::test]
138 async fn test_reasoning_stream() {
139 let responses = collect_responses(
140 bedrock_stream()
141 .reasoning(0, &["thinking"])
142 .text(1, &["answer"])
143 .message_stop(BedrockStopReason::EndTurn)
144 .build(),
145 )
146 .await;
147
148 assert_eq!(
149 responses,
150 llm_response().reasoning(&["thinking"]).text(&["answer"]).build_with_stop_reason(StopReason::EndTurn)
151 );
152 }
153
154 #[tokio::test]
155 async fn test_tool_call_stream() {
156 let deltas = [r#"{"query":"#, r#""test"}"#];
157
158 let responses = collect_responses(
159 bedrock_stream()
160 .tool_call(0, "tool_123", "search", &deltas)
161 .message_stop(BedrockStopReason::ToolUse)
162 .build(),
163 )
164 .await;
165
166 assert_eq!(
167 responses,
168 llm_response().tool_call("tool_123", "search", &deltas).build_with_stop_reason(StopReason::ToolCalls)
169 );
170 }
171
172 #[tokio::test]
173 async fn stream_closed_mid_tool_call_does_not_complete_it() {
174 let responses =
175 process_events(bedrock_stream().tool_start(0, "tool_123", "search").tool_delta(0, r#"{"query":"#).build())
176 .await;
177
178 assert!(
179 matches!(
180 responses.as_slice(),
181 [
182 Ok(LlmResponse::Start),
183 Ok(LlmResponse::ToolRequestStart { .. }),
184 Ok(LlmResponse::ToolRequestArg { .. }),
185 Err(error)
186 ] if error.provider().map(|provider| provider.kind) == Some(ProviderErrorKind::StreamInterrupted)
187 ),
188 "{responses:?}"
189 );
190 }
191
192 #[tokio::test]
193 async fn test_metadata_after_message_stop_reports_cache_usage() {
194 let usage = BedrockTokenUsage::builder()
195 .input_tokens(100)
196 .output_tokens(50)
197 .total_tokens(150)
198 .cache_read_input_tokens(40)
199 .cache_write_input_tokens(20)
200 .build()
201 .unwrap();
202
203 let responses =
204 collect_responses(bedrock_stream().message_stop(BedrockStopReason::EndTurn).metadata(usage).build()).await;
205
206 let usage = responses.iter().find_map(|response| match response {
207 LlmResponse::Usage { tokens } => Some(*tokens),
208 _ => None,
209 });
210 assert_eq!(
211 usage,
212 Some(TokenUsage {
213 input_tokens: 160.into(),
214 output_tokens: 50.into(),
215 cache_read_tokens: Some(40.into()),
216 cache_creation_tokens: Some(20.into()),
217 ..TokenUsage::default()
218 }),
219 "cached tokens count toward the prompt"
220 );
221 assert_eq!(responses.last(), Some(&LlmResponse::done_with_stop_reason(StopReason::EndTurn)));
222 }
223
224 #[tokio::test]
225 async fn test_metadata_without_cache_fields() {
226 let usage = BedrockTokenUsage::builder().input_tokens(10).output_tokens(5).total_tokens(15).build().unwrap();
227
228 let responses =
229 collect_responses(bedrock_stream().message_stop(BedrockStopReason::EndTurn).metadata(usage).build()).await;
230
231 assert_eq!(responses, llm_response().usage(10, 5).build_with_stop_reason(StopReason::EndTurn));
232 }
233
234 #[tokio::test]
235 async fn test_stop_reasons_map_to_llm_stop_reasons() {
236 for (bedrock_stop_reason, stop_reason) in [
237 (BedrockStopReason::EndTurn, StopReason::EndTurn),
238 (BedrockStopReason::StopSequence, StopReason::EndTurn),
239 (BedrockStopReason::ToolUse, StopReason::ToolCalls),
240 (BedrockStopReason::MaxTokens, StopReason::Length),
241 (BedrockStopReason::ModelContextWindowExceeded, StopReason::Length),
242 (BedrockStopReason::ContentFiltered, StopReason::ContentFilter),
243 (BedrockStopReason::GuardrailIntervened, StopReason::ContentFilter),
244 ] {
245 let responses = collect_responses(bedrock_stream().message_stop(bedrock_stop_reason).build()).await;
246
247 assert_eq!(responses, llm_response().build_with_stop_reason(stop_reason));
248 }
249 }
250
251 async fn collect_responses(events: Vec<ConverseStreamOutput>) -> Vec<LlmResponse> {
252 process_events(events).await.into_iter().map(Result::unwrap).collect()
253 }
254
255 async fn process_events(events: Vec<ConverseStreamOutput>) -> Vec<crate::Result<LlmResponse>> {
256 process_bedrock_stream(stream::iter(events.into_iter().map(Ok))).collect().await
257 }
258
259 fn bedrock_stream() -> BedrockStreamBuilder {
260 BedrockStreamBuilder::default().push(ConverseStreamOutput::MessageStart(
261 MessageStartEvent::builder().role(ConversationRole::Assistant).build().unwrap(),
262 ))
263 }
264
265 #[derive(Default)]
266 struct BedrockStreamBuilder {
267 events: Vec<ConverseStreamOutput>,
268 }
269
270 impl BedrockStreamBuilder {
271 fn text(self, index: i32, chunks: &[&str]) -> Self {
272 chunks
273 .iter()
274 .fold(self, |builder, chunk| builder.delta(index, ContentBlockDelta::Text((*chunk).to_string())))
275 .block_stop(index)
276 }
277
278 fn reasoning(self, index: i32, chunks: &[&str]) -> Self {
279 chunks
280 .iter()
281 .fold(self, |builder, chunk| {
282 builder.delta(
283 index,
284 ContentBlockDelta::ReasoningContent(ReasoningContentBlockDelta::Text((*chunk).to_string())),
285 )
286 })
287 .block_stop(index)
288 }
289
290 fn tool_call(self, index: i32, id: &str, name: &str, argument_deltas: &[&str]) -> Self {
291 argument_deltas
292 .iter()
293 .fold(self.tool_start(index, id, name), |builder, delta| builder.tool_delta(index, delta))
294 .block_stop(index)
295 }
296
297 fn tool_start(self, index: i32, id: &str, name: &str) -> Self {
298 let tool = ToolUseBlockStart::builder().tool_use_id(id).name(name).build().unwrap();
299 self.push(ConverseStreamOutput::ContentBlockStart(
300 ContentBlockStartEvent::builder()
301 .content_block_index(index)
302 .start(ContentBlockStart::ToolUse(tool))
303 .build()
304 .unwrap(),
305 ))
306 }
307
308 fn tool_delta(self, index: i32, input: &str) -> Self {
309 self.delta(index, ContentBlockDelta::ToolUse(ToolUseBlockDelta::builder().input(input).build().unwrap()))
310 }
311
312 fn message_stop(self, stop_reason: BedrockStopReason) -> Self {
313 self.push(ConverseStreamOutput::MessageStop(
314 MessageStopEvent::builder().stop_reason(stop_reason).build().unwrap(),
315 ))
316 }
317
318 fn metadata(self, usage: BedrockTokenUsage) -> Self {
319 self.push(ConverseStreamOutput::Metadata(ConverseStreamMetadataEvent::builder().usage(usage).build()))
320 }
321
322 fn delta(self, index: i32, delta: ContentBlockDelta) -> Self {
323 self.push(ConverseStreamOutput::ContentBlockDelta(
324 ContentBlockDeltaEvent::builder().content_block_index(index).delta(delta).build().unwrap(),
325 ))
326 }
327
328 fn block_stop(self, index: i32) -> Self {
329 self.push(ConverseStreamOutput::ContentBlockStop(
330 ContentBlockStopEvent::builder().content_block_index(index).build().unwrap(),
331 ))
332 }
333
334 fn push(mut self, event: ConverseStreamOutput) -> Self {
335 self.events.push(event);
336 self
337 }
338
339 fn build(self) -> Vec<ConverseStreamOutput> {
340 self.events
341 }
342 }
343}