lc_chains/
conversation_chain.rs

1// lc-chains/src/conversation_chain.rs
2//! Conversation Chain
3//!
4//! A Chain with memory, supporting multi-turn conversations.
5
6use async_trait::async_trait;
7use futures_util::StreamExt;
8use lc_core::BaseChatModel;
9use lc_memory::{BaseMemory, ConversationBufferMemory};
10use lc_providers::{wrap_chat_model, ProviderError};
11use lc_schema::Message;
12use serde_json::Value;
13use std::collections::HashMap;
14use std::sync::Arc;
15use tokio::sync::Mutex;
16
17use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
18use crate::BoxedChatModel;
19
20/// Conversation Chain
21///
22/// A Chain with memory that automatically saves and loads conversation history.
23pub struct ConversationChain {
24    llm: BoxedChatModel,
25    memory: Arc<Mutex<dyn BaseMemory>>,
26    system_prompt: Option<String>,
27    input_key: String,
28    output_key: String,
29    name: String,
30    verbose: bool,
31}
32
33impl ConversationChain {
34    /// Create a new ConversationChain.
35    ///
36    /// # Arguments
37    /// * `llm` - LLM client (any type implementing BaseChatModel)
38    /// * `memory` - Conversation memory
39    pub fn new<L>(llm: L, memory: ConversationBufferMemory) -> Self
40    where
41        L: BaseChatModel + Send + Sync + 'static,
42        L::Error: Into<ProviderError>,
43    {
44        Self::from_memory(llm, Arc::new(Mutex::new(memory.with_return_messages(true))))
45    }
46
47    /// Create a ConversationChain from any [`BaseMemory`] implementation.
48    ///
49    /// Unlike [`ConversationChain::new`] (which takes the concrete
50    /// `ConversationBufferMemory`), this accepts any memory — window, summary,
51    /// vector-store, persistent — so the chain's memory is pluggable without
52    /// changing the chain source.
53    pub fn from_memory<L>(llm: L, memory: Arc<Mutex<dyn BaseMemory>>) -> Self
54    where
55        L: BaseChatModel + Send + Sync + 'static,
56        L::Error: Into<ProviderError>,
57    {
58        Self::from_wrapped_memory(wrap_chat_model(llm), memory)
59    }
60
61    /// Construct from an already-wrapped model (internal builder path).
62    pub(crate) fn from_wrapped_memory(
63        llm: BoxedChatModel,
64        memory: Arc<Mutex<dyn BaseMemory>>,
65    ) -> Self {
66        Self {
67            llm,
68            memory,
69            system_prompt: None,
70            input_key: "input".to_string(),
71            output_key: "output".to_string(),
72            name: "conversation_chain".to_string(),
73            verbose: false,
74        }
75    }
76
77    /// Set system prompt.
78    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
79        self.system_prompt = Some(prompt.into());
80        self
81    }
82
83    /// Set input key name.
84    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
85        self.input_key = key.into();
86        self
87    }
88
89    /// Set output key name.
90    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
91        self.output_key = key.into();
92        self
93    }
94
95    /// Set chain name.
96    pub fn with_name(mut self, name: impl Into<String>) -> Self {
97        self.name = name.into();
98        self
99    }
100
101    /// Set verbose mode.
102    pub fn with_verbose(mut self, verbose: bool) -> Self {
103        self.verbose = verbose;
104        self
105    }
106
107    /// Get memory reference.
108    pub fn memory(&self) -> &Arc<Mutex<dyn BaseMemory>> {
109        &self.memory
110    }
111
112    /// Create a [`ConversationChainBuilder`] from an LLM.
113    pub fn builder<L>(llm: L) -> ConversationChainBuilder
114    where
115        L: BaseChatModel + Send + Sync + 'static,
116        L::Error: Into<ProviderError>,
117    {
118        ConversationChainBuilder::new(llm)
119    }
120
121    /// Clear memory.
122    pub async fn clear_memory(&self) -> Result<(), ChainError> {
123        let mut memory = self.memory.lock().await;
124        memory
125            .clear()
126            .await
127            .map_err(|e| ChainError::ExecutionError(format!("Failed to clear memory: {}", e)))?;
128        Ok(())
129    }
130
131    /// Simplified prediction interface.
132    ///
133    /// Takes a user input string, returns AI response string.
134    pub async fn predict(&self, input: impl Into<String>) -> Result<String, ChainError> {
135        let inputs = HashMap::from([(self.input_key.clone(), Value::String(input.into()))]);
136
137        let result = self.invoke(inputs).await?;
138
139        result
140            .get(&self.output_key)
141            .and_then(|v| v.as_str())
142            .map(|s| s.to_string())
143            .ok_or_else(|| ChainError::OutputError("Missing output".to_string()))
144    }
145
146    /// Prepare message list.
147    ///
148    /// Combines system prompt, history messages, and current user input.
149    pub fn prepare_messages(&self, input: &str, history_messages: &[Message]) -> Vec<Message> {
150        let mut messages = Vec::new();
151
152        if let Some(system_prompt) = &self.system_prompt {
153            messages.push(Message::system(system_prompt));
154        }
155
156        for msg in history_messages {
157            messages.push(msg.clone());
158        }
159
160        messages.push(Message::human(input));
161
162        messages
163    }
164
165    /// Load history messages through the memory trait's `load_memory_variables`.
166    ///
167    /// Accepts any memory shape (array of `Message` objects or a rendered
168    /// history string); the current input is forwarded so input-aware memories
169    /// (e.g. `VectorStoreRetrieverMemory`) can use it as the retrieval query.
170    async fn load_history(&self, input: &str) -> Result<Vec<Message>, ChainError> {
171        let memory = self.memory.lock().await;
172        let inputs = HashMap::from([(self.input_key.clone(), input.to_string())]);
173        let vars = memory
174            .load_memory_variables(&inputs)
175            .await
176            .map_err(|e| ChainError::ExecutionError(format!("Failed to load memory: {}", e)))?;
177        Ok(crate::base::variables_to_messages(&vars))
178    }
179
180    /// Save conversation context.
181    async fn save_context(&self, input: &str, output: &str) -> Result<(), ChainError> {
182        let mut memory = self.memory.lock().await;
183
184        let inputs = HashMap::from([(self.input_key.clone(), input.to_string())]);
185        let outputs = HashMap::from([(self.output_key.clone(), output.to_string())]);
186
187        memory
188            .save_context(&inputs, &outputs)
189            .await
190            .map_err(|e| ChainError::ExecutionError(format!("Failed to save context: {}", e)))?;
191
192        Ok(())
193    }
194}
195
196#[async_trait]
197impl BaseChain for ConversationChain {
198    fn input_keys(&self) -> Vec<&str> {
199        vec![&self.input_key]
200    }
201
202    fn output_keys(&self) -> Vec<&str> {
203        vec![&self.output_key]
204    }
205
206    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
207        self.validate_inputs(&inputs)?;
208
209        let input = inputs
210            .get(&self.input_key)
211            .and_then(|v| v.as_str())
212            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
213
214        if self.verbose {
215            println!("\n=== ConversationChain execution ===");
216            println!("User input: {}", input);
217        }
218
219        let history_messages = self.load_history(input).await?;
220
221        if self.verbose && !history_messages.is_empty() {
222            println!("History message count: {}", history_messages.len());
223        }
224
225        let messages = self.prepare_messages(input, &history_messages);
226
227        if self.verbose {
228            println!("Total message count: {}", messages.len());
229        }
230
231        let result = self
232            .llm
233            .invoke(messages, None)
234            .await
235            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
236
237        let output = result.content;
238
239        if self.verbose {
240            println!("AI response: {}", output);
241        }
242
243        self.save_context(input, &output).await?;
244
245        if self.verbose {
246            println!("=== ConversationChain complete ===\n");
247        }
248
249        let mut result = HashMap::new();
250        result.insert(self.output_key.clone(), Value::String(output));
251
252        Ok(result)
253    }
254
255    /// Stream execution for ConversationChain -- token by token output.
256    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
257        self.validate_inputs(&inputs)?;
258
259        let input = inputs
260            .get(&self.input_key)
261            .and_then(|v| v.as_str())
262            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
263
264        let history_messages = self.load_history(input).await?;
265
266        let messages = self.prepare_messages(input, &history_messages);
267
268        let llm_stream = self
269            .llm
270            .stream_chat(messages, None)
271            .await
272            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
273
274        let memory = self.memory.clone();
275        let input_key = self.input_key.clone();
276        let output_key = self.output_key.clone();
277        let input_str = input.to_string();
278
279        // P1-4: queue tokens through an unbounded channel instead of `try_lock`
280        // on a shared mutex. The map closure's sync sender never drops a token,
281        // so the memory write sees the full output rather than a truncated one.
282        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<String>();
283
284        let stream = llm_stream.map(move |result| match result {
285            Ok(token) => {
286                let _ = tx.send(token.clone());
287                Ok(StreamToken {
288                    token,
289                    is_final: false,
290                })
291            }
292            Err(e) => Err(ChainError::StreamError(format!(
293                "Stream token error: {}",
294                e
295            ))),
296        });
297
298        let finalizer_stream = async move {
299            // The channel closes once the map closure's sender is dropped (stream
300            // exhausted); drain every queued token into the memory write.
301            let mut output = String::new();
302            let mut rx = rx;
303            while let Some(token) = rx.recv().await {
304                output.push_str(&token);
305            }
306
307            if !output.is_empty() {
308                let mut mem = memory.lock().await;
309                let ctx_inputs = HashMap::from([(input_key.clone(), input_str.clone())]);
310                let ctx_outputs = HashMap::from([(output_key.clone(), output)]);
311                if let Err(e) = mem.save_context(&ctx_inputs, &ctx_outputs).await {
312                    log::error!("[ConversationChain] failed to save context: {}", e);
313                }
314            }
315        };
316
317        let final_stream = stream.chain(futures_util::stream::once(async move {
318            finalizer_stream.await;
319            Ok(StreamToken {
320                token: String::new(),
321                is_final: true,
322            })
323        }));
324
325        Ok(Box::pin(final_stream))
326    }
327
328    fn name(&self) -> &str {
329        &self.name
330    }
331}
332
333/// ConversationChain Builder.
334///
335/// Convenience builder for ConversationChain.
336pub struct ConversationChainBuilder {
337    llm: BoxedChatModel,
338    memory: Option<Arc<Mutex<dyn BaseMemory>>>,
339    system_prompt: Option<String>,
340    input_key: Option<String>,
341    output_key: Option<String>,
342    name: Option<String>,
343    verbose: Option<bool>,
344}
345
346impl ConversationChainBuilder {
347    /// Create a new [`ConversationChainBuilder`] with the given chat model.
348    pub fn new<L>(llm: L) -> Self
349    where
350        L: BaseChatModel + Send + Sync + 'static,
351        L::Error: Into<ProviderError>,
352    {
353        Self {
354            llm: wrap_chat_model(llm),
355            memory: None,
356            system_prompt: None,
357            input_key: None,
358            output_key: None,
359            name: None,
360            verbose: None,
361        }
362    }
363
364    /// Set the conversation memory.
365    pub fn memory<Mem: BaseMemory + 'static>(mut self, memory: Mem) -> Self {
366        self.memory = Some(Arc::new(Mutex::new(memory)));
367        self
368    }
369
370    /// Set the system prompt.
371    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
372        self.system_prompt = Some(prompt.into());
373        self
374    }
375
376    /// Set the input key.
377    pub fn input_key(mut self, key: impl Into<String>) -> Self {
378        self.input_key = Some(key.into());
379        self
380    }
381
382    /// Set the output key.
383    pub fn output_key(mut self, key: impl Into<String>) -> Self {
384        self.output_key = Some(key.into());
385        self
386    }
387
388    /// Set the chain name.
389    pub fn name(mut self, name: impl Into<String>) -> Self {
390        self.name = Some(name.into());
391        self
392    }
393
394    /// Set verbose mode.
395    pub fn verbose(mut self, verbose: bool) -> Self {
396        self.verbose = Some(verbose);
397        self
398    }
399
400    /// Build the final [`ConversationChain`].
401    pub fn build(self) -> ConversationChain {
402        let mut chain = match self.memory {
403            Some(memory) => ConversationChain::from_wrapped_memory(self.llm, memory),
404            None => ConversationChain::from_wrapped_memory(
405                self.llm,
406                Arc::new(Mutex::new(
407                    ConversationBufferMemory::new().with_return_messages(true),
408                )),
409            ),
410        };
411
412        if let Some(prompt) = self.system_prompt {
413            chain = chain.with_system_prompt(prompt);
414        }
415
416        if let Some(key) = self.input_key {
417            chain = chain.with_input_key(key);
418        }
419
420        if let Some(key) = self.output_key {
421            chain = chain.with_output_key(key);
422        }
423
424        if let Some(name) = self.name {
425            chain = chain.with_name(name);
426        }
427
428        if let Some(verbose) = self.verbose {
429            chain = chain.with_verbose(verbose);
430        }
431
432        chain
433    }
434}
lc_chains/conversation_chain.rs

lc_chains/
conversation_chain.rs
