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

lc_chains/
conversation_chain.rs
