lc_chains/
retrieval_qa.rs

1// lc-chains/src/retrieval_qa.rs
2//! RetrievalQA Chain
3//!
4//! One-stop retrieval QA chain that encapsulates the complete RAG workflow.
5
6use async_trait::async_trait;
7use futures_util::StreamExt;
8use lc_core::BaseChatModel;
9use lc_providers::{wrap_chat_model, ProviderError};
10use lc_rag::retriever::RetrieverTrait;
11use lc_schema::Message;
12use lc_shared::document::Document;
13use serde_json::Value;
14use std::collections::HashMap;
15use std::sync::Arc;
16
17use crate::base::{BaseChain, ChainError, ChainResult, ChainStream, StreamToken};
18use crate::BoxedChatModel;
19
20/// Default QA prompt template.
21const DEFAULT_QA_PROMPT: &str = "Answer the question based on the following context. If the context does not contain relevant information, say 'I don't know'.
22
23Context:
24{context}
25
26Question: {question}
27
28Answer:";
29
30/// RetrievalQA Chain
31///
32/// One-stop retrieval QA chain that automatically:
33/// 1. Retrieves relevant documents
34/// 2. Assembles prompt (context + question)
35/// 3. LLM generates answer
36pub struct RetrievalQA {
37    llm: BoxedChatModel,
38    retriever: Arc<dyn RetrieverTrait>,
39
40    prompt_template: String,
41    input_key: String,
42    output_key: String,
43    name: String,
44
45    k: usize,
46    verbose: bool,
47
48    return_source_documents: bool,
49    source_document_key: String,
50}
51
52impl RetrievalQA {
53    /// Create a new [`RetrievalQA`] chain with the given LLM and retriever.
54    pub fn new<L>(llm: L, retriever: Arc<dyn RetrieverTrait>) -> Self
55    where
56        L: BaseChatModel + Send + Sync + 'static,
57        L::Error: Into<ProviderError>,
58    {
59        Self {
60            llm: wrap_chat_model(llm),
61            retriever,
62            prompt_template: DEFAULT_QA_PROMPT.to_string(),
63            input_key: "query".to_string(),
64            output_key: "result".to_string(),
65            name: "retrieval_qa".to_string(),
66            k: 4,
67            verbose: false,
68            return_source_documents: false,
69            source_document_key: "source_documents".to_string(),
70        }
71    }
72
73    /// Set the prompt template.
74    pub fn with_prompt_template(mut self, template: impl Into<String>) -> Self {
75        self.prompt_template = template.into();
76        self
77    }
78
79    /// Set the input key.
80    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
81        self.input_key = key.into();
82        self
83    }
84
85    /// Set the output key.
86    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
87        self.output_key = key.into();
88        self
89    }
90
91    /// Set the chain name.
92    pub fn with_name(mut self, name: impl Into<String>) -> Self {
93        self.name = name.into();
94        self
95    }
96
97    /// Set the number of documents to retrieve.
98    pub fn with_k(mut self, k: usize) -> Self {
99        self.k = k;
100        self
101    }
102
103    /// Set verbose mode.
104    pub fn with_verbose(mut self, verbose: bool) -> Self {
105        self.verbose = verbose;
106        self
107    }
108
109    /// Set whether to return the source documents with the answer.
110    pub fn with_return_source_documents(mut self, return_source: bool) -> Self {
111        self.return_source_documents = return_source;
112        self
113    }
114
115    /// Set the key under which source documents are placed in the output.
116    pub fn with_source_document_key(mut self, key: impl Into<String>) -> Self {
117        self.source_document_key = key.into();
118        self
119    }
120
121    /// Get the retriever reference.
122    pub fn retriever(&self) -> &Arc<dyn RetrieverTrait> {
123        &self.retriever
124    }
125
126    /// Get the number of documents retrieved (`k`).
127    pub fn k(&self) -> usize {
128        self.k
129    }
130
131    fn format_context(&self, documents: &[Document]) -> String {
132        documents
133            .iter()
134            .map(|doc| doc.content.clone())
135            .collect::<Vec<_>>()
136            .join("\n\n")
137    }
138
139    fn build_prompt(&self, context: &str, question: &str) -> String {
140        self.prompt_template
141            .replace("{context}", context)
142            .replace("{question}", question)
143    }
144
145    /// Simplified query interface returning the answer string.
146    pub async fn query(&self, question: impl Into<String>) -> Result<String, ChainError> {
147        let inputs = HashMap::from([(self.input_key.clone(), Value::String(question.into()))]);
148
149        let result = self.invoke(inputs).await?;
150
151        result
152            .get(&self.output_key)
153            .and_then(|v| v.as_str())
154            .map(|s| s.to_string())
155            .ok_or_else(|| ChainError::OutputError("Missing output result".to_string()))
156    }
157
158    /// Query the chain, returning both the answer and the retrieved source documents.
159    pub async fn query_with_sources(
160        &self,
161        question: impl Into<String>,
162    ) -> Result<(String, Vec<Document>), ChainError> {
163        // Reuses the single `run` pipeline instead of duplicating retrieval +
164        // prompt assembly here (P1-5: previously re-implemented the whole chain
165        // when `return_source_documents` was false).
166        let question = question.into();
167        self.run(&question).await
168    }
169
170    /// Shared execution pipeline: retrieve → assemble prompt → LLM.
171    ///
172    /// Used by both [`BaseChain::invoke`] and [`Self::query_with_sources`] so the
173    /// RAG path is defined exactly once.
174    async fn run(&self, question: &str) -> Result<(String, Vec<Document>), ChainError> {
175        if self.verbose {
176            println!("\n=== RetrievalQA Execution ===");
177            println!("Question: {}", question);
178            println!("Retrieval count (k): {}", self.k);
179            println!("\n--- Step 1: Retrieve relevant documents ---");
180        }
181
182        let documents = self
183            .retriever
184            .retrieve(question, self.k)
185            .await
186            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
187
188        if self.verbose {
189            println!("Retrieved {} documents", documents.len());
190            for (i, doc) in documents.iter().enumerate() {
191                let preview: String = doc.content.chars().take(100).collect();
192                println!("Document {}: {}", i + 1, preview);
193            }
194            if documents.is_empty() {
195                println!("Warning: No relevant documents retrieved");
196            }
197            println!("\n--- Step 2: Assemble Prompt ---");
198        }
199
200        let context = self.format_context(&documents);
201        let prompt = self.build_prompt(&context, question);
202
203        if self.verbose {
204            println!("Context length: {} characters", context.len());
205            println!("Prompt length: {} characters", prompt.len());
206            println!("\n--- Step 3: LLM generates answer ---");
207        }
208
209        let messages = vec![Message::human(&prompt)];
210        let response = self
211            .llm
212            .invoke(messages, None)
213            .await
214            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
215
216        if self.verbose {
217            println!("Answer: {}", response.content);
218            println!("=== RetrievalQA Complete ===\n");
219        }
220
221        Ok((response.content, documents))
222    }
223}
224
225#[async_trait]
226impl BaseChain for RetrievalQA {
227    fn input_keys(&self) -> Vec<&str> {
228        vec![&self.input_key]
229    }
230
231    fn output_keys(&self) -> Vec<&str> {
232        if self.return_source_documents {
233            vec![&self.output_key, &self.source_document_key]
234        } else {
235            vec![&self.output_key]
236        }
237    }
238
239    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
240        self.validate_inputs(&inputs)?;
241
242        let question = inputs
243            .get(&self.input_key)
244            .and_then(|v| v.as_str())
245            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
246
247        let (answer, documents) = self.run(question).await?;
248
249        let mut result = HashMap::new();
250        result.insert(self.output_key.clone(), Value::String(answer));
251
252        if self.return_source_documents {
253            // Explicit error instead of silently inserting Value::Null (P1-2).
254            let sources = crate::base::documents_to_values(&documents)?;
255            result.insert(self.source_document_key.clone(), Value::Array(sources));
256        }
257
258        Ok(result)
259    }
260
261    /// Stream execution for RetrievalQA -- token by token output.
262    ///
263    /// P2-2: real streaming — retrieves and assembles the prompt first, then
264    /// pushes LLM tokens via `stream_chat` instead of wrapping `invoke` in a
265    /// single chunk (the base default).
266    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
267        self.validate_inputs(&inputs)?;
268
269        let question = inputs
270            .get(&self.input_key)
271            .and_then(|v| v.as_str())
272            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
273
274        if self.verbose {
275            println!("\n=== RetrievalQA Stream ===");
276            println!("Question: {}", question);
277            println!("Retrieval count (k): {}", self.k);
278        }
279
280        let documents = self
281            .retriever
282            .retrieve(question, self.k)
283            .await
284            .map_err(|e| ChainError::ExecutionError(format!("Retrieval failed: {}", e)))?;
285
286        if self.verbose {
287            println!("Retrieved {} documents", documents.len());
288        }
289
290        let context = self.format_context(&documents);
291        let prompt = self.build_prompt(&context, question);
292
293        let messages = vec![Message::human(&prompt)];
294        let llm_stream = self
295            .llm
296            .stream_chat(messages, None)
297            .await
298            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
299
300        let stream = llm_stream.map(move |result| match result {
301            Ok(chunk) => Ok(StreamToken {
302                token: chunk.text,
303                is_final: false,
304            }),
305            Err(e) => Err(ChainError::StreamError(format!(
306                "Stream token error: {}",
307                e
308            ))),
309        });
310
311        let final_stream = stream.chain(futures_util::stream::once(async move {
312            Ok(StreamToken {
313                token: String::new(),
314                is_final: true,
315            })
316        }));
317
318        Ok(Box::pin(final_stream))
319    }
320
321    fn name(&self) -> &str {
322        &self.name
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use async_trait::async_trait;
330    use futures_util::Stream;
331    use lc_core::language_models::{LLMResult, StreamChunk};
332    use lc_core::runnables::RunnableConfig;
333    use lc_core::{BaseLanguageModel, Runnable};
334    use lc_rag::retriever::RetrieverError;
335    use lc_shared::document::SearchResult;
336    use std::pin::Pin;
337
338    /// Mock retriever that returns the preloaded documents (up to `k`).
339    struct MockRetriever(Vec<Document>);
340
341    #[async_trait]
342    impl RetrieverTrait for MockRetriever {
343        async fn retrieve(&self, _query: &str, k: usize) -> Result<Vec<Document>, RetrieverError> {
344            Ok(self.0.iter().take(k).cloned().collect())
345        }
346        async fn retrieve_with_scores(
347            &self,
348            _query: &str,
349            _k: usize,
350        ) -> Result<Vec<SearchResult>, RetrieverError> {
351            Ok(Vec::new())
352        }
353        async fn add_documents(&self, _documents: Vec<Document>) -> Result<(), RetrieverError> {
354            Ok(())
355        }
356    }
357
358    /// Mock chat model with a deterministic token stream.
359    struct MockLLM;
360
361    #[async_trait]
362    impl Runnable<Vec<Message>, LLMResult> for MockLLM {
363        type Error = ProviderError;
364        async fn invoke(
365            &self,
366            _input: Vec<Message>,
367            _config: Option<RunnableConfig>,
368        ) -> Result<LLMResult, Self::Error> {
369            Ok(LLMResult {
370                content: "hello world".to_string(),
371                model: "mock".to_string(),
372                token_usage: None,
373                tool_calls: None,
374                thinking_content: None,
375            })
376        }
377    }
378
379    #[async_trait]
380    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockLLM {
381        fn model_name(&self) -> &str {
382            "mock"
383        }
384        fn get_num_tokens(&self, t: &str) -> usize {
385            t.len()
386        }
387        fn with_temperature(self, _: f32) -> Self {
388            self
389        }
390        fn with_max_tokens(self, _: usize) -> Self {
391            self
392        }
393    }
394
395    #[async_trait]
396    impl BaseChatModel for MockLLM {
397        async fn chat(
398            &self,
399            _messages: Vec<Message>,
400            _config: Option<RunnableConfig>,
401        ) -> Result<LLMResult, Self::Error> {
402            Ok(LLMResult {
403                content: "hello world".to_string(),
404                model: "mock".to_string(),
405                token_usage: None,
406                tool_calls: None,
407                thinking_content: None,
408            })
409        }
410        async fn stream_chat(
411            &self,
412            _messages: Vec<Message>,
413            _config: Option<RunnableConfig>,
414        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
415        {
416            let tokens = [
417                Ok(StreamChunk::new("hello")),
418                Ok(StreamChunk::new(" ")),
419                Ok(StreamChunk::new("world")),
420            ];
421            Ok(Box::pin(futures_util::stream::iter(tokens)))
422        }
423    }
424
425    fn doc(content: &str) -> Document {
426        Document::new(content.to_string())
427    }
428
429    /// P2-2: RetrievalQA streams real tokens from the LLM instead of wrapping
430    /// `invoke` in a single chunk.
431    #[tokio::test]
432    async fn test_retrieval_qa_stream_emits_tokens() {
433        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![doc("ctx")]));
434        let chain = RetrievalQA::new(MockLLM, retriever);
435        let inputs = HashMap::from([("query".to_string(), Value::String("q".to_string()))]);
436
437        let mut stream = chain.stream(inputs).await.unwrap();
438        let mut tokens = Vec::new();
439        while let Some(item) = stream.next().await {
440            tokens.push(item.unwrap());
441        }
442        let text: String = tokens.iter().map(|t| t.token.as_str()).collect();
443        assert_eq!(text, "hello world");
444        assert!(tokens.last().unwrap().is_final);
445        // Multiple non-final tokens prove the stream is token-by-token, not one
446        // wrapped chunk.
447        assert!(tokens.iter().filter(|t| !t.is_final).count() >= 2);
448    }
449
450    #[tokio::test]
451    async fn test_retrieval_qa_stream_missing_input() {
452        let retriever: Arc<dyn RetrieverTrait> = Arc::new(MockRetriever(vec![]));
453        let chain = RetrievalQA::new(MockLLM, retriever);
454        let err = match chain.stream(HashMap::new()).await {
455            Ok(_) => panic!("expected a missing-input error"),
456            Err(e) => e,
457        };
458        assert!(matches!(err, ChainError::MissingInput(_)));
459    }
460}
lc_chains/retrieval_qa.rs

lc_chains/
retrieval_qa.rs
