lc_chains/document_chains/
stuff.rs

1// lc-chains/src/document_chains/stuff.rs
2//! StuffDocumentsChain - stuffs all documents into a single prompt.
3
4use async_trait::async_trait;
5use lc_core::language_models::LLMResult;
6use lc_core::{BaseChatModel, Runnable};
7use lc_schema::Message;
8use lc_shared::document::Document;
9use serde_json::Value;
10use std::collections::HashMap;
11
12use crate::base::{BaseChain, ChainError, ChainResult};
13
14/// Default Stuff prompt template.
15pub(crate) const DEFAULT_STUFF_PROMPT: &str =
16    "Answer the user's question based on the following reference information.
17
18Reference information:
19{context}
20
21Question: {input}
22
23Answer:";
24
25/// StuffDocumentsChain
26///
27/// Stuffs all documents into a single prompt for LLM processing.
28/// Suitable when the total document content fits within the LLM context window.
29pub struct StuffDocumentsChain<M: BaseChatModel> {
30    llm: M,
31    prompt_template: String,
32    document_variable_name: String,
33    input_key: String,
34    output_key: String,
35    name: String,
36    verbose: bool,
37    /// Maximum character count per document (truncated if exceeded).
38    max_doc_length: Option<usize>,
39}
40
41impl<M: BaseChatModel> StuffDocumentsChain<M> {
42    pub fn new(llm: M) -> Self {
43        Self {
44            llm,
45            prompt_template: DEFAULT_STUFF_PROMPT.to_string(),
46            document_variable_name: "context".to_string(),
47            input_key: "input".to_string(),
48            output_key: "output".to_string(),
49            name: "stuff_documents".to_string(),
50            verbose: false,
51            max_doc_length: None,
52        }
53    }
54
55    pub fn with_prompt_template(mut self, template: impl Into<String>) -> Self {
56        self.prompt_template = template.into();
57        self
58    }
59
60    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
61        self.document_variable_name = name.into();
62        self
63    }
64
65    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
66        self.input_key = key.into();
67        self
68    }
69
70    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
71        self.output_key = key.into();
72        self
73    }
74
75    pub fn with_name(mut self, name: impl Into<String>) -> Self {
76        self.name = name.into();
77        self
78    }
79
80    pub fn with_verbose(mut self, verbose: bool) -> Self {
81        self.verbose = verbose;
82        self
83    }
84
85    pub fn with_max_doc_length(mut self, max: usize) -> Self {
86        self.max_doc_length = Some(max);
87        self
88    }
89
90    /// Format document list into context text.
91    pub fn format_documents(&self, documents: &[Document]) -> String {
92        let mut parts = Vec::new();
93        for (i, doc) in documents.iter().enumerate() {
94            let mut content = doc.content.clone();
95            if let Some(max_len) = self.max_doc_length {
96                let char_count: usize = content.chars().count();
97                if char_count > max_len {
98                    content = content.chars().take(max_len).collect::<String>();
99                    content.push_str("...\n[document truncated]");
100                }
101            }
102            parts.push(format!("Document {}:\n{}", i + 1, content));
103        }
104        parts.join("\n\n---\n\n")
105    }
106
107    /// Build prompt.
108    pub fn build_prompt(&self, context: &str, input: &str) -> String {
109        self.prompt_template
110            .replace(&format!("{{{}}}", self.document_variable_name), context)
111            .replace("{input}", input)
112    }
113
114    /// Invoke with documents and input directly.
115    pub async fn invoke_with_documents(
116        &self,
117        documents: Vec<Document>,
118        input: &str,
119    ) -> Result<String, ChainError>
120    where
121        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
122    {
123        let context = self.format_documents(&documents);
124
125        if self.verbose {
126            println!("\n=== StuffDocumentsChain ===");
127            println!("Document count: {}", documents.len());
128            println!("Context length: {} characters", context.len());
129        }
130
131        let prompt = self.build_prompt(&context, input);
132
133        if self.verbose {
134            println!("Prompt length: {} characters", prompt.len());
135        }
136
137        let messages = vec![Message::human(&prompt)];
138        let response = self
139            .llm
140            .invoke(messages, None)
141            .await
142            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
143
144        let output = response.content;
145
146        if self.verbose {
147            println!("Output: {}", output);
148            println!("=== StuffDocumentsChain complete ===\n");
149        }
150
151        Ok(output)
152    }
153}
154
155#[async_trait]
156impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for StuffDocumentsChain<M>
157where
158    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
159{
160    fn input_keys(&self) -> Vec<&str> {
161        vec![&self.input_key, "documents"]
162    }
163
164    fn output_keys(&self) -> Vec<&str> {
165        vec![&self.output_key]
166    }
167
168    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
169        let input = inputs
170            .get(&self.input_key)
171            .and_then(|v| v.as_str())
172            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
173
174        let documents: Vec<Document> = inputs
175            .get("documents")
176            .and_then(|v| v.as_array())
177            .map(|arr| {
178                arr.iter()
179                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
180                    .collect()
181            })
182            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
183
184        let output = self.invoke_with_documents(documents, input).await?;
185
186        let mut result = HashMap::new();
187        result.insert(self.output_key.clone(), Value::String(output));
188        Ok(result)
189    }
190
191    fn name(&self) -> &str {
192        &self.name
193    }
194}
lc_chains/document_chains/stuff.rs

lc_chains/document_chains/
stuff.rs
