1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::ProviderStreamEvent;
7
8#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
9pub struct ToolCallFragment {
10 pub choice_index: u64,
11 pub tool_index: u64,
12 #[serde(default, skip_serializing_if = "Option::is_none")]
13 pub id: Option<String>,
14 #[serde(default, skip_serializing_if = "Option::is_none")]
15 pub name: Option<String>,
16 #[serde(default)]
17 pub arguments: String,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct AssembledToolCall {
22 pub choice_index: u64,
23 pub tool_index: u64,
24 pub id: String,
25 pub name: String,
26 pub arguments: String,
27}
28
29impl AssembledToolCall {
30 pub fn arguments_json(&self) -> Result<Value, serde_json::Error> {
31 serde_json::from_str(&self.arguments)
32 }
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
36pub enum ToolCallAssemblyError {
37 #[error("tool call {choice_index}/{tool_index} changed {field} from '{first}' to '{next}'")]
38 ConflictingMetadata {
39 choice_index: u64,
40 tool_index: u64,
41 field: &'static str,
42 first: String,
43 next: String,
44 },
45 #[error("tool call {choice_index}/{tool_index} is missing {field}")]
46 MissingMetadata {
47 choice_index: u64,
48 tool_index: u64,
49 field: &'static str,
50 },
51}
52
53#[derive(Debug, Clone, Default)]
54struct PendingToolCall {
55 id: Option<String>,
56 name: Option<String>,
57 arguments: String,
58}
59
60#[derive(Debug, Clone, Default)]
61pub struct ToolCallAssembler {
62 pending: BTreeMap<(u64, u64), PendingToolCall>,
63}
64
65impl ToolCallAssembler {
66 pub fn new() -> Self {
67 Self::default()
68 }
69
70 pub fn push_fragment(
71 &mut self,
72 fragment: ToolCallFragment,
73 ) -> Result<(), ToolCallAssemblyError> {
74 let key = (fragment.choice_index, fragment.tool_index);
75 let pending = self.pending.entry(key).or_default();
76 merge_metadata(&mut pending.id, fragment.id, key, "id")?;
77 merge_metadata(&mut pending.name, fragment.name, key, "name")?;
78 pending.arguments.push_str(&fragment.arguments);
79 Ok(())
80 }
81
82 pub fn push_event(
83 &mut self,
84 event: &ProviderStreamEvent,
85 ) -> Result<usize, ToolCallAssemblyError> {
86 let ProviderStreamEvent::Data { data, .. } = event else {
87 return Ok(0);
88 };
89 let mut count = 0;
90 for choice in data
91 .get("choices")
92 .and_then(Value::as_array)
93 .into_iter()
94 .flatten()
95 {
96 let choice_index = choice.get("index").and_then(Value::as_u64).unwrap_or(0);
97 for call in choice
98 .pointer("/delta/tool_calls")
99 .and_then(Value::as_array)
100 .into_iter()
101 .flatten()
102 {
103 let fragment = ToolCallFragment {
104 choice_index,
105 tool_index: call.get("index").and_then(Value::as_u64).unwrap_or(0),
106 id: call.get("id").and_then(Value::as_str).map(str::to_owned),
107 name: call
108 .pointer("/function/name")
109 .and_then(Value::as_str)
110 .map(str::to_owned),
111 arguments: call
112 .pointer("/function/arguments")
113 .and_then(Value::as_str)
114 .unwrap_or_default()
115 .to_owned(),
116 };
117 self.push_fragment(fragment)?;
118 count += 1;
119 }
120 }
121 Ok(count)
122 }
123
124 pub fn finish(self) -> Result<Vec<AssembledToolCall>, ToolCallAssemblyError> {
125 self.pending
126 .into_iter()
127 .map(|((choice_index, tool_index), pending)| {
128 let id = pending.id.ok_or(ToolCallAssemblyError::MissingMetadata {
129 choice_index,
130 tool_index,
131 field: "id",
132 })?;
133 let name = pending.name.ok_or(ToolCallAssemblyError::MissingMetadata {
134 choice_index,
135 tool_index,
136 field: "name",
137 })?;
138 Ok(AssembledToolCall {
139 choice_index,
140 tool_index,
141 id,
142 name,
143 arguments: pending.arguments,
144 })
145 })
146 .collect()
147 }
148}
149
150fn merge_metadata(
151 current: &mut Option<String>,
152 next: Option<String>,
153 (choice_index, tool_index): (u64, u64),
154 field: &'static str,
155) -> Result<(), ToolCallAssemblyError> {
156 let Some(next) = next.filter(|value| !value.is_empty()) else {
157 return Ok(());
158 };
159 if let Some(first) = current {
160 if first != &next {
161 return Err(ToolCallAssemblyError::ConflictingMetadata {
162 choice_index,
163 tool_index,
164 field,
165 first: first.clone(),
166 next,
167 });
168 }
169 } else {
170 *current = Some(next);
171 }
172 Ok(())
173}
174
175#[cfg(test)]
176mod tests {
177 use serde_json::json;
178
179 use super::*;
180
181 #[test]
182 fn assembles_interleaved_calls_in_index_order() {
183 let mut assembler = ToolCallAssembler::new();
184 for data in [
185 json!({"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"b","type":"function","function":{"name":"second","arguments":""}}]}}]}),
186 json!({"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"a","type":"function","function":{"name":"first","arguments":"{\"x\":"}}]}}]}),
187 json!({"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"{}"}},{"index":0,"function":{"arguments":"1}"}}]}}]}),
188 ] {
189 assembler
190 .push_event(&ProviderStreamEvent::Data { event: None, data })
191 .unwrap();
192 }
193 let calls = assembler.finish().unwrap();
194 assert_eq!(
195 calls.iter().map(|call| call.tool_index).collect::<Vec<_>>(),
196 vec![0, 1]
197 );
198 assert_eq!(calls[0].arguments_json().unwrap(), json!({"x": 1}));
199 assert_eq!(calls[1].arguments_json().unwrap(), json!({}));
200 }
201}