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

lc_chains/
retrieval_qa.rs
