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

lc_chains/
conversation_chain.rs
