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

codei_llm/
stream.rs
