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

lc_chains/document_chains/
stuff.rs
