Skip to main content

bamboo_engine/runtime/stream/
handler.rs

1use tokio::sync::mpsc;
2use tokio_util::sync::CancellationToken;
3
4use bamboo_agent_core::tools::ToolCall;
5use bamboo_agent_core::{AgentError, AgentEvent};
6use bamboo_infrastructure::LLMStream;
7
8mod chunk_handling;
9mod consume;
10mod stream_state;
11
12pub struct StreamHandlingOutput {
13    pub response_id: Option<String>,
14    pub content: String,
15    pub reasoning_content: String,
16    pub token_count: usize,
17    pub tool_calls: Vec<ToolCall>,
18    pub output_tokens: u64,
19    pub thinking_tokens: u64,
20    pub cache_creation_input_tokens: u64,
21    pub cache_read_input_tokens: u64,
22}
23
24pub async fn consume_llm_stream(
25    stream: LLMStream,
26    event_tx: &mpsc::Sender<AgentEvent>,
27    cancel_token: &CancellationToken,
28    session_id: &str,
29) -> Result<StreamHandlingOutput, AgentError> {
30    consume::consume_llm_stream_internal(stream, Some(event_tx), cancel_token, session_id).await
31}
32
33pub async fn consume_llm_stream_silent(
34    stream: LLMStream,
35    cancel_token: &CancellationToken,
36    session_id: &str,
37) -> Result<StreamHandlingOutput, AgentError> {
38    consume::consume_llm_stream_internal(stream, None, cancel_token, session_id).await
39}
40
41#[cfg(test)]
42mod tests;