1use std::collections::BTreeMap;
2
3use serde_json::Value;
4
5use crate::stream::UniversalEvent;
6
7use super::{ContentBlock, Extensions, Role, UniversalItem, UniversalResponse};
8
9impl UniversalResponse {
10 pub fn from_events(events: &[UniversalEvent]) -> Self {
11 let mut response = UniversalResponse::default();
12 let mut current_message: Option<PartialMessage> = None;
13 let mut pending_tool_calls: Vec<PartialToolCall> = Vec::new();
14 let mut pending_reasoning: BTreeMap<usize, String> = BTreeMap::new();
15
16 for event in events {
17 match event {
18 UniversalEvent::ResponseStart {
19 id,
20 model,
21 extensions,
22 } => {
23 response.id = id.clone();
24 response.model = model.clone();
25 response.extensions.extend(extensions.clone());
26 response
27 .status
28 .get_or_insert_with(|| "in_progress".to_string());
29 }
30 UniversalEvent::MessageStart {
31 id,
32 role,
33 extensions,
34 } => {
35 flush_partial_message(&mut response.output, &mut current_message);
36 flush_pending_tool_calls(&mut response.output, &mut pending_tool_calls);
37 current_message = Some(PartialMessage {
38 id: Some(id.clone()),
39 role: *role,
40 content: BTreeMap::new(),
41 extensions: extensions.clone(),
42 });
43 }
44 UniversalEvent::ContentDone { index, final_block } => {
45 if let Some(ContentBlock::ToolCall {
46 id,
47 name,
48 arguments,
49 extensions,
50 }) = final_block
51 {
52 remember_tool_call_metadata(
53 &mut pending_tool_calls,
54 id,
55 Some(name),
56 extensions,
57 );
58 fill_tool_call_arguments_if_empty(&mut pending_tool_calls, id, arguments);
59 } else if let (Some(message), Some(block)) = (&mut current_message, final_block)
60 {
61 message.content.insert(*index, block.clone());
62 }
63 }
64 UniversalEvent::TextDelta { index, text } => {
65 if let Some(message) = &mut current_message {
66 append_text_block(&mut message.content, *index, text);
67 }
68 }
69 UniversalEvent::ReasoningDelta { index, text } => {
70 if let Some(message) = &mut current_message {
71 append_reasoning_block(&mut message.content, *index, text);
72 } else {
73 pending_reasoning
74 .entry(*index)
75 .and_modify(|existing| existing.push_str(text))
76 .or_insert_with(|| text.clone());
77 }
78 }
79 UniversalEvent::ToolCallDelta {
80 id,
81 name,
82 arguments_delta,
83 } => {
84 flush_partial_message(&mut response.output, &mut current_message);
85 append_tool_call_delta(
86 &mut pending_tool_calls,
87 id,
88 name.as_deref(),
89 arguments_delta,
90 );
91 }
92 UniversalEvent::MessageDone {
93 finish_reason,
94 usage,
95 extensions: _,
96 } => {
97 response.finish_reason = *finish_reason;
98 if usage.is_some() {
99 response.usage = usage.clone();
100 }
101 flush_partial_message(&mut response.output, &mut current_message);
102 flush_pending_tool_calls(&mut response.output, &mut pending_tool_calls);
103 }
104 UniversalEvent::ResponseDone { usage, extensions } => {
105 if usage.is_some() {
106 response.usage = usage.clone();
107 }
108 flush_pending_tool_calls(&mut response.output, &mut pending_tool_calls);
109 response.status = Some("completed".to_string());
110 response.extensions.extend(extensions.clone());
111 }
112 UniversalEvent::Error { message, raw } => {
113 response.status = Some("failed".to_string());
114 response.extensions.insert(
115 "error".to_string(),
116 raw.clone()
117 .unwrap_or_else(|| Value::String(message.clone())),
118 );
119 }
120 UniversalEvent::Unknown { raw, .. } => {
121 flush_partial_message(&mut response.output, &mut current_message);
122 flush_pending_tool_calls(&mut response.output, &mut pending_tool_calls);
123 response
124 .output
125 .push(UniversalItem::Unknown { raw: raw.clone() });
126 }
127 UniversalEvent::ContentStart { block, .. } => {
128 if let ContentBlock::ToolCall {
129 id,
130 name,
131 extensions,
132 ..
133 } = block
134 {
135 remember_tool_call_metadata(
136 &mut pending_tool_calls,
137 id,
138 Some(name),
139 extensions,
140 );
141 }
142 }
143 }
144 }
145
146 flush_partial_message(&mut response.output, &mut current_message);
147 flush_pending_tool_calls(&mut response.output, &mut pending_tool_calls);
148 for (index, text) in pending_reasoning {
149 response.output.insert(
150 index.min(response.output.len()),
151 UniversalItem::Reasoning {
152 id: None,
153 text: Some(text),
154 encrypted: None,
155 extensions: Extensions::new(),
156 },
157 );
158 }
159 response
160 }
161
162 pub fn to_events(&self) -> Vec<UniversalEvent> {
163 let mut events = Vec::new();
164 events.push(UniversalEvent::ResponseStart {
165 id: self.id.clone(),
166 model: self.model.clone(),
167 extensions: self.extensions.clone(),
168 });
169 for (item_index, item) in self.output.iter().enumerate() {
170 match item {
171 UniversalItem::Message {
172 role,
173 id,
174 content,
175 extensions,
176 } => {
177 events.push(UniversalEvent::MessageStart {
178 id: id
179 .clone()
180 .unwrap_or_else(|| format!("message_{item_index}")),
181 role: *role,
182 extensions: extensions.clone(),
183 });
184 for (content_index, block) in content.iter().cloned().enumerate() {
185 events.push(UniversalEvent::ContentStart {
186 index: content_index,
187 block: block.clone(),
188 });
189 if let ContentBlock::Text { text } = &block {
190 events.push(UniversalEvent::TextDelta {
191 index: content_index,
192 text: text.clone(),
193 });
194 }
195 if let ContentBlock::Reasoning {
196 text: Some(text), ..
197 } = &block
198 {
199 events.push(UniversalEvent::ReasoningDelta {
200 index: content_index,
201 text: text.clone(),
202 });
203 }
204 events.push(UniversalEvent::ContentDone {
205 index: content_index,
206 final_block: Some(block),
207 });
208 }
209 events.push(UniversalEvent::MessageDone {
210 finish_reason: self.finish_reason,
211 usage: self.usage.clone(),
212 extensions: Extensions::new(),
213 });
214 }
215 UniversalItem::ToolCall {
216 id,
217 name,
218 arguments,
219 ..
220 } => events.push(UniversalEvent::ToolCallDelta {
221 id: id.clone(),
222 name: Some(name.clone()),
223 arguments_delta: match arguments {
224 Value::String(value) => value.clone(),
225 value => serde_json::to_string(value).unwrap_or_default(),
226 },
227 }),
228 UniversalItem::Reasoning { text, .. } => {
229 if let Some(text) = text {
230 events.push(UniversalEvent::ReasoningDelta {
231 index: item_index,
232 text: text.clone(),
233 });
234 }
235 }
236 UniversalItem::ToolResult { .. } | UniversalItem::Unknown { .. } => {}
237 }
238 }
239 events.push(UniversalEvent::ResponseDone {
240 usage: self.usage.clone(),
241 extensions: Extensions::new(),
242 });
243 events
244 }
245}
246
247struct PartialMessage {
248 id: Option<String>,
249 role: Role,
250 content: BTreeMap<usize, ContentBlock>,
251 extensions: Extensions,
252}
253
254struct PartialToolCall {
255 id: String,
256 name: Option<String>,
257 arguments: String,
258 saw_delta: bool,
259 extensions: Extensions,
260}
261
262fn flush_partial_message(output: &mut Vec<UniversalItem>, message: &mut Option<PartialMessage>) {
263 let Some(message) = message.take() else {
264 return;
265 };
266 output.push(UniversalItem::Message {
267 role: message.role,
268 id: message.id,
269 content: message.content.into_values().collect(),
270 extensions: message.extensions,
271 });
272}
273
274fn append_tool_call_delta(
275 pending_tool_calls: &mut Vec<PartialToolCall>,
276 id: &str,
277 name: Option<&str>,
278 arguments_delta: &str,
279) {
280 remember_tool_call_metadata(pending_tool_calls, id, name, &Extensions::new());
281 let Some(tool_call) = pending_tool_calls
282 .iter_mut()
283 .find(|tool_call| tool_call.id == id)
284 else {
285 return;
286 };
287
288 tool_call.arguments.push_str(arguments_delta);
289 tool_call.saw_delta = true;
290}
291
292fn remember_tool_call_metadata(
293 pending_tool_calls: &mut Vec<PartialToolCall>,
294 id: &str,
295 name: Option<&str>,
296 extensions: &Extensions,
297) {
298 let Some(tool_call) = pending_tool_calls
299 .iter_mut()
300 .find(|tool_call| tool_call.id == id)
301 else {
302 pending_tool_calls.push(PartialToolCall {
303 id: id.to_string(),
304 name: name
305 .filter(|name| !name.is_empty())
306 .map(ToString::to_string),
307 arguments: String::new(),
308 saw_delta: false,
309 extensions: extensions.clone(),
310 });
311 return;
312 };
313
314 if let Some(name) = name.filter(|name| !name.is_empty()) {
315 tool_call.name = Some(name.to_string());
316 }
317 tool_call.extensions.extend(extensions.clone());
318}
319
320fn fill_tool_call_arguments_if_empty(
321 pending_tool_calls: &mut Vec<PartialToolCall>,
322 id: &str,
323 arguments: &Value,
324) {
325 let Some(tool_call) = pending_tool_calls
326 .iter_mut()
327 .find(|tool_call| tool_call.id == id)
328 else {
329 return;
330 };
331 if tool_call.saw_delta || !tool_call.arguments.is_empty() {
332 return;
333 }
334 tool_call.arguments = stringify_tool_arguments(arguments);
335}
336
337fn flush_pending_tool_calls(
338 output: &mut Vec<UniversalItem>,
339 pending_tool_calls: &mut Vec<PartialToolCall>,
340) {
341 for tool_call in std::mem::take(pending_tool_calls) {
342 output.push(UniversalItem::ToolCall {
343 id: tool_call.id,
344 name: tool_call.name.unwrap_or_default(),
345 arguments: tool_call
346 .arguments
347 .parse::<Value>()
348 .unwrap_or_else(|_| Value::String(tool_call.arguments)),
349 extensions: tool_call.extensions,
350 });
351 }
352}
353
354fn stringify_tool_arguments(arguments: &Value) -> String {
355 match arguments {
356 Value::String(value) => value.clone(),
357 Value::Null => String::new(),
358 value => serde_json::to_string(value).unwrap_or_default(),
359 }
360}
361
362fn append_text_block(content: &mut BTreeMap<usize, ContentBlock>, index: usize, text: &str) {
363 match content.get_mut(&index) {
364 Some(ContentBlock::Text { text: existing }) => existing.push_str(text),
365 Some(_) => {}
366 None => {
367 content.insert(
368 index,
369 ContentBlock::Text {
370 text: text.to_string(),
371 },
372 );
373 }
374 }
375}
376
377fn append_reasoning_block(content: &mut BTreeMap<usize, ContentBlock>, index: usize, text: &str) {
378 match content.get_mut(&index) {
379 Some(ContentBlock::Reasoning {
380 text: Some(existing),
381 ..
382 }) => existing.push_str(text),
383 Some(_) => {}
384 None => {
385 content.insert(
386 index,
387 ContentBlock::Reasoning {
388 text: Some(text.to_string()),
389 encrypted: None,
390 extensions: Extensions::new(),
391 },
392 );
393 }
394 }
395}
396#[cfg(test)]
397mod tests {
398 use serde_json::json;
399
400 use crate::{UniversalEvent, UniversalItem, UniversalResponse};
401
402 #[test]
403 fn aggregates_split_tool_call_deltas_by_id() {
404 let response = UniversalResponse::from_events(&[
405 UniversalEvent::ToolCallDelta {
406 id: "call_pwd".to_string(),
407 name: Some("exec_command".to_string()),
408 arguments_delta: "{\"cmd\"".to_string(),
409 },
410 UniversalEvent::ToolCallDelta {
411 id: "call_pwd".to_string(),
412 name: None,
413 arguments_delta: ":\"pwd\"}".to_string(),
414 },
415 ]);
416
417 assert_eq!(response.output.len(), 1);
418 assert!(matches!(
419 &response.output[0],
420 UniversalItem::ToolCall {
421 id,
422 name,
423 arguments,
424 ..
425 } if id == "call_pwd" && name == "exec_command" && arguments == &json!({ "cmd": "pwd" })
426 ));
427 }
428}