langchain_rust/chain/combine_documents/
retrieval_qa.rs

1use std::collections::HashMap;
2
3use serde_json::{json, Value};
4
5use crate::{
6    chain::{Chain, ChainError, ChainEventStream, StuffDocument, DEFAULT_RESULT_KEY},
7    language_models::{llm::LLM, GenerateResult},
8    prompt::PromptArgs,
9    schemas::{messages::value_to_text, Document, Retriever},
10};
11
12const RETRIEVAL_QA_DEFAULT_INPUT_KEY: &str = "query";
13const RETRIEVAL_QA_DEFAULT_OUTPUT_KEY: &str = "output";
14/// Output key of [`RetrievalQA::execute`] holding the retrieved documents, when
15/// [`with_return_source_documents`](RetrievalQA::with_return_source_documents) is on.
16pub const RETRIEVAL_QA_SOURCE_DOCUMENTS_KEY: &str = "source_documents";
17
18/// Question answering over a [`Retriever`]: fetches the documents relevant to the query,
19/// then hands them with the question to a document-combining chain.
20///
21/// Stateless, unlike [`ConversationalRetrieverChain`](crate::chain::ConversationalRetrieverChain):
22/// no memory, no question rephrasing. The combining chain receives `input_documents` and
23/// `question`, so [`StuffDocument`], [`RefineDocuments`](crate::chain::RefineDocuments) and
24/// [`MapReduceDocuments`](crate::chain::MapReduceDocuments) QA chains all fit.
25///
26/// ```rust,ignore
27/// let qa = RetrievalQA::from_llm(OpenAI::default(), store_retriever);
28/// let answer = qa.invoke(prompt_args! { "query" => "What is our refund policy?" }).await?;
29/// ```
30pub struct RetrievalQA {
31    retriever: Box<dyn Retriever>,
32    combine_documents_chain: Box<dyn Chain>,
33    input_key: String,
34    output_key: String,
35    return_source_documents: bool,
36}
37
38impl RetrievalQA {
39    pub fn new<R: Into<Box<dyn Retriever>>, C: Into<Box<dyn Chain>>>(
40        retriever: R,
41        combine_documents_chain: C,
42    ) -> Self {
43        Self {
44            retriever: retriever.into(),
45            combine_documents_chain: combine_documents_chain.into(),
46            input_key: RETRIEVAL_QA_DEFAULT_INPUT_KEY.into(),
47            output_key: RETRIEVAL_QA_DEFAULT_OUTPUT_KEY.into(),
48            return_source_documents: false,
49        }
50    }
51
52    /// Stuffs every retrieved document into one QA prompt.
53    pub fn from_llm<L: Into<Box<dyn LLM>>, R: Into<Box<dyn Retriever>>>(
54        llm: L,
55        retriever: R,
56    ) -> Self {
57        Self::new(retriever, StuffDocument::load_stuff_qa(llm))
58    }
59
60    /// Input key holding the query (default `query`).
61    pub fn with_input_key<S: Into<String>>(mut self, key: S) -> Self {
62        self.input_key = key.into();
63        self
64    }
65
66    /// Output key of the answer in [`execute`](Chain::execute) (default `output`).
67    pub fn with_output_key<S: Into<String>>(mut self, key: S) -> Self {
68        self.output_key = key.into();
69        self
70    }
71
72    /// Adds the retrieved documents to [`execute`](Chain::execute)'s output under
73    /// [`RETRIEVAL_QA_SOURCE_DOCUMENTS_KEY`].
74    pub fn with_return_source_documents(mut self, yes: bool) -> Self {
75        self.return_source_documents = yes;
76        self
77    }
78
79    async fn retrieve(
80        &self,
81        input_variables: &PromptArgs,
82    ) -> Result<(PromptArgs, Vec<Document>), ChainError> {
83        let query = input_variables
84            .get(&self.input_key)
85            .map(value_to_text)
86            .ok_or_else(|| ChainError::MissingInputVariable(self.input_key.clone()))?;
87        let documents = self
88            .retriever
89            .get_relevant_documents(&query)
90            .await
91            .map_err(|e| ChainError::RetrieverError(e.to_string()))?;
92        let inputs = crate::prompt_args! {
93            "question" => query,
94            "input_documents" => documents,
95        };
96        Ok((inputs, documents))
97    }
98}
99
100#[async_trait::async_trait]
101impl Chain for RetrievalQA {
102    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
103        let (inputs, _) = self.retrieve(&input_variables).await?;
104        self.combine_documents_chain.call(inputs).await
105    }
106
107    async fn execute(
108        &self,
109        input_variables: PromptArgs,
110    ) -> Result<HashMap<String, Value>, ChainError> {
111        let (inputs, documents) = self.retrieve(&input_variables).await?;
112        let result = self.combine_documents_chain.call(inputs).await?;
113        let mut output = HashMap::new();
114        output.insert(self.output_key.clone(), json!(result.generation));
115        output.insert(DEFAULT_RESULT_KEY.to_string(), json!(result));
116        if self.return_source_documents {
117            output.insert(
118                RETRIEVAL_QA_SOURCE_DOCUMENTS_KEY.to_string(),
119                json!(documents),
120            );
121        }
122        Ok(output)
123    }
124
125    /// Retrieves, then streams the combining chain's answer.
126    fn call_stream<'a>(&'a self, input_variables: PromptArgs) -> ChainEventStream<'a> {
127        use futures::StreamExt;
128        Box::pin(
129            futures::stream::once(async move {
130                match self.retrieve(&input_variables).await {
131                    Ok((inputs, _)) => self.combine_documents_chain.call_stream(inputs),
132                    Err(e) => Box::pin(futures::stream::iter([Err(e)])) as ChainEventStream<'a>,
133                }
134            })
135            .flatten(),
136        )
137    }
138
139    fn get_input_keys(&self) -> Vec<String> {
140        vec![self.input_key.clone()]
141    }
142
143    fn get_output_keys(&self) -> Vec<String> {
144        let mut keys = vec![self.output_key.clone(), DEFAULT_RESULT_KEY.to_string()];
145        if self.return_source_documents {
146            keys.push(RETRIEVAL_QA_SOURCE_DOCUMENTS_KEY.to_string());
147        }
148        keys
149    }
150}
langchain_rust/chain/combine_documents/retrieval_qa.rs

langchain_rust/chain/combine_documents/
retrieval_qa.rs
