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                // 0.22.0 C6 fix: char-boundary preview (byte slicing panicked on
298                // CJK documents; matches retrieval_qa.rs's chars().take(100)).
299                let preview: String = doc.content.chars().take(100).collect();
300                println!("Document {}: {}", i + 1, preview);
301            }
302        }
303
304        // Step 3: Assemble Prompt
305        if self.verbose {
306            println!("\n--- Step 3: Assemble Prompt ---");
307        }
308
309        let context = self.format_context(&documents);
310
311        if self.verbose {
312            println!("History length: {} characters", history.len());
313            println!("Context length: {} characters", context.len());
314        }
315
316        // Step 4: LLM generates answer
317        if self.verbose {
318            println!("\n--- Step 4: LLM generates answer ---");
319        }
320
321        let context_str = self.format_context(&documents);
322        let messages = self.build_messages(&history_messages, &context_str, question);
323
324        let response = self
325            .llm
326            .invoke(messages, None)
327            .await
328            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
329
330        let answer = response.content;
331
332        if self.verbose {
333            println!("Answer: {}", answer);
334        }
335
336        // Step 5: Save to memory
337        self.save_context(question, &answer).await?;
338
339        if self.verbose {
340            println!("=== ConversationRetrievalChain Complete ===\n");
341        }
342
343        let mut result = HashMap::new();
344        result.insert(self.output_key.clone(), Value::String(answer));
345
346        if self.return_source_documents {
347            // Explicit error instead of silently inserting Value::Null (P1-2).
348            let sources = crate::base::documents_to_values(&documents)?;
349            result.insert(self.source_document_key.clone(), Value::Array(sources));
350        }
351
352        Ok(result)
353    }
354
355    /// Stream execution for ConversationRetrievalChain -- token by token output.
356    ///
357    /// P2-2: real streaming — loads history, retrieves and assembles the prompt,
358    /// then pushes LLM tokens via `stream_chat`. The full answer is accumulated
359    /// through an unbounded channel (the P1-4 pattern) and written to memory
360    /// once the stream completes, matching the invoke path's `save_context`.
361    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
362        self.validate_inputs(&inputs)?;
363
364        let question = inputs
365            .get(&self.input_key)
366            .and_then(|v| v.as_str())
367            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
368
369        if self.verbose {
370            println!("\n=== ConversationRetrievalChain Stream ===");
371            println!("Question: {}", question);
372        }
373
374        // Step 1: Load conversation history
375        let history_messages = self.load_history(question).await?;
376
377        // Step 2: Retrieve relevant documents
378        let documents = self
379            .retriever
380            .retrieve(question, self.k)
381            .await
382            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
383
384        if self.verbose {
385            println!("Retrieved {} documents", documents.len());
386        }
387
388        // Step 3: Assemble messages (history + context + question)
389        let context = self.format_context(&documents);
390        let messages = self.build_messages(&history_messages, &context, question);
391
392        // Step 4: Stream LLM tokens
393        let llm_stream = self
394            .llm
395            .stream_chat(messages, None)
396            .await
397            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
398
399        let memory = self.memory.clone();
400        let input_key = self.input_key.clone();
401        let output_key = self.output_key.clone();
402        let question_str = question.to_string();
403
404        // Queue tokens through an unbounded channel; the finalizer drains every
405        // token and writes the full output to memory (never a truncated one).
406        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<String>();
407
408        let stream = llm_stream.map(move |result| match result {
409            Ok(chunk) => {
410                let _ = tx.send(chunk.text.clone());
411                Ok(StreamToken {
412                    token: chunk.text,
413                    is_final: false,
414                })
415            }
416            Err(e) => Err(ChainError::StreamError(format!(
417                "Stream token error: {}",
418                e
419            ))),
420        });
421
422        let finalizer_stream = async move {
423            let mut output = String::new();
424            let mut rx = rx;
425            while let Some(token) = rx.recv().await {
426                output.push_str(&token);
427            }
428
429            // Step 5: Save to memory
430            if !output.is_empty() {
431                let mut mem = memory.lock().await;
432                let ctx_inputs = HashMap::from([(input_key.clone(), question_str.clone())]);
433                let ctx_outputs = HashMap::from([(output_key.clone(), output)]);
434                if let Err(e) = mem.save_context(&ctx_inputs, &ctx_outputs).await {
435                    log::error!("[ConversationRetrievalChain] failed to save context: {}", e);
436                }
437            }
438        };
439
440        let final_stream = stream.chain(futures_util::stream::once(async move {
441            finalizer_stream.await;
442            Ok(StreamToken {
443                token: String::new(),
444                is_final: true,
445            })
446        }));
447
448        Ok(Box::pin(final_stream))
449    }
450
451    fn name(&self) -> &str {
452        &self.name
453    }
454}
455
456#[cfg(test)]
457mod tests {
458    use super::*;
459    use async_trait::async_trait;
460    use futures_util::Stream;
461    use lc_core::language_models::{LLMResult, StreamChunk};
462    use lc_core::runnables::RunnableConfig;
463    use lc_core::{BaseLanguageModel, Runnable};
464    use lc_rag::retriever::RetrieverError;
465    use lc_shared::document::SearchResult;
466    use std::pin::Pin;
467
468    /// Mock retriever that returns the preloaded documents (up to `k`).
469    struct MockRetriever(Vec<Document>);
470
471    #[async_trait]
472    impl RetrieverTrait for MockRetriever {
473        async fn retrieve(&self, _query: &str, k: usize) -> Result<Vec<Document>, RetrieverError> {
474            Ok(self.0.iter().take(k).cloned().collect())
475        }
476        async fn retrieve_with_scores(
477            &self,
478            _query: &str,
479            _k: usize,
480        ) -> Result<Vec<SearchResult>, RetrieverError> {
481            Ok(Vec::new())
482        }
483        async fn add_documents(&self, _documents: Vec<Document>) -> Result<(), RetrieverError> {
484            Ok(())
485        }
486    }
487
488    /// Mock chat model with a deterministic token stream.
489    struct MockLLM;
490
491    #[async_trait]
492    impl Runnable<Vec<Message>, LLMResult> for MockLLM {
493        type Error = ProviderError;
494        async fn invoke(
495            &self,
496            _input: Vec<Message>,
497            _config: Option<RunnableConfig>,
498        ) -> Result<LLMResult, Self::Error> {
499            Ok(LLMResult {
500                content: "hello world".to_string(),
501                model: "mock".to_string(),
502                token_usage: None,
503                tool_calls: None,
504                thinking_content: None,
505            })
506        }
507    }
508
509    #[async_trait]
510    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
511        fn model_name(&self) -> &str {
512            "mock"
513        }
514        fn get_num_tokens(&self, t: &str) -> usize {
515            t.len()
516        }
517        fn with_temperature(self, _: f32) -> Self {
518            self
519        }
520        fn with_max_tokens(self, _: usize) -> Self {
521            self
522        }
523    }
524
525    #[async_trait]
526    impl BaseChatModel for MockLLM {
527        async fn chat(
528            &self,
529            _messages: Vec<Message>,
530            _config: Option<RunnableConfig>,
531        ) -> Result<LLMResult, Self::Error> {
532            Ok(LLMResult {
533                content: "hello world".to_string(),
534                model: "mock".to_string(),
535                token_usage: None,
536                tool_calls: None,
537                thinking_content: None,
538            })
539        }
540        async fn stream_chat(
541            &self,
542            _messages: Vec<Message>,
543            _config: Option<RunnableConfig>,
544        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
545        {
546            let tokens = [
547                Ok(StreamChunk::new("hello")),
548                Ok(StreamChunk::new(" ")),
549                Ok(StreamChunk::new("world")),
550            ];
551            Ok(Box::pin(futures_util::stream::iter(tokens)))
552        }
553    }
554
555    fn doc(content: &str) -> Document {
556        Document::new(content.to_string())
557    }
558
559    /// P2-2: ConversationRetrieval streams real tokens and persists the full
560    /// streamed answer to memory once the stream completes.
561    #[tokio::test]
562    async fn test_conversation_retrieval_stream_saves_to_memory() {
563        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![doc("ctx")]));
564        let chain =
565            ConversationRetrievalChain::new(MockLLM, retriever, ConversationBufferMemory::new());
566        let inputs = HashMap::from([("query".to_string(), Value::String("q".to_string()))]);
567
568        let mut stream = chain.stream(inputs).await.unwrap();
569        let mut tokens = Vec::new();
570        while let Some(item) = stream.next().await {
571            tokens.push(item.unwrap());
572        }
573        let text: String = tokens.iter().map(|t| t.token.as_str()).collect();
574        assert_eq!(text, "hello world");
575        assert!(tokens.last().unwrap().is_final);
576
577        // Step 5: the streamed answer is written to memory (never truncated).
578        let memory = chain.memory().clone();
579        let mem = memory.lock().await;
580        let vars = mem.load_memory_variables(&HashMap::new()).await.unwrap();
581        let messages = crate::base::variables_to_messages(&vars);
582        assert!(
583            messages.iter().any(|m| m.content.contains("hello world")),
584            "memory should contain the streamed answer, got {:?}",
585            messages
586        );
587    }
588
589    #[tokio::test]
590    async fn test_conversation_retrieval_stream_missing_input() {
591        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![]));
592        let chain =
593            ConversationRetrievalChain::new(MockLLM, retriever, ConversationBufferMemory::new());
594        let err = match chain.stream(HashMap::new()).await {
595            Ok(_) => panic!("expected a missing-input error"),
596            Err(e) => e,
597        };
598        assert!(matches!(err, ChainError::MissingInput(_)));
599    }
600}
lc_chains/conversation_retrieval.rs

lc_chains/
conversation_retrieval.rs
