langchain_rust/chain/conversational_retrieval_qa/
conversational_retrieval_qa.rs

1use futures::Stream;
2use futures_util::{pin_mut, StreamExt};
3use std::{collections::HashMap, pin::Pin, sync::Arc};
4
5use async_stream::stream;
6use async_trait::async_trait;
7use serde_json::{json, Value};
8use tokio::sync::Mutex;
9
10use crate::{
11    chain::{
12        Chain, ChainError, CondenseQuestionPromptBuilder, StuffQAPromptBuilder, DEFAULT_RESULT_KEY,
13    },
14    language_models::{GenerateResult, TokenUsage},
15    prompt::PromptArgs,
16    schemas::{BaseMemory, Message, Retriever, StreamData},
17};
18// _conversationalRetrievalQADefaultInputKey             = "question"
19// _conversationalRetrievalQADefaultSourceDocumentKey    = "source_documents"
20// 	_conversationalRetrievalQADefaultGeneratedQuestionKey = "generated_question"
21// )
22
23const CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_SOURCE_DOCUMENT_KEY: &str = "source_documents";
24const CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_GENERATED_QUESTION_KEY: &str = "generated_question";
25
26pub struct ConversationalRetrieverChain {
27    pub(crate) retriever: Box<dyn Retriever>,
28    pub memory: Arc<Mutex<dyn BaseMemory>>,
29    pub(crate) combine_documents_chain: Box<dyn Chain>,
30    pub(crate) condense_question_chain: Box<dyn Chain>,
31    pub(crate) rephrase_question: bool,
32    pub(crate) return_source_documents: bool,
33    pub(crate) input_key: String,  //Default is `question`
34    pub(crate) output_key: String, //default is output
35}
36impl ConversationalRetrieverChain {
37    async fn get_question(
38        &self,
39        history: &[Message],
40        input: &str,
41    ) -> Result<(String, Option<TokenUsage>), ChainError> {
42        if history.is_empty() {
43            return Ok((input.to_string(), None));
44        }
45        let mut token_usage: Option<TokenUsage> = None;
46        let question = match self.rephrase_question {
47            true => {
48                let result = self
49                    .condense_question_chain
50                    .call(
51                        CondenseQuestionPromptBuilder::new()
52                            .question(input)
53                            .chat_history(history)
54                            .build(),
55                    )
56                    .await?;
57                if let Some(tokens) = result.tokens {
58                    token_usage = Some(tokens);
59                };
60                result.generation
61            }
62            false => input.to_string(),
63        };
64
65        Ok((question, token_usage))
66    }
67}
68
69#[async_trait]
70impl Chain for ConversationalRetrieverChain {
71    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
72        let output = self.execute(input_variables).await?;
73        let result: GenerateResult = serde_json::from_value(output[DEFAULT_RESULT_KEY].clone())?;
74        Ok(result)
75    }
76
77    async fn execute(
78        &self,
79        input_variables: PromptArgs,
80    ) -> Result<HashMap<String, Value>, ChainError> {
81        let mut token_usage: Option<TokenUsage> = None;
82        let input_variable = &input_variables
83            .get(&self.input_key)
84            .ok_or(ChainError::MissingInputVariable(self.input_key.clone()))?;
85
86        let human_message =
87            Message::new_human_message(crate::schemas::messages::value_to_text(input_variable));
88        let history = self.memory.lock().await.messages().await?;
89
90        let (question, token) = self.get_question(&history, &human_message.content).await?;
91        if let Some(token) = token {
92            token_usage = Some(token);
93        }
94
95        let documents = self
96            .retriever
97            .get_relevant_documents(&question)
98            .await
99            .map_err(|e| ChainError::RetrieverError(e.to_string()))?;
100
101        let mut output = self
102            .combine_documents_chain
103            .call(
104                StuffQAPromptBuilder::new()
105                    .documents(&documents)
106                    .question(question.clone())
107                    .build(),
108            )
109            .await?;
110
111        match &output.tokens {
112            Some(tokens) => {
113                if let Some(mut token_usage) = token_usage {
114                    token_usage.add(tokens);
115                    output.tokens = Some(token_usage)
116                }
117            }
118            None => {}
119        }
120
121        {
122            let mut memory = self.memory.lock().await;
123            memory.add_message(human_message).await?;
124            memory
125                .add_message(Message::new_ai_message(&output.generation))
126                .await?;
127        }
128
129        let mut result = HashMap::new();
130        result.insert(self.output_key.clone(), json!(output.generation));
131
132        result.insert(DEFAULT_RESULT_KEY.to_string(), json!(output));
133
134        if self.return_source_documents {
135            result.insert(
136                CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_SOURCE_DOCUMENT_KEY.to_string(),
137                json!(documents),
138            );
139        }
140
141        if self.rephrase_question {
142            result.insert(
143                CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_GENERATED_QUESTION_KEY.to_string(),
144                json!(question),
145            );
146        }
147
148        Ok(result)
149    }
150
151    async fn stream(
152        &self,
153        input_variables: PromptArgs,
154    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, ChainError>> + Send>>, ChainError>
155    {
156        let input_variable = &input_variables
157            .get(&self.input_key)
158            .ok_or(ChainError::MissingInputVariable(self.input_key.clone()))?;
159
160        let human_message =
161            Message::new_human_message(crate::schemas::messages::value_to_text(input_variable));
162        let history = self.memory.lock().await.messages().await?;
163
164        let (question, _) = self.get_question(&history, &human_message.content).await?;
165
166        let documents = self
167            .retriever
168            .get_relevant_documents(&question)
169            .await
170            .map_err(|e| ChainError::RetrieverError(e.to_string()))?;
171
172        let stream = self
173            .combine_documents_chain
174            .stream(
175                StuffQAPromptBuilder::new()
176                    .documents(&documents)
177                    .question(question.clone())
178                    .build(),
179            )
180            .await?;
181
182        let memory = self.memory.clone();
183        let complete_ai_message = Arc::new(Mutex::new(String::new()));
184        let complete_ai_message_clone = complete_ai_message.clone();
185        let output_stream = stream! {
186            pin_mut!(stream);
187            while let Some(result) = stream.next().await {
188                match result {
189                    Ok(data) => {
190                        let mut complete_ai_message_clone =
191                            complete_ai_message_clone.lock().await;
192                        complete_ai_message_clone.push_str(&data.content);
193
194                        yield Ok(data);
195                    },
196                    Err(e) => {
197                        yield Err(e);
198                    }
199                }
200            }
201
202            let mut memory = memory.lock().await;
203            let ai_message = Message::new_ai_message(&complete_ai_message.lock().await);
204            for message in [human_message, ai_message] {
205                if let Err(e) = memory.add_message(message).await {
206                    yield Err(e.into());
207                    break;
208                }
209            }
210        };
211
212        Ok(Box::pin(output_stream))
213    }
214
215    fn get_input_keys(&self) -> Vec<String> {
216        vec![self.input_key.clone()]
217    }
218
219    fn get_output_keys(&self) -> Vec<String> {
220        let mut keys = Vec::new();
221        if self.return_source_documents {
222            keys.push(CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_SOURCE_DOCUMENT_KEY.to_string());
223        }
224
225        if self.rephrase_question {
226            keys.push(CONVERSATIONAL_RETRIEVAL_QA_DEFAULT_GENERATED_QUESTION_KEY.to_string());
227        }
228
229        keys.push(self.output_key.clone());
230        keys.push(DEFAULT_RESULT_KEY.to_string());
231
232        keys
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use std::error::Error;
239
240    use crate::{
241        chain::ConversationalRetrieverChainBuilder,
242        llm::openai::{OpenAI, OpenAIModel},
243        memory::SimpleMemory,
244        prompt_args,
245        schemas::Document,
246    };
247
248    use super::*;
249
250    struct RetrieverTest {}
251    #[async_trait]
252    impl Retriever for RetrieverTest {
253        async fn get_relevant_documents(
254            &self,
255            _question: &str,
256        ) -> Result<Vec<Document>, Box<dyn Error>> {
257            Ok(vec![
258                Document::new(format!(
259                    "\nQuestion: {}\nAnswer: {}\n",
260                    "Which is the favorite text editor of luis", "Nvim"
261                )),
262                Document::new(format!(
263                    "\nQuestion: {}\nAnswer: {}\n",
264                    "How old is Luis", "24"
265                )),
266                Document::new(format!(
267                    "\nQuestion: {}\nAnswer: {}\n",
268                    "Where do luis live", "Peru"
269                )),
270                Document::new(format!(
271                    "\nQuestion: {}\nAnswer: {}\n",
272                    "Whts his favorite food", "Pan con chicharron"
273                )),
274            ])
275        }
276    }
277
278    #[tokio::test]
279    #[ignore]
280    async fn test_invoke_retriever_conversational() {
281        let llm = OpenAI::default().with_model(OpenAIModel::Gpt35.to_string());
282        let chain = ConversationalRetrieverChainBuilder::new()
283            .llm(llm)
284            .retriever(RetrieverTest {})
285            .memory(SimpleMemory::new().into())
286            .build()
287            .expect("Error building ConversationalChain");
288
289        let input_variables_first = prompt_args! {
290            "question" => "Hola",
291        };
292        // Execute the first `chain.invoke` and assert that it should succeed
293        let result_first = chain.invoke(input_variables_first).await;
294        assert!(
295            result_first.is_ok(),
296            "Error invoking LLMChain: {:?}",
297            result_first.err()
298        );
299
300        // Optionally, if you want to print the successful result, you can do so like this:
301        if let Ok(result) = result_first {
302            println!("Result: {:?}", result);
303        }
304
305        let input_variables_second = prompt_args! {
306            "question" => "Cual es la comida favorita de luis",
307        };
308        // Execute the second `chain.invoke` and assert that it should succeed
309        let result_second = chain.invoke(input_variables_second).await;
310        assert!(
311            result_second.is_ok(),
312            "Error invoking LLMChain: {:?}",
313            result_second.err()
314        );
315
316        if let Ok(result) = result_second {
317            println!("Result: {:?}", result);
318        }
319    }
320}
langchain_rust/chain/conversational_retrieval_qa/conversational_retrieval_qa.rs

langchain_rust/chain/conversational_retrieval_qa/
conversational_retrieval_qa.rs
