langchain_rust/chain/
question_answering.rs

1use std::pin::Pin;
2
3use async_trait::async_trait;
4use futures::Stream;
5
6use crate::{
7    language_models::{llm::LLM, GenerateResult},
8    prompt::PromptArgs,
9    prompt_args,
10    schemas::{messages::Message, Document, StreamData},
11    template_jinja2,
12};
13
14use super::{
15    options::ChainCallOptions, Chain, ChainError, LLMChain, LLMChainBuilder, StuffDocument,
16};
17
18const DEFAULTCONDENSEQUESTIONTEMPLATE: &str = r#"Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
19
20Chat History:
21{{chat_history}}
22Follow Up Input: {{question}}
23Standalone question:"#;
24
25pub struct CondenseQuestionPromptBuilder {
26    chat_history: String,
27    question: String,
28}
29impl CondenseQuestionPromptBuilder {
30    pub fn new() -> Self {
31        Self {
32            chat_history: "".to_string(),
33            question: "".to_string(),
34        }
35    }
36
37    pub fn question<S: Into<String>>(mut self, question: S) -> Self {
38        self.question = question.into();
39        self
40    }
41
42    pub fn chat_history(mut self, chat_history: &[Message]) -> Self {
43        self.chat_history = Message::messages_to_string(chat_history);
44        self
45    }
46
47    pub fn build(self) -> PromptArgs {
48        prompt_args! {
49            "chat_history" => self.chat_history,
50            "question" => self.question
51        }
52    }
53}
54
55pub struct CondenseQuestionGeneratorChain {
56    chain: LLMChain,
57}
58
59impl CondenseQuestionGeneratorChain {
60    pub fn new<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
61        let condense_question_prompt_template =
62            template_jinja2!(DEFAULTCONDENSEQUESTIONTEMPLATE, "chat_history", "question");
63
64        let chain = LLMChainBuilder::new()
65            .llm(llm)
66            .prompt(condense_question_prompt_template)
67            .build()
68            .unwrap(); //Its safe to unwrap here because we are sure that the prompt and the LLM are
69                       //set.
70        Self { chain }
71    }
72
73    pub fn prompt_builder(&self) -> CondenseQuestionPromptBuilder {
74        CondenseQuestionPromptBuilder::new()
75    }
76}
77
78#[async_trait]
79impl Chain for CondenseQuestionGeneratorChain {
80    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
81        self.chain.call(input_variables).await
82    }
83
84    async fn stream(
85        &self,
86        input_variables: PromptArgs,
87    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, ChainError>> + Send>>, ChainError>
88    {
89        self.chain.stream(input_variables).await
90    }
91}
92
93const DEFAULT_STUFF_QA_TEMPLATE: &str = crate::chain::combine_documents::STUFF_QA_TEMPLATE;
94
95pub struct StuffQAPromptBuilder<'a> {
96    input_documents: Vec<&'a Document>,
97    question: String,
98}
99
100impl<'a> StuffQAPromptBuilder<'a> {
101    pub fn new() -> Self {
102        Self {
103            input_documents: vec![],
104            question: "".to_string(),
105        }
106    }
107
108    pub fn documents(mut self, documents: &'a [Document]) -> Self {
109        self.input_documents = documents.iter().collect();
110        self
111    }
112
113    pub fn question<S: Into<String>>(mut self, question: S) -> Self {
114        self.question = question.into();
115        self
116    }
117
118    pub fn build(self) -> PromptArgs {
119        prompt_args! {
120            "input_documents" => self.input_documents,
121            "question" => self.question
122        }
123    }
124}
125
126pub(crate) fn load_stuff_qa<L: Into<Box<dyn LLM>>>(
127    llm: L,
128    options: Option<ChainCallOptions>,
129) -> StuffDocument {
130    let default_qa_prompt_template =
131        template_jinja2!(DEFAULT_STUFF_QA_TEMPLATE, "context", "question");
132
133    let llm_chain_builder = LLMChainBuilder::new()
134        .prompt(default_qa_prompt_template)
135        .options(options.unwrap_or_default())
136        .llm(llm)
137        .build()
138        .unwrap();
139
140    let llm_chain = llm_chain_builder;
141
142    StuffDocument::new(llm_chain)
143}
144
145#[cfg(test)]
146mod tests {
147    use crate::{
148        chain::{Chain, StuffDocument},
149        llm::openai::OpenAI,
150        schemas::Document,
151    };
152
153    #[tokio::test]
154    #[ignore]
155    async fn test_qa() {
156        let llm = OpenAI::default();
157        let chain = StuffDocument::load_stuff_qa(llm);
158        let input = chain
159            .qa_prompt_builder()
160            .documents(&[
161                Document::new(format!(
162                    "\nQuestion: {}\nAnswer: {}\n",
163                    "Which is the favorite text editor of luis", "Nvim"
164                )),
165                Document::new(format!(
166                    "\nQuestion: {}\nAnswer: {}\n",
167                    "How old is Luis", "24"
168                )),
169            ])
170            .question("How old is luis and whats his favorite text editor")
171            .build();
172
173        let ouput = chain.invoke(input).await.unwrap();
174
175        println!("{}", ouput);
176    }
177}
langchain_rust/chain/question_answering.rs

langchain_rust/chain/
question_answering.rs
