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