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::runnables::RunnableConfig;
7use lc_core::BaseChatModel;
8use lc_providers::{wrap_chat_model, ProviderError};
9use lc_schema::Message;
10use lc_shared::document::Document;
11use serde_json::Value;
12use std::collections::HashMap;
13
14use crate::base::{
15    stream_chain_with_callbacks, substitute_template, BaseChain, ChainError, ChainResult,
16    ChainStream, StreamToken,
17};
18use crate::BoxedChatModel;
19
20/// Default Stuff prompt template.
21pub(crate) const DEFAULT_STUFF_PROMPT: &str =
22    "Answer the user's question based on the following reference information.
23
24Reference information:
25{context}
26
27Question: {input}
28
29Answer:";
30
31/// StuffDocumentsChain
32///
33/// Stuffs all documents into a single prompt for LLM processing.
34/// Suitable when the total document content fits within the LLM context window.
35pub struct StuffDocumentsChain {
36    llm: BoxedChatModel,
37    prompt_template: String,
38    document_variable_name: String,
39    input_key: String,
40    output_key: String,
41    name: String,
42    verbose: bool,
43    /// Maximum character count per document (truncated if exceeded).
44    max_doc_length: Option<usize>,
45}
46
47impl StuffDocumentsChain {
48    /// Shared streaming body used by `stream` (config-less) and
49    /// `stream_with_config` (config threaded into the LLM stream).
50    async fn stream_body(
51        &self,
52        inputs: HashMap<String, Value>,
53        config: Option<RunnableConfig>,
54    ) -> Result<ChainStream, ChainError> {
55        self.validate_inputs(&inputs)?;
56
57        if config.as_ref().is_some_and(|c| c.is_cancelled()) {
58            return Err(ChainError::StreamError("Operation cancelled".to_string()));
59        }
60
61        let input = inputs
62            .get(&self.input_key)
63            .and_then(|v| v.as_str())
64            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
65
66        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
67
68        // P2-7: guard empty documents on the stream path too, mirroring the
69        // other document chains — streaming with zero context would still call
70        // the LLM and emit a fabricated, reference-free answer.
71        if documents.is_empty() {
72            return Err(ChainError::ExecutionError(
73                "Document list is empty".to_string(),
74            ));
75        }
76
77        let context = self.format_documents(&documents);
78        let prompt = self.build_prompt(&context, input);
79        let messages = vec![Message::human(&prompt)];
80
81        let llm_stream = self
82            .llm
83            .stream_chat(messages, config)
84            .await
85            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
86
87        let stream = llm_stream.map(|result| match result {
88            Ok(chunk) => Ok(StreamToken {
89                token: chunk.text,
90                is_final: false,
91            }),
92            Err(e) => Err(ChainError::StreamError(format!(
93                "Stream token error: {}",
94                e
95            ))),
96        });
97
98        let final_stream = stream.chain(futures_util::stream::once(async move {
99            Ok(StreamToken {
100                token: String::new(),
101                is_final: true,
102            })
103        }));
104
105        Ok(Box::pin(final_stream))
106    }
107
108    /// Create a new [`StuffDocumentsChain`] with the given LLM.
109    pub fn new<L>(llm: L) -> Self
110    where
111        L: BaseChatModel + Send + Sync + 'static,
112        L::Error: Into<ProviderError>,
113    {
114        Self {
115            llm: wrap_chat_model(llm),
116            prompt_template: DEFAULT_STUFF_PROMPT.to_string(),
117            document_variable_name: "context".to_string(),
118            input_key: "input".to_string(),
119            output_key: "output".to_string(),
120            name: "stuff_documents".to_string(),
121            verbose: false,
122            max_doc_length: None,
123        }
124    }
125
126    /// Set the prompt template.
127    pub fn with_prompt_template(mut self, template: impl Into<String>) -> Self {
128        self.prompt_template = template.into();
129        self
130    }
131
132    /// Set the document variable name used in the prompt.
133    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
134        self.document_variable_name = name.into();
135        self
136    }
137
138    /// Set the input key.
139    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
140        self.input_key = key.into();
141        self
142    }
143
144    /// Set the output key.
145    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
146        self.output_key = key.into();
147        self
148    }
149
150    /// Set the chain name.
151    pub fn with_name(mut self, name: impl Into<String>) -> Self {
152        self.name = name.into();
153        self
154    }
155
156    /// Set verbose mode.
157    pub fn with_verbose(mut self, verbose: bool) -> Self {
158        self.verbose = verbose;
159        self
160    }
161
162    /// Set the maximum character count per document (documents are truncated if exceeded).
163    pub fn with_max_doc_length(mut self, max: usize) -> Self {
164        self.max_doc_length = Some(max);
165        self
166    }
167
168    /// Format document list into context text.
169    pub fn format_documents(&self, documents: &[Document]) -> String {
170        let mut parts = Vec::new();
171        for (i, doc) in documents.iter().enumerate() {
172            let mut content = doc.content.clone();
173            if let Some(max_len) = self.max_doc_length {
174                let char_count: usize = content.chars().count();
175                if char_count > max_len {
176                    content = content.chars().take(max_len).collect::<String>();
177                    content.push_str("...\n[document truncated]");
178                }
179            }
180            parts.push(format!("Document {}:\n{}", i + 1, content));
181        }
182        parts.join("\n\n---\n\n")
183    }
184
185    /// Build prompt.
186    pub fn build_prompt(&self, context: &str, input: &str) -> String {
187        // 0.22.0 audit fix (H-C2): single-pass substitution — document
188        // contents containing literal `{input}` / `{context}` are never
189        // re-replaced (the old two-pass `String::replace` let them inject).
190        let vars = HashMap::from([
191            (self.document_variable_name.clone(), context.to_string()),
192            ("input".to_string(), input.to_string()),
193        ]);
194        substitute_template(&self.prompt_template, &vars).0
195    }
196
197    /// Invoke with documents and input directly.
198    pub async fn invoke_with_documents(
199        &self,
200        documents: Vec<Document>,
201        input: &str,
202    ) -> Result<String, ChainError> {
203        // P2-7: same loud empty-documents guard as map_reduce/refine/map_rerank —
204        // without it the LLM would be called with an empty context and fabricate
205        // an answer that has no reference information at all.
206        if documents.is_empty() {
207            return Err(ChainError::ExecutionError(
208                "Document list is empty".to_string(),
209            ));
210        }
211
212        let context = self.format_documents(&documents);
213
214        if self.verbose {
215            println!("\n=== StuffDocumentsChain ===");
216            println!("Document count: {}", documents.len());
217            println!("Context length: {} characters", context.len());
218        }
219
220        let prompt = self.build_prompt(&context, input);
221
222        if self.verbose {
223            println!("Prompt length: {} characters", prompt.len());
224        }
225
226        let messages = vec![Message::human(&prompt)];
227        let response = self
228            .llm
229            .invoke(messages, None)
230            .await
231            .map_err(|e| ChainError::ExecutionError(format!("LLM call failed: {}", e)))?;
232
233        let output = response.content;
234
235        if self.verbose {
236            println!("Output: {}", output);
237            println!("=== StuffDocumentsChain complete ===\n");
238        }
239
240        Ok(output)
241    }
242}
243
244#[async_trait]
245impl BaseChain for StuffDocumentsChain {
246    fn input_keys(&self) -> Vec<&str> {
247        vec![&self.input_key, "documents"]
248    }
249
250    fn output_keys(&self) -> Vec<&str> {
251        vec![&self.output_key]
252    }
253
254    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
255        // P2-8: validate inputs on the invoke path too, matching stream (and the
256        // crate-wide invoke+stream convention) so missing keys fail identically.
257        self.validate_inputs(&inputs)?;
258
259        let input = inputs
260            .get(&self.input_key)
261            .and_then(|v| v.as_str())
262            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
263
264        let documents = crate::base::documents_from_input(inputs.get("documents"))?;
265
266        let output = self.invoke_with_documents(documents, input).await?;
267
268        let mut result = HashMap::new();
269        result.insert(self.output_key.clone(), Value::String(output));
270        Ok(result)
271    }
272
273    /// Stream execution for StuffDocumentsChain — token by token output.
274    ///
275    /// Stuffs all documents into a single prompt, then streams the LLM
276    /// response token by token.
277    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
278        self.stream_body(inputs, None).await
279    }
280
281    /// Stream with config propagation.
282    ///
283    /// 0.22.0 audit fix (H-C3): thread the chain's `RunnableConfig` into
284    /// `stream_chat` (sampling overrides / cancellation token / callbacks)
285    /// instead of passing `None`.
286    async fn stream_with_config(
287        &self,
288        inputs: HashMap<String, Value>,
289        config: Option<RunnableConfig>,
290    ) -> Result<ChainStream, ChainError> {
291        let output_key = Some(self.output_key.clone());
292        stream_chain_with_callbacks(
293            self.name(),
294            inputs,
295            config.clone(),
296            output_key,
297            |inputs| async move { self.stream_body(inputs, config).await },
298        )
299        .await
300    }
301
302    fn name(&self) -> &str {
303        &self.name
304    }
305}
306
307#[cfg(test)]
308mod tests {
309    use super::*;
310    use async_trait::async_trait;
311    use futures_util::Stream;
312    use lc_core::language_models::{LLMResult, StreamChunk};
313    use lc_core::runnables::RunnableConfig;
314    use lc_core::{BaseLanguageModel, Runnable};
315    use std::pin::Pin;
316    use std::sync::atomic::{AtomicUsize, Ordering};
317    use std::sync::Arc;
318
319    /// Counting chat model that records `invoke` calls so the P2-7 typed entry
320    /// is provable: non-empty documents reach the LLM exactly once, and empty
321    /// documents never reach it at all.
322    struct CountingLLM {
323        invokes: Arc<AtomicUsize>,
324    }
325
326    #[async_trait]
327    impl Runnable<Vec<Message>, LLMResult> for CountingLLM {
328        type Error = ProviderError;
329        async fn invoke(
330            &self,
331            _input: Vec<Message>,
332            _config: Option<RunnableConfig>,
333        ) -> Result<LLMResult, Self::Error> {
334            self.invokes.fetch_add(1, Ordering::SeqCst);
335            Ok(LLMResult {
336                content: "stuffed answer".to_string(),
337                model: "mock".to_string(),
338                token_usage: None,
339                tool_calls: None,
340                thinking_content: None,
341            })
342        }
343    }
344
345    #[async_trait]
346    impl BaseLanguageModel<Vec<Message>, LLMResult> for CountingLLM {
347        fn model_name(&self) -> &str {
348            "mock"
349        }
350        fn get_num_tokens(&self, t: &str) -> usize {
351            t.len()
352        }
353        fn with_temperature(self, _: f32) -> Self {
354            self
355        }
356        fn with_max_tokens(self, _: usize) -> Self {
357            self
358        }
359    }
360
361    #[async_trait]
362    impl BaseChatModel for CountingLLM {
363        async fn chat(
364            &self,
365            _messages: Vec<Message>,
366            _config: Option<RunnableConfig>,
367        ) -> Result<LLMResult, Self::Error> {
368            self.invokes.fetch_add(1, Ordering::SeqCst);
369            Ok(LLMResult {
370                content: "stuffed answer".to_string(),
371                model: "mock".to_string(),
372                token_usage: None,
373                tool_calls: None,
374                thinking_content: None,
375            })
376        }
377        async fn stream_chat(
378            &self,
379            _messages: Vec<Message>,
380            _config: Option<RunnableConfig>,
381        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
382        {
383            let tokens = [Ok(StreamChunk::new("streamed answer"))];
384            Ok(Box::pin(futures_util::stream::iter(tokens)))
385        }
386    }
387
388    /// P2-7: the typed `invoke_with_documents` entry skips the HashMap roundtrip
389    /// and runs the LLM over the documents directly.
390    #[tokio::test]
391    async fn test_stuff_typed_invoke_runs_llm() {
392        let invokes = Arc::new(AtomicUsize::new(0));
393        let chain = StuffDocumentsChain::new(CountingLLM {
394            invokes: invokes.clone(),
395        });
396
397        let out = chain
398            .invoke_with_documents(vec![Document::new("doc one")], "question")
399            .await
400            .unwrap();
401        assert_eq!(out, "stuffed answer");
402        assert_eq!(invokes.load(Ordering::SeqCst), 1);
403    }
404
405    /// P2-7: empty documents error loudly on the typed invoke path and the LLM
406    /// is never called with an empty context (consistent with
407    /// map_reduce/refine/map_rerank).
408    #[tokio::test]
409    async fn test_stuff_typed_invoke_empty_documents_errors() {
410        let invokes = Arc::new(AtomicUsize::new(0));
411        let chain = StuffDocumentsChain::new(CountingLLM {
412            invokes: invokes.clone(),
413        });
414
415        let err = match chain.invoke_with_documents(vec![], "q").await {
416            Ok(_) => panic!("expected an execution error"),
417            Err(e) => e,
418        };
419        assert!(matches!(err, ChainError::ExecutionError(_)));
420        assert_eq!(invokes.load(Ordering::SeqCst), 0);
421    }
422
423    /// P2-8: the HashMap `invoke` path validates inputs before any LLM call,
424    /// matching the stream path — a missing `input` key yields `MissingInput`
425    /// with zero invokes, instead of reaching the LLM first.
426    #[tokio::test]
427    async fn test_stuff_invoke_missing_input_errors_before_llm() {
428        let invokes = Arc::new(AtomicUsize::new(0));
429        let chain = StuffDocumentsChain::new(CountingLLM {
430            invokes: invokes.clone(),
431        });
432        let mut inputs = HashMap::new();
433        inputs.insert("documents".to_string(), serde_json::json!([]));
434
435        let err = match chain.invoke(inputs).await {
436            Ok(_) => panic!("expected a missing-input error"),
437            Err(e) => e,
438        };
439        assert!(matches!(err, ChainError::MissingInput(_)));
440        assert_eq!(invokes.load(Ordering::SeqCst), 0);
441    }
442
443    /// P2-7: the stream path guards empty documents the same way, erroring
444    /// before any `stream_chat` call.
445    #[tokio::test]
446    async fn test_stuff_stream_empty_documents_errors() {
447        let chain = StuffDocumentsChain::new(CountingLLM {
448            invokes: Arc::new(AtomicUsize::new(0)),
449        });
450        let mut inputs = HashMap::new();
451        inputs.insert("input".to_string(), Value::String("q".to_string()));
452        inputs.insert("documents".to_string(), serde_json::json!([]));
453
454        let err = match chain.stream(inputs).await {
455            Ok(_) => panic!("expected an execution error"),
456            Err(e) => e,
457        };
458        assert!(matches!(err, ChainError::ExecutionError(_)));
459    }
460}
lc_chains/document_chains/stuff.rs

lc_chains/document_chains/
stuff.rs
