lc_chains/
conversation_retrieval.rs

1// lc-chains/src/conversation_retrieval.rs
2//! ConversationRetrieval Chain
3//!
4//! Retrieval-augmented generation chain with memory, combining conversation
5//! history with document retrieval.
6
7use async_trait::async_trait;
8use lc_core::language_models::LLMResult;
9use lc_core::{BaseChatModel, Runnable};
10use lc_memory::{BaseMemory, ConversationBufferMemory};
11use lc_rag::retriever::RetrieverTrait;
12use lc_schema::{Message, MessageType};
13use lc_shared::document::Document;
14use serde_json::Value;
15use std::collections::HashMap;
16use std::sync::Arc;
17use tokio::sync::Mutex;
18
19use crate::base::{BaseChain, ChainError, ChainResult};
20
21/// Default retrieval-augmented conversation prompt template.
22const DEFAULT_QA_PROMPT: &str = "You are an AI assistant. Please answer the user's question based on the conversation history and reference information.
23
24Conversation history:
25{history}
26
27Reference information:
28{context}
29
30Question: {question}
31
32Answer:";
33
34/// ConversationRetrievalChain
35///
36/// Retrieval-augmented conversation chain with memory that automatically:
37/// 1. Loads conversation history
38/// 2. Retrieves relevant documents
39/// 3. Combines history + context + question
40/// 4. LLM generates answer
41/// 5. Saves to conversation memory
42pub struct ConversationRetrievalChain<M: BaseChatModel> {
43    llm: M,
44    retriever: Arc<dyn RetrieverTrait>,
45    memory: Arc<Mutex<ConversationBufferMemory>>,
46
47    system_prompt: Option<String>,
48    qa_prompt_template: String,
49    input_key: String,
50    output_key: String,
51    memory_key: String,
52    name: String,
53
54    k: usize,
55    verbose: bool,
56    return_source_documents: bool,
57    source_document_key: String,
58}
59
60impl<M: BaseChatModel + 'static> ConversationRetrievalChain<M> {
61    pub fn new(
62        llm: M,
63        retriever: Arc<dyn RetrieverTrait>,
64        memory: ConversationBufferMemory,
65    ) -> Self {
66        Self {
67            llm,
68            retriever,
69            memory: Arc::new(Mutex::new(memory.with_return_messages(true))),
70            system_prompt: None,
71            qa_prompt_template: DEFAULT_QA_PROMPT.to_string(),
72            input_key: "query".to_string(),
73            output_key: "result".to_string(),
74            memory_key: "history".to_string(),
75            name: "conversation_retrieval".to_string(),
76            k: 4,
77            verbose: false,
78            return_source_documents: false,
79            source_document_key: "source_documents".to_string(),
80        }
81    }
82
83    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
84        self.system_prompt = Some(prompt.into());
85        self
86    }
87
88    pub fn with_qa_prompt(mut self, template: impl Into<String>) -> Self {
89        self.qa_prompt_template = template.into();
90        self
91    }
92
93    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
94        self.input_key = key.into();
95        self
96    }
97
98    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
99        self.output_key = key.into();
100        self
101    }
102
103    pub fn with_memory_key(mut self, key: impl Into<String>) -> Self {
104        self.memory_key = key.into();
105        self
106    }
107
108    pub fn with_name(mut self, name: impl Into<String>) -> Self {
109        self.name = name.into();
110        self
111    }
112
113    pub fn with_k(mut self, k: usize) -> Self {
114        self.k = k;
115        self
116    }
117
118    pub fn with_verbose(mut self, verbose: bool) -> Self {
119        self.verbose = verbose;
120        self
121    }
122
123    pub fn with_return_source_documents(mut self, return_source: bool) -> Self {
124        self.return_source_documents = return_source;
125        self
126    }
127
128    pub fn memory(&self) -> &Arc<Mutex<ConversationBufferMemory>> {
129        &self.memory
130    }
131
132    pub async fn clear_memory(&self) -> Result<(), ChainError> {
133        let mut memory = self.memory.lock().await;
134        memory
135            .clear()
136            .await
137            .map_err(|e| ChainError::ExecutionError(format!("Failed to clear memory: {}", e)))?;
138        Ok(())
139    }
140
141    /// Simplified query interface.
142    pub async fn query(&self, question: impl Into<String>) -> Result<String, ChainError> {
143        let inputs = HashMap::from([(self.input_key.clone(), Value::String(question.into()))]);
144        let result = self.invoke(inputs).await?;
145        result
146            .get(&self.output_key)
147            .and_then(|v| v.as_str())
148            .map(|s| s.to_string())
149            .ok_or_else(|| ChainError::OutputError("Missing output result".to_string()))
150    }
151
152    fn format_context(&self, documents: &[Document]) -> String {
153        documents
154            .iter()
155            .map(|doc| doc.content.clone())
156            .collect::<Vec<_>>()
157            .join("\n\n---\n\n")
158    }
159
160    /// Build structured messages for a given history, context, and question.
161    pub fn build_messages(
162        &self,
163        history: &[Message],
164        context: &str,
165        question: &str,
166    ) -> Vec<Message> {
167        let mut messages = Vec::new();
168
169        if let Some(system) = &self.system_prompt {
170            messages.push(Message::system(system));
171        } else {
172            messages.push(Message::system(
173                "You are an AI assistant. Answer the user's question based on the conversation history and reference information."
174            ));
175        }
176
177        for msg in history {
178            messages.push(msg.clone());
179        }
180
181        let human_content = if context.is_empty() {
182            question.to_string()
183        } else {
184            format!(
185                "Reference information:\n{}\n\nQuestion: {}",
186                context, question
187            )
188        };
189        messages.push(Message::human(&human_content));
190
191        messages
192    }
193
194    fn format_history(&self, messages: &[Message]) -> String {
195        messages
196            .iter()
197            .map(|msg| {
198                let role = match msg.message_type {
199                    MessageType::Human => "User",
200                    MessageType::AI => "Assistant",
201                    _ => "System",
202                };
203                format!("{}: {}", role, msg.content)
204            })
205            .collect::<Vec<_>>()
206            .join("\n")
207    }
208
209    async fn load_history(&self) -> Result<Vec<Message>, ChainError> {
210        let memory = self.memory.lock().await;
211        Ok(memory.chat_memory().messages().to_vec())
212    }
213
214    async fn save_context(&self, input: &str, output: &str) -> Result<(), ChainError> {
215        let mut memory = self.memory.lock().await;
216        let inputs = HashMap::from([(self.input_key.clone(), input.to_string())]);
217        let outputs = HashMap::from([(self.output_key.clone(), output.to_string())]);
218        memory
219            .save_context(&inputs, &outputs)
220            .await
221            .map_err(|e| ChainError::ExecutionError(format!("Failed to save context: {}", e)))?;
222        Ok(())
223    }
224}
225
226#[async_trait]
227impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for ConversationRetrievalChain<M>
228where
229    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
230{
231    fn input_keys(&self) -> Vec<&str> {
232        vec![&self.input_key]
233    }
234
235    fn output_keys(&self) -> Vec<&str> {
236        if self.return_source_documents {
237            vec![&self.output_key, &self.source_document_key]
238        } else {
239            vec![&self.output_key]
240        }
241    }
242
243    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
244        self.validate_inputs(&inputs)?;
245
246        let question = inputs
247            .get(&self.input_key)
248            .and_then(|v| v.as_str())
249            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
250
251        if self.verbose {
252            println!("\n=== ConversationRetrievalChain Execution ===");
253            println!("Question: {}", question);
254        }
255
256        // Step 1: Load conversation history
257        let history_messages = self.load_history().await?;
258        let history = self.format_history(&history_messages);
259
260        if self.verbose {
261            println!("History messages: {}", history_messages.len());
262        }
263
264        // Step 2: Retrieve relevant documents
265        if self.verbose {
266            println!("\n--- Step 2: Retrieve relevant documents ---");
267        }
268
269        let documents = self
270            .retriever
271            .retrieve(question, self.k)
272            .await
273            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
274
275        if self.verbose {
276            println!("Retrieved {} documents", documents.len());
277            for (i, doc) in documents.iter().enumerate() {
278                let preview = if doc.content.len() > 100 {
279                    &doc.content[..100]
280                } else {
281                    &doc.content
282                };
283                println!("Document {}: {}", i + 1, preview);
284            }
285        }
286
287        // Step 3: Assemble Prompt
288        if self.verbose {
289            println!("\n--- Step 3: Assemble Prompt ---");
290        }
291
292        let context = self.format_context(&documents);
293
294        if self.verbose {
295            println!("History length: {} characters", history.len());
296            println!("Context length: {} characters", context.len());
297        }
298
299        // Step 4: LLM generates answer
300        if self.verbose {
301            println!("\n--- Step 4: LLM generates answer ---");
302        }
303
304        let mut messages = Vec::new();
305
306        if let Some(system) = &self.system_prompt {
307            messages.push(Message::system(system));
308        } else {
309            messages.push(Message::system(
310                "You are an AI assistant. Answer the user's question based on the conversation history and reference information."
311            ));
312        }
313
314        for msg in &history_messages {
315            messages.push(msg.clone());
316        }
317
318        let context_str = self.format_context(&documents);
319        let human_content = if context_str.is_empty() {
320            question.to_string()
321        } else {
322            format!(
323                "Reference information:\n{}\n\nQuestion: {}",
324                context_str, question
325            )
326        };
327        messages.push(Message::human(&human_content));
328
329        let response = self
330            .llm
331            .invoke(messages, None)
332            .await
333            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
334
335        let answer = response.content;
336
337        if self.verbose {
338            println!("Answer: {}", answer);
339        }
340
341        // Step 5: Save to memory
342        self.save_context(question, &answer).await?;
343
344        if self.verbose {
345            println!("=== ConversationRetrievalChain Complete ===\n");
346        }
347
348        let mut result = HashMap::new();
349        result.insert(self.output_key.clone(), Value::String(answer));
350
351        if self.return_source_documents {
352            let sources: Vec<Value> = documents
353                .iter()
354                .map(|doc| serde_json::to_value(doc).unwrap_or(Value::Null))
355                .collect();
356            result.insert(self.source_document_key.clone(), Value::Array(sources));
357        }
358
359        Ok(result)
360    }
361
362    fn name(&self) -> &str {
363        &self.name
364    }
365}
lc_chains/conversation_retrieval.rs

lc_chains/
conversation_retrieval.rs
