lc_agents/streaming/
tool_call_stream.rs

1//! StreamingFunctionCallingAgent - streaming output agent
2
3use std::pin::Pin;
4use std::sync::Arc;
5
6use futures_util::{Stream, StreamExt};
7use tokio::sync::mpsc;
8use tokio_stream::wrappers::ReceiverStream;
9
10use lc_core::language_models::BaseChatModel;
11use lc_core::runnables::RunnableConfig;
12use lc_providers::ProviderError;
13use lc_schema::Message;
14
15use super::state::AgentStreamEvent;
16
17/// Streaming Function Calling Agent
18///
19/// Streams LLM text (token by token), then emits FinalAnswer at the end.
20/// Tool-call state is exposed via `AgentStreamEvent::ToolCall`.
21/// Works with any LLM provider implementing `BaseChatModel`.
22pub struct StreamingFunctionCallingAgent {
23    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
24}
25
26impl StreamingFunctionCallingAgent {
27    /// Creates a new Streaming Function Calling Agent
28    ///
29    /// # Backward compatibility
30    /// Old code `StreamingFunctionCallingAgent::new(openai_chat)` still works.
31    pub fn new<L>(llm: L) -> Self
32    where
33        L: BaseChatModel + Send + Sync + 'static,
34        L::Error: Into<ProviderError>,
35    {
36        Self {
37            llm: lc_providers::wrap_chat_model(llm),
38        }
39    }
40
41    /// Creates an agent from a wrapped `Arc<dyn BaseChatModel>`
42    pub fn from_arc(llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>) -> Self {
43        Self { llm }
44    }
45
46    /// Streams execution: returns an event stream
47    pub async fn invoke_stream(
48        &self,
49        input: String,
50    ) -> Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send>> {
51        self.invoke_stream_with_config(input, None).await
52    }
53
54    /// Streams execution with a [`RunnableConfig`] so the streamed LLM call
55    /// emits `on_llm_start/end` to the configured callbacks/OTel backend
56    /// (T6, v0.23.0).
57    pub async fn invoke_stream_with_config(
58        &self,
59        input: String,
60        config: Option<&RunnableConfig>,
61    ) -> Pin<Box<dyn Stream<Item = AgentStreamEvent> + Send>> {
62        let (tx, rx) = mpsc::channel(32);
63        let llm = self.llm.clone();
64        let messages = vec![Message::human(input)];
65        // The spawned task is `'static`, so the config must be owned here.
66        let config = config.cloned();
67
68        tokio::spawn(async move {
69            let mut stream = match llm.stream_chat(messages, config).await {
70                Ok(s) => s,
71                Err(e) => {
72                    let _ = tx
73                        .send(AgentStreamEvent::Error {
74                            message: format!("Stream initialization failed: {}", e),
75                        })
76                        .await;
77                    return;
78                }
79            };
80
81            let mut full = String::new();
82            while let Some(chunk) = stream.next().await {
83                match chunk {
84                    Ok(chunk) => {
85                        full.push_str(&chunk.text);
86                        if tx
87                            .send(AgentStreamEvent::Text {
88                                content: chunk.text,
89                            })
90                            .await
91                            .is_err()
92                        {
93                            break;
94                        }
95                    }
96                    Err(e) => {
97                        let _ = tx
98                            .send(AgentStreamEvent::Error {
99                                message: format!("Stream error: {}", e),
100                            })
101                            .await;
102                        break;
103                    }
104                }
105            }
106
107            let _ = tx
108                .send(AgentStreamEvent::FinalAnswer { content: full })
109                .await;
110        });
111
112        Box::pin(ReceiverStream::new(rx))
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119    use lc_providers::{OpenAIChat, OpenAIConfig};
120
121    #[test]
122    fn test_new() {
123        let llm = OpenAIChat::new(OpenAIConfig::default());
124        let _agent = StreamingFunctionCallingAgent::new(llm);
125    }
126}
lc_agents/streaming/tool_call_stream.rs

lc_agents/streaming/
tool_call_stream.rs
