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 futures_util::StreamExt;
9use lc_core::language_models::LLMResult;
10use lc_core::{BaseChatModel, Runnable};
11use lc_memory::{BaseMemory, ConversationBufferMemory};
12use lc_rag::retriever::RetrieverTrait;
13use lc_schema::{Message, MessageType};
14use lc_shared::document::Document;
15use serde_json::Value;
16use std::collections::HashMap;
17use std::sync::Arc;
18use tokio::sync::Mutex;
19
20use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
21
22/// ConversationRetrievalChain
23///
24/// Retrieval-augmented conversation chain with memory that automatically:
25/// 1. Loads conversation history
26/// 2. Retrieves relevant documents
27/// 3. Combines history + context + question
28/// 4. LLM generates answer
29/// 5. Saves to conversation memory
30pub struct ConversationRetrievalChain<M: BaseChatModel> {
31    llm: M,
32    retriever: Arc<dyn RetrieverTrait>,
33    memory: Arc<Mutex<dyn BaseMemory>>,
34
35    system_prompt: Option<String>,
36    input_key: String,
37    output_key: String,
38    name: String,
39
40    k: usize,
41    verbose: bool,
42    return_source_documents: bool,
43    source_document_key: String,
44}
45
46impl<M: BaseChatModel + 'static> ConversationRetrievalChain<M> {
47    pub fn new(
48        llm: M,
49        retriever: Arc<dyn RetrieverTrait>,
50        memory: ConversationBufferMemory,
51    ) -> Self {
52        // Align the memory's input/output keys with this chain's defaults
53        // ("query"/"result"). `save_context` addresses the memory by these keys,
54        // so without alignment persistence silently fails with `Missing input
55        // key 'input'` on both the invoke and stream paths.
56        Self::from_memory(
57            llm,
58            retriever,
59            Arc::new(Mutex::new(
60                memory
61                    .with_return_messages(true)
62                    .with_input_key("query".to_string())
63                    .with_output_key("result".to_string()),
64            )),
65        )
66    }
67
68    /// Create from any [`BaseMemory`] implementation (window / summary /
69    /// vector-store / persistent), mirroring `ConversationChain::from_memory`.
70    pub fn from_memory(
71        llm: M,
72        retriever: Arc<dyn RetrieverTrait>,
73        memory: Arc<Mutex<dyn BaseMemory>>,
74    ) -> Self {
75        Self {
76            llm,
77            retriever,
78            memory,
79            system_prompt: None,
80            input_key: "query".to_string(),
81            output_key: "result".to_string(),
82            name: "conversation_retrieval".to_string(),
83            k: 4,
84            verbose: false,
85            return_source_documents: false,
86            source_document_key: "source_documents".to_string(),
87        }
88    }
89
90    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
91        self.system_prompt = Some(prompt.into());
92        self
93    }
94
95    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
96        self.input_key = key.into();
97        self
98    }
99
100    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
101        self.output_key = key.into();
102        self
103    }
104
105    pub fn with_name(mut self, name: impl Into<String>) -> Self {
106        self.name = name.into();
107        self
108    }
109
110    pub fn with_k(mut self, k: usize) -> Self {
111        self.k = k;
112        self
113    }
114
115    pub fn with_verbose(mut self, verbose: bool) -> Self {
116        self.verbose = verbose;
117        self
118    }
119
120    pub fn with_return_source_documents(mut self, return_source: bool) -> Self {
121        self.return_source_documents = return_source;
122        self
123    }
124
125    pub fn memory(&self) -> &Arc<Mutex<dyn BaseMemory>> {
126        &self.memory
127    }
128
129    pub async fn clear_memory(&self) -> Result<(), ChainError> {
130        let mut memory = self.memory.lock().await;
131        memory
132            .clear()
133            .await
134            .map_err(|e| ChainError::ExecutionError(format!("Failed to clear memory: {}", e)))?;
135        Ok(())
136    }
137
138    /// Simplified query interface.
139    pub async fn query(&self, question: impl Into<String>) -> Result<String, ChainError> {
140        let inputs = HashMap::from([(self.input_key.clone(), Value::String(question.into()))]);
141        let result = self.invoke(inputs).await?;
142        result
143            .get(&self.output_key)
144            .and_then(|v| v.as_str())
145            .map(|s| s.to_string())
146            .ok_or_else(|| ChainError::OutputError("Missing output result".to_string()))
147    }
148
149    fn format_context(&self, documents: &[Document]) -> String {
150        documents
151            .iter()
152            .map(|doc| doc.content.clone())
153            .collect::<Vec<_>>()
154            .join("\n\n---\n\n")
155    }
156
157    /// Build structured messages for a given history, context, and question.
158    pub fn build_messages(
159        &self,
160        history: &[Message],
161        context: &str,
162        question: &str,
163    ) -> Vec<Message> {
164        let mut messages = Vec::new();
165
166        if let Some(system) = &self.system_prompt {
167            messages.push(Message::system(system));
168        } else {
169            messages.push(Message::system(
170                "You are an AI assistant. Answer the user's question based on the conversation history and reference information."
171            ));
172        }
173
174        for msg in history {
175            messages.push(msg.clone());
176        }
177
178        let human_content = if context.is_empty() {
179            question.to_string()
180        } else {
181            format!(
182                "Reference information:\n{}\n\nQuestion: {}",
183                context, question
184            )
185        };
186        messages.push(Message::human(&human_content));
187
188        messages
189    }
190
191    fn format_history(&self, messages: &[Message]) -> String {
192        messages
193            .iter()
194            .map(|msg| {
195                let role = match msg.message_type {
196                    MessageType::Human => "User",
197                    MessageType::AI => "Assistant",
198                    _ => "System",
199                };
200                format!("{}: {}", role, msg.content)
201            })
202            .collect::<Vec<_>>()
203            .join("\n")
204    }
205
206    async fn load_history(&self, question: &str) -> Result<Vec<Message>, ChainError> {
207        let memory = self.memory.lock().await;
208        let inputs = HashMap::from([(self.input_key.clone(), question.to_string())]);
209        let vars = memory
210            .load_memory_variables(&inputs)
211            .await
212            .map_err(|e| ChainError::ExecutionError(format!("Failed to load memory: {}", e)))?;
213        Ok(crate::base::variables_to_messages(&vars))
214    }
215
216    async fn save_context(&self, input: &str, output: &str) -> Result<(), ChainError> {
217        let mut memory = self.memory.lock().await;
218        let inputs = HashMap::from([(self.input_key.clone(), input.to_string())]);
219        let outputs = HashMap::from([(self.output_key.clone(), output.to_string())]);
220        memory
221            .save_context(&inputs, &outputs)
222            .await
223            .map_err(|e| ChainError::ExecutionError(format!("Failed to save context: {}", e)))?;
224        Ok(())
225    }
226}
227
228#[async_trait]
229impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for ConversationRetrievalChain<M>
230where
231    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
232{
233    fn input_keys(&self) -> Vec<&str> {
234        vec![&self.input_key]
235    }
236
237    fn output_keys(&self) -> Vec<&str> {
238        if self.return_source_documents {
239            vec![&self.output_key, &self.source_document_key]
240        } else {
241            vec![&self.output_key]
242        }
243    }
244
245    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
246        self.validate_inputs(&inputs)?;
247
248        let question = inputs
249            .get(&self.input_key)
250            .and_then(|v| v.as_str())
251            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
252
253        if self.verbose {
254            println!("\n=== ConversationRetrievalChain Execution ===");
255            println!("Question: {}", question);
256        }
257
258        // Step 1: Load conversation history
259        let history_messages = self.load_history(question).await?;
260        let history = self.format_history(&history_messages);
261
262        if self.verbose {
263            println!("History messages: {}", history_messages.len());
264        }
265
266        // Step 2: Retrieve relevant documents
267        if self.verbose {
268            println!("\n--- Step 2: Retrieve relevant documents ---");
269        }
270
271        let documents = self
272            .retriever
273            .retrieve(question, self.k)
274            .await
275            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
276
277        if self.verbose {
278            println!("Retrieved {} documents", documents.len());
279            for (i, doc) in documents.iter().enumerate() {
280                let preview = if doc.content.len() > 100 {
281                    &doc.content[..100]
282                } else {
283                    &doc.content
284                };
285                println!("Document {}: {}", i + 1, preview);
286            }
287        }
288
289        // Step 3: Assemble Prompt
290        if self.verbose {
291            println!("\n--- Step 3: Assemble Prompt ---");
292        }
293
294        let context = self.format_context(&documents);
295
296        if self.verbose {
297            println!("History length: {} characters", history.len());
298            println!("Context length: {} characters", context.len());
299        }
300
301        // Step 4: LLM generates answer
302        if self.verbose {
303            println!("\n--- Step 4: LLM generates answer ---");
304        }
305
306        let context_str = self.format_context(&documents);
307        let messages = self.build_messages(&history_messages, &context_str, question);
308
309        let response = self
310            .llm
311            .invoke(messages, None)
312            .await
313            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
314
315        let answer = response.content;
316
317        if self.verbose {
318            println!("Answer: {}", answer);
319        }
320
321        // Step 5: Save to memory
322        self.save_context(question, &answer).await?;
323
324        if self.verbose {
325            println!("=== ConversationRetrievalChain Complete ===\n");
326        }
327
328        let mut result = HashMap::new();
329        result.insert(self.output_key.clone(), Value::String(answer));
330
331        if self.return_source_documents {
332            // Explicit error instead of silently inserting Value::Null (P1-2).
333            let sources = crate::base::documents_to_values(&documents)?;
334            result.insert(self.source_document_key.clone(), Value::Array(sources));
335        }
336
337        Ok(result)
338    }
339
340    /// Stream execution for ConversationRetrievalChain -- token by token output.
341    ///
342    /// P2-2: real streaming — loads history, retrieves and assembles the prompt,
343    /// then pushes LLM tokens via `stream_chat`. The full answer is accumulated
344    /// through an unbounded channel (the P1-4 pattern) and written to memory
345    /// once the stream completes, matching the invoke path's `save_context`.
346    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
347        self.validate_inputs(&inputs)?;
348
349        let question = inputs
350            .get(&self.input_key)
351            .and_then(|v| v.as_str())
352            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
353
354        if self.verbose {
355            println!("\n=== ConversationRetrievalChain Stream ===");
356            println!("Question: {}", question);
357        }
358
359        // Step 1: Load conversation history
360        let history_messages = self.load_history(question).await?;
361
362        // Step 2: Retrieve relevant documents
363        let documents = self
364            .retriever
365            .retrieve(question, self.k)
366            .await
367            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
368
369        if self.verbose {
370            println!("Retrieved {} documents", documents.len());
371        }
372
373        // Step 3: Assemble messages (history + context + question)
374        let context = self.format_context(&documents);
375        let messages = self.build_messages(&history_messages, &context, question);
376
377        // Step 4: Stream LLM tokens
378        let llm_stream = self
379            .llm
380            .stream_chat(messages, None)
381            .await
382            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
383
384        let memory = self.memory.clone();
385        let input_key = self.input_key.clone();
386        let output_key = self.output_key.clone();
387        let question_str = question.to_string();
388
389        // Queue tokens through an unbounded channel; the finalizer drains every
390        // token and writes the full output to memory (never a truncated one).
391        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<String>();
392
393        let stream = llm_stream.map(move |result| match result {
394            Ok(token) => {
395                let _ = tx.send(token.clone());
396                Ok(StreamToken {
397                    token,
398                    is_final: false,
399                })
400            }
401            Err(e) => Err(ChainError::StreamError(format!(
402                "Stream token error: {}",
403                e
404            ))),
405        });
406
407        let finalizer_stream = async move {
408            let mut output = String::new();
409            let mut rx = rx;
410            while let Some(token) = rx.recv().await {
411                output.push_str(&token);
412            }
413
414            // Step 5: Save to memory
415            if !output.is_empty() {
416                let mut mem = memory.lock().await;
417                let ctx_inputs = HashMap::from([(input_key.clone(), question_str.clone())]);
418                let ctx_outputs = HashMap::from([(output_key.clone(), output)]);
419                if let Err(e) = mem.save_context(&ctx_inputs, &ctx_outputs).await {
420                    log::error!("[ConversationRetrievalChain] failed to save context: {}", e);
421                }
422            }
423        };
424
425        let final_stream = stream.chain(futures_util::stream::once(async move {
426            finalizer_stream.await;
427            Ok(StreamToken {
428                token: String::new(),
429                is_final: true,
430            })
431        }));
432
433        Ok(Box::pin(final_stream))
434    }
435
436    fn name(&self) -> &str {
437        &self.name
438    }
439}
440
441#[cfg(test)]
442mod tests {
443    use super::*;
444    use async_trait::async_trait;
445    use futures_util::Stream;
446    use lc_core::language_models::LLMResult;
447    use lc_core::runnables::RunnableConfig;
448    use lc_core::{BaseLanguageModel, Runnable};
449    use lc_rag::retriever::RetrieverError;
450    use lc_shared::document::SearchResult;
451    use std::pin::Pin;
452
453    /// Mock retriever that returns the preloaded documents (up to `k`).
454    struct MockRetriever(Vec<Document>);
455
456    #[async_trait]
457    impl RetrieverTrait for MockRetriever {
458        async fn retrieve(&self, _query: &str, k: usize) -> Result<Vec<Document>, RetrieverError> {
459            Ok(self.0.iter().take(k).cloned().collect())
460        }
461        async fn retrieve_with_scores(
462            &self,
463            _query: &str,
464            _k: usize,
465        ) -> Result<Vec<SearchResult>, RetrieverError> {
466            Ok(Vec::new())
467        }
468        async fn add_documents(&self, _documents: Vec<Document>) -> Result<(), RetrieverError> {
469            Ok(())
470        }
471    }
472
473    /// Mock chat model with a deterministic token stream.
474    #[derive(Debug)]
475    struct MockError(String);
476    impl std::fmt::Display for MockError {
477        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
478            write!(f, "{}", self.0)
479        }
480    }
481    impl std::error::Error for MockError {}
482
483    struct MockLLM;
484
485    #[async_trait]
486    impl Runnable<Vec<Message>, LLMResult> for MockLLM {
487        type Error = MockError;
488        async fn invoke(
489            &self,
490            _input: Vec<Message>,
491            _config: Option<RunnableConfig>,
492        ) -> Result<LLMResult, Self::Error> {
493            Ok(LLMResult {
494                content: "hello world".to_string(),
495                model: "mock".to_string(),
496                token_usage: None,
497                tool_calls: None,
498                thinking_content: None,
499            })
500        }
501    }
502
503    #[async_trait]
504    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
505        fn model_name(&self) -> &str {
506            "mock"
507        }
508        fn get_num_tokens(&self, t: &str) -> usize {
509            t.len()
510        }
511        fn with_temperature(self, _: f32) -> Self {
512            self
513        }
514        fn with_max_tokens(self, _: usize) -> Self {
515            self
516        }
517    }
518
519    #[async_trait]
520    impl BaseChatModel for MockLLM {
521        async fn chat(
522            &self,
523            _messages: Vec<Message>,
524            _config: Option<RunnableConfig>,
525        ) -> Result<LLMResult, Self::Error> {
526            Ok(LLMResult {
527                content: "hello world".to_string(),
528                model: "mock".to_string(),
529                token_usage: None,
530                tool_calls: None,
531                thinking_content: None,
532            })
533        }
534        async fn stream_chat(
535            &self,
536            _messages: Vec<Message>,
537            _config: Option<RunnableConfig>,
538        ) -> Result<Pin<Box<dyn Stream<Item = Result<String, Self::Error>> + Send>>, Self::Error>
539        {
540            let tokens = [
541                Ok("hello".to_string()),
542                Ok(" ".to_string()),
543                Ok("world".to_string()),
544            ];
545            Ok(Box::pin(futures_util::stream::iter(tokens)))
546        }
547    }
548
549    fn doc(content: &str) -> Document {
550        Document::new(content.to_string())
551    }
552
553    /// P2-2: ConversationRetrieval streams real tokens and persists the full
554    /// streamed answer to memory once the stream completes.
555    #[tokio::test]
556    async fn test_conversation_retrieval_stream_saves_to_memory() {
557        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![doc("ctx")]));
558        let chain =
559            ConversationRetrievalChain::new(MockLLM, retriever, ConversationBufferMemory::new());
560        let inputs = HashMap::from([("query".to_string(), Value::String("q".to_string()))]);
561
562        let mut stream = chain.stream(inputs).await.unwrap();
563        let mut tokens = Vec::new();
564        while let Some(item) = stream.next().await {
565            tokens.push(item.unwrap());
566        }
567        let text: String = tokens.iter().map(|t| t.token.as_str()).collect();
568        assert_eq!(text, "hello world");
569        assert!(tokens.last().unwrap().is_final);
570
571        // Step 5: the streamed answer is written to memory (never truncated).
572        let memory = chain.memory().clone();
573        let mem = memory.lock().await;
574        let vars = mem.load_memory_variables(&HashMap::new()).await.unwrap();
575        let messages = crate::base::variables_to_messages(&vars);
576        assert!(
577            messages.iter().any(|m| m.content.contains("hello world")),
578            "memory should contain the streamed answer, got {:?}",
579            messages
580        );
581    }
582
583    #[tokio::test]
584    async fn test_conversation_retrieval_stream_missing_input() {
585        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![]));
586        let chain =
587            ConversationRetrievalChain::new(MockLLM, retriever, ConversationBufferMemory::new());
588        let err = match chain.stream(HashMap::new()).await {
589            Ok(_) => panic!("expected a missing-input error"),
590            Err(e) => e,
591        };
592        assert!(matches!(err, ChainError::MissingInput(_)));
593    }
594}
lc_chains/conversation_retrieval.rs

lc_chains/
conversation_retrieval.rs
