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

lc_chains/
conversation_retrieval.rs
