codei_llm/
stream.rs

1use std::pin::Pin;
2
3use futures::Stream;
4use serde::{Deserialize, Serialize};
5
6use crate::AssistantResponse;
7
8pub type ChatStream = Pin<Box<dyn Stream<Item = Result<StreamEvent, crate::LlmError>> + Send>>;
9
10#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
11pub struct Usage {
12    pub input_tokens: u32,
13    pub output_tokens: u32,
14}
15
16impl Usage {
17    pub fn add_assign(&mut self, other: Self) {
18        self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
19        self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
20    }
21
22    pub fn total(&self) -> u32 {
23        self.input_tokens.saturating_add(self.output_tokens)
24    }
25}
26
27#[derive(Debug, Clone)]
28pub enum StreamEvent {
29    TextDelta(String),
30    ToolCallDelta {
31        index: u32,
32        id: Option<String>,
33        name: Option<String>,
34        arguments: Option<String>,
35    },
36    Usage(Usage),
37    Done,
38}
39
40/// Collect a chat stream into a single assistant response.
41pub async fn collect_response<S>(mut stream: S) -> Result<AssistantResponse, crate::LlmError>
42where
43    S: Stream<Item = Result<StreamEvent, crate::LlmError>> + Unpin,
44{
45    use futures_util::StreamExt;
46
47    let mut response = AssistantResponse::default();
48    let mut pending_tools: std::collections::BTreeMap<
49        u32,
50        (Option<String>, Option<String>, String),
51    > = std::collections::BTreeMap::new();
52
53    while let Some(item) = stream.next().await {
54        match item? {
55            StreamEvent::TextDelta(text) => response.content.push_str(&text),
56            StreamEvent::ToolCallDelta {
57                index,
58                id,
59                name,
60                arguments,
61            } => {
62                let entry = pending_tools.entry(index).or_default();
63                if let Some(id) = id {
64                    entry.0 = Some(id);
65                }
66                if let Some(name) = name {
67                    entry.1 = Some(name);
68                }
69                if let Some(args) = arguments {
70                    entry.2.push_str(&args);
71                }
72            }
73            StreamEvent::Usage(usage) => response.usage = Some(usage),
74            StreamEvent::Done => {}
75        }
76    }
77
78    for (_, (id, name, arguments)) in pending_tools {
79        if let (Some(id), Some(name)) = (id, name) {
80            response.tool_calls.push(crate::ToolCall {
81                id,
82                name,
83                arguments,
84            });
85        }
86    }
87
88    Ok(response)
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use futures::stream;
95
96    #[tokio::test]
97    async fn collects_text_and_tool_calls() {
98        let events = vec![
99            Ok(StreamEvent::TextDelta("hi".into())),
100            Ok(StreamEvent::ToolCallDelta {
101                index: 0,
102                id: Some("call_1".into()),
103                name: Some("read".into()),
104                arguments: None,
105            }),
106            Ok(StreamEvent::ToolCallDelta {
107                index: 0,
108                id: None,
109                name: None,
110                arguments: Some(r#"{"path":"a.rs"}"#.into()),
111            }),
112            Ok(StreamEvent::Done),
113        ];
114        let response = collect_response(stream::iter(events)).await.unwrap();
115        assert_eq!(response.content, "hi");
116        assert_eq!(response.tool_calls.len(), 1);
117        assert_eq!(response.tool_calls[0].name, "read");
118    }
119}
codei_llm/stream.rs

codei_llm/
stream.rs
