langchain_rust/chain/stuff_documents/
chain.rs

1use std::pin::Pin;
2
3use async_trait::async_trait;
4use futures::Stream;
5use serde_json::Value;
6
7use crate::{
8    chain::{
9        load_stuff_qa, options::ChainCallOptions, Chain, ChainError, ChainEventStream, LLMChain,
10        StuffQAPromptBuilder,
11    },
12    language_models::{llm::LLM, GenerateResult},
13    prompt::PromptArgs,
14    schemas::{Document, StreamData},
15};
16
17pub(crate) const COMBINE_DOCUMENTS_DEFAULT_INPUT_KEY: &str = "input_documents";
18const COMBINE_DOCUMENTS_DEFAULT_OUTPUT_KEY: &str = "text";
19pub(crate) const COMBINE_DOCUMENTS_DEFAULT_DOCUMENT_VARIABLE_NAME: &str = "context";
20const STUFF_DOCUMENTS_DEFAULT_SEPARATOR: &str = "\n\n";
21
22/// Reads the documents to combine from `input_variables[key]`.
23pub(crate) fn input_documents(
24    input_variables: &PromptArgs,
25    key: &str,
26) -> Result<Vec<Document>, ChainError> {
27    let docs = input_variables
28        .get(key)
29        .ok_or_else(|| ChainError::MissingInputVariable(key.to_string()))?;
30    serde_json::from_value(docs.clone()).map_err(|e| ChainError::IncorrectInputVariable {
31        source: e,
32        expected_type: "Vec<Document>".to_string(),
33    })
34}
35
36pub struct StuffDocument {
37    llm_chain: LLMChain,
38    input_key: String,
39    document_variable_name: String,
40    separator: String,
41}
42
43impl StuffDocument {
44    pub fn new(llm_chain: LLMChain) -> Self {
45        Self {
46            llm_chain,
47            input_key: COMBINE_DOCUMENTS_DEFAULT_INPUT_KEY.to_string(),
48            document_variable_name: COMBINE_DOCUMENTS_DEFAULT_DOCUMENT_VARIABLE_NAME.to_string(),
49            separator: STUFF_DOCUMENTS_DEFAULT_SEPARATOR.to_string(),
50        }
51    }
52
53    fn join_documents(&self, docs: Vec<Document>) -> String {
54        docs.iter()
55            .map(|doc| doc.page_content.clone())
56            .collect::<Vec<_>>()
57            .join(&self.separator)
58    }
59
60    ///Inly use thi if you use the deafult prompt
61    pub fn qa_prompt_builder<'a>(&self) -> StuffQAPromptBuilder<'a> {
62        StuffQAPromptBuilder::new()
63    }
64
65    /// load_stuff_qa return an instance of StuffDocument
66    /// with a prompt desiged for question ansering
67    ///
68    /// # Example
69    /// ```rust,ignore
70    ///
71    /// let llm = OpenAI::default();
72    /// let chain = StuffDocument::load_stuff_qa(llm);
73    ///
74    /// let input = chain
75    /// .qa_prompt_builder()
76    /// .documents(&[
77    /// Document::new(format!(
78    /// "\nQuestion: {}\nAnswer: {}\n",
79    /// "Which is the favorite text editor of luis", "Nvim"
80    /// )),
81    /// Document::new(format!(
82    /// "\nQuestion: {}\nAnswer: {}\n",
83    /// "How old is Luis", "24"
84    /// )),
85    /// ])
86    /// .question("How old is luis and whats his favorite text editor")
87    /// .build();
88    ///
89    /// let ouput = chain.invoke(input).await.unwrap();
90    ///
91    /// println!("{}", ouput);
92    /// ```
93    ///
94    pub fn load_stuff_qa<L: Into<Box<dyn LLM>>>(llm: L) -> Self {
95        load_stuff_qa(llm, None)
96    }
97
98    /// load_stuff_qa_with_options return an instance of StuffDocument
99    /// with a prompt desiged for question ansering
100    ///
101    /// # Example
102    /// ```rust,ignore
103    ///
104    /// let llm = OpenAI::default();
105    /// let chain = StuffDocument::load_stuff_qa_with_options(llm,ChainCallOptions::default());
106    ///
107    /// let input = chain
108    /// .qa_prompt_builder()
109    /// .documents(&[
110    /// Document::new(format!(
111    /// "\nQuestion: {}\nAnswer: {}\n",
112    /// "Which is the favorite text editor of luis", "Nvim"
113    /// )),
114    /// Document::new(format!(
115    /// "\nQuestion: {}\nAnswer: {}\n",
116    /// "How old is Luis", "24"
117    /// )),
118    /// ])
119    /// .question("How old is luis and whats his favorite text editor")
120    /// .build();
121    ///
122    /// let ouput = chain.invoke(input).await.unwrap();
123    ///
124    /// println!("{}", ouput);
125    /// ```
126    ///
127    pub fn load_stuff_qa_with_options<L: LLM + 'static>(llm: L, opt: ChainCallOptions) -> Self {
128        load_stuff_qa(llm, Some(opt))
129    }
130}
131
132#[async_trait]
133impl Chain for StuffDocument {
134    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError> {
135        let input_values = self.stuffed_inputs(input_variables)?;
136        self.llm_chain.call(input_values).await
137    }
138
139    async fn stream(
140        &self,
141        input_variables: PromptArgs,
142    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, ChainError>> + Send>>, ChainError>
143    {
144        let input_values = self.stuffed_inputs(input_variables)?;
145        self.llm_chain.stream(input_values).await
146    }
147
148    fn call_stream<'a>(&'a self, input_variables: PromptArgs) -> ChainEventStream<'a> {
149        match self.stuffed_inputs(input_variables) {
150            Ok(input_values) => self.llm_chain.call_stream(input_values),
151            Err(e) => Box::pin(futures::stream::iter([Err(e)])),
152        }
153    }
154
155    fn get_input_keys(&self) -> Vec<String> {
156        vec![self.input_key.clone()]
157    }
158}
159
160impl StuffDocument {
161    /// The input variables plus the documents joined into `document_variable_name`.
162    fn stuffed_inputs(&self, input_variables: PromptArgs) -> Result<PromptArgs, ChainError> {
163        let documents = input_documents(&input_variables, &self.input_key)?;
164        let mut input_values = input_variables;
165        input_values.insert(
166            self.document_variable_name.clone(),
167            Value::String(self.join_documents(documents)),
168        );
169        Ok(input_values)
170    }
171}
langchain_rust/chain/stuff_documents/chain.rs

langchain_rust/chain/stuff_documents/
chain.rs
