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