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

lc_chains/
retrieval_qa.rs
