kcode_k1_chat_chatend/
lib.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
4pub struct ToolCallId {
5    session: [u8; 12],
6    sequence: u64,
7}
8
9impl ToolCallId {
10    pub const fn new(session: [u8; 12], sequence: u64) -> Self {
11        Self { session, sequence }
12    }
13
14    pub const fn session(self) -> [u8; 12] {
15        self.session
16    }
17
18    pub const fn sequence(self) -> u64 {
19        self.sequence
20    }
21}
22
23#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
24pub struct BoxId(u64);
25
26impl BoxId {
27    pub const fn new(value: u64) -> Self {
28        Self(value)
29    }
30
31    pub const fn get(self) -> u64 {
32        self.0
33    }
34}
35
36#[derive(Clone, Debug, Eq, PartialEq)]
37pub enum BoxContent {
38    System(String),
39    User(String),
40    Kennedy {
41        text: String,
42    },
43    Attachment,
44    KtoolCall {
45        tool_call_id: ToolCallId,
46        name: String,
47        arguments: String,
48    },
49    KtoolReturn {
50        tool_call_id: ToolCallId,
51        originating_call: BoxId,
52        result: Result<String, String>,
53    },
54}
55
56#[derive(Clone, Debug, Eq, PartialEq)]
57pub struct ChatBox {
58    id: BoxId,
59    content: BoxContent,
60}
61
62impl ChatBox {
63    pub fn new(id: BoxId, content: BoxContent) -> Self {
64        Self { id, content }
65    }
66
67    pub const fn id(&self) -> BoxId {
68        self.id
69    }
70
71    pub const fn content(&self) -> &BoxContent {
72        &self.content
73    }
74}
75
76#[derive(Clone, Debug, Eq, PartialEq)]
77pub struct ProviderCall {
78    pub tool_call_id: ToolCallId,
79    pub name: String,
80    pub arguments: String,
81}
82
83#[derive(Clone, Debug, Eq, PartialEq)]
84pub struct DispatchCall {
85    pub tool_call_id: ToolCallId,
86    pub call_box_id: BoxId,
87    pub name: String,
88    pub arguments: String,
89}
90
91#[derive(Clone, Copy, Debug, Eq, PartialEq)]
92pub enum TransitionError {
93    InvalidPhase,
94    DuplicateToolCall,
95    UnknownToolCall,
96    DuplicateReturn,
97    BoxIdOverflow,
98}
99
100#[derive(Clone, Copy, Debug, Eq, PartialEq)]
101pub enum RecoveryError {
102    NonContiguousBoxId,
103    EmptyKennedy,
104    DuplicateToolCall,
105    UnknownToolCall,
106    WrongOriginatingCall,
107    DuplicateReturn,
108}
109
110#[derive(Default)]
111pub struct Chatend {
112    boxes: Vec<ChatBox>,
113    provider: bool,
114    queued: Vec<BoxContent>,
115    tool_calls: BTreeMap<ToolCallId, ToolCallRecord>,
116    last_id: u64,
117}
118
119struct ToolCallRecord {
120    call_box_id: BoxId,
121    returned: bool,
122}
123
124impl Chatend {
125    pub fn new() -> Self {
126        Self::default()
127    }
128
129    pub fn recover(boxes: Vec<ChatBox>) -> Result<Self, RecoveryError> {
130        let mut expected_id = 1_u64;
131        let mut tool_calls: BTreeMap<ToolCallId, ToolCallRecord> = BTreeMap::new();
132        for chat_box in &boxes {
133            if chat_box.id.get() != expected_id {
134                return Err(RecoveryError::NonContiguousBoxId);
135            }
136            match &chat_box.content {
137                BoxContent::Kennedy { text } if text.is_empty() => {
138                    return Err(RecoveryError::EmptyKennedy);
139                }
140                BoxContent::KtoolCall { tool_call_id, .. } => {
141                    if tool_calls.contains_key(tool_call_id) {
142                        return Err(RecoveryError::DuplicateToolCall);
143                    }
144                    tool_calls.insert(
145                        *tool_call_id,
146                        ToolCallRecord {
147                            call_box_id: chat_box.id,
148                            returned: false,
149                        },
150                    );
151                }
152                BoxContent::KtoolReturn {
153                    tool_call_id,
154                    originating_call,
155                    ..
156                } => {
157                    let record = tool_calls
158                        .get_mut(tool_call_id)
159                        .ok_or(RecoveryError::UnknownToolCall)?;
160                    if record.call_box_id != *originating_call {
161                        return Err(RecoveryError::WrongOriginatingCall);
162                    }
163                    if record.returned {
164                        return Err(RecoveryError::DuplicateReturn);
165                    }
166                    record.returned = true;
167                }
168                _ => {}
169            }
170            expected_id = expected_id
171                .checked_add(1)
172                .ok_or(RecoveryError::NonContiguousBoxId)?;
173        }
174        Ok(Self {
175            last_id: expected_id - 1,
176            boxes,
177            provider: false,
178            queued: Vec::new(),
179            tool_calls,
180        })
181    }
182
183    pub fn boxes(&self) -> &[ChatBox] {
184        &self.boxes
185    }
186
187    pub fn accept_system(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
188        self.accept_arrival(BoxContent::System(text))
189    }
190
191    pub fn accept_user(&mut self, text: String) -> Result<Option<BoxId>, TransitionError> {
192        self.accept_arrival(BoxContent::User(text))
193    }
194
195    pub fn accept_attachment(&mut self) -> Result<Option<BoxId>, TransitionError> {
196        self.accept_arrival(BoxContent::Attachment)
197    }
198
199    pub fn accept_async_return(
200        &mut self,
201        tool_call_id: ToolCallId,
202        result: Result<String, String>,
203    ) -> Result<Option<BoxId>, TransitionError> {
204        let record = self
205            .tool_calls
206            .get(&tool_call_id)
207            .ok_or(TransitionError::UnknownToolCall)?;
208        if record.returned {
209            return Err(TransitionError::DuplicateReturn);
210        }
211        let originating_call = record.call_box_id;
212        if !self.provider {
213            self.ensure_capacity(1)?;
214        }
215        self.tool_calls.get_mut(&tool_call_id).unwrap().returned = true;
216        let content = BoxContent::KtoolReturn {
217            tool_call_id,
218            originating_call,
219            result,
220        };
221        if self.provider {
222            self.queued.push(content);
223            Ok(None)
224        } else {
225            Ok(Some(self.append(content)))
226        }
227    }
228
229    pub fn start_round(&mut self) -> Result<Option<BoxId>, TransitionError> {
230        if self.provider {
231            return Err(TransitionError::InvalidPhase);
232        }
233        let frontier = self.boxes.last().map(ChatBox::id);
234        self.provider = true;
235        Ok(frontier)
236    }
237
238    pub fn append_stage(
239        &mut self,
240        text: String,
241        calls: Vec<ProviderCall>,
242    ) -> Result<Vec<DispatchCall>, TransitionError> {
243        if !self.provider {
244            return Err(TransitionError::InvalidPhase);
245        }
246        let mut seen = BTreeSet::new();
247        for call in &calls {
248            if !seen.insert(call.tool_call_id) || self.tool_calls.contains_key(&call.tool_call_id) {
249                return Err(TransitionError::DuplicateToolCall);
250            }
251        }
252        let text_growth = if text.is_empty() { 0 } else { 1 };
253        let growth = calls
254            .len()
255            .checked_add(text_growth)
256            .ok_or(TransitionError::BoxIdOverflow)?;
257        self.ensure_capacity(growth)?;
258        if !text.is_empty() {
259            self.append(BoxContent::Kennedy { text });
260        }
261        let mut dispatches = Vec::with_capacity(calls.len());
262        for call in calls {
263            let call_box_id = self.append(BoxContent::KtoolCall {
264                tool_call_id: call.tool_call_id,
265                name: call.name.clone(),
266                arguments: call.arguments.clone(),
267            });
268            self.tool_calls.insert(
269                call.tool_call_id,
270                ToolCallRecord {
271                    call_box_id,
272                    returned: false,
273                },
274            );
275            dispatches.push(DispatchCall {
276                tool_call_id: call.tool_call_id,
277                call_box_id,
278                name: call.name,
279                arguments: call.arguments,
280            });
281        }
282        Ok(dispatches)
283    }
284
285    pub fn flush_active_arrivals(&mut self) -> Result<Vec<ChatBox>, TransitionError> {
286        if !self.provider {
287            return Err(TransitionError::InvalidPhase);
288        }
289        self.ensure_capacity(self.queued.len())?;
290        let first = self.boxes.len();
291        for content in std::mem::take(&mut self.queued) {
292            self.append(content);
293        }
294        Ok(self.boxes[first..].to_vec())
295    }
296
297    pub fn done(&mut self, text: String) -> Result<(), TransitionError> {
298        if !self.provider {
299            return Err(TransitionError::InvalidPhase);
300        }
301        let text_growth = if text.is_empty() { 0 } else { 1 };
302        let growth = self
303            .queued
304            .len()
305            .checked_add(text_growth)
306            .ok_or(TransitionError::BoxIdOverflow)?;
307        self.ensure_capacity(growth)?;
308        if !text.is_empty() {
309            self.append(BoxContent::Kennedy { text });
310        }
311        self.provider = false;
312        for content in std::mem::take(&mut self.queued) {
313            self.append(content);
314        }
315        Ok(())
316    }
317
318    fn accept_arrival(&mut self, content: BoxContent) -> Result<Option<BoxId>, TransitionError> {
319        if self.provider {
320            self.queued.push(content);
321            Ok(None)
322        } else {
323            self.ensure_capacity(1)?;
324            Ok(Some(self.append(content)))
325        }
326    }
327
328    fn ensure_capacity(&self, additional: usize) -> Result<(), TransitionError> {
329        let additional = u64::try_from(additional).map_err(|_| TransitionError::BoxIdOverflow)?;
330        self.last_id
331            .checked_add(additional)
332            .ok_or(TransitionError::BoxIdOverflow)?;
333        Ok(())
334    }
335
336    fn append(&mut self, content: BoxContent) -> BoxId {
337        self.last_id += 1;
338        let id = BoxId(self.last_id);
339        self.boxes.push(ChatBox { id, content });
340        id
341    }
342}
343
344#[cfg(test)]
345mod tests;
kcode_k1_chat_chatend/lib.rs

kcode_k1_chat_chatend/
lib.rs
