lc_chains/document_chains/
refine.rs

1// lc-chains/src/document_chains/refine.rs
2//! RefineDocumentsChain - iteratively refines the answer document by document.
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 initial processing prompt template.
16pub(crate) const DEFAULT_REFINE_INITIAL_PROMPT: &str =
17    "Answer the question based on the following reference information.
18
19Reference information:
20{context}
21
22Question: {input}
23
24Answer:";
25
26/// Default iterative refinement prompt template.
27pub(crate) const DEFAULT_REFINE_PROMPT: &str = "You have provided an answer based on partial information. Here is additional reference information.
28
29Existing answer:
30{existing_answer}
31
32New reference information:
33{context}
34
35Please refine or modify your answer based on the new information. If the new information does not conflict with the existing answer, merge them. If the new information conflicts with the existing answer, prioritize the new information.
36
37Question: {input}
38
39Refined answer:";
40
41/// RefineDocumentsChain
42///
43/// Iteratively refines the answer document by document.
44/// Generates an initial answer from the first document, then refines with each subsequent document.
45pub struct RefineDocumentsChain<M: BaseChatModel> {
46    llm: M,
47    initial_prompt_template: String,
48    refine_prompt_template: String,
49    document_variable_name: String,
50    input_key: String,
51    output_key: String,
52    name: String,
53    verbose: bool,
54}
55
56impl<M: BaseChatModel> RefineDocumentsChain<M> {
57    pub fn new(llm: M) -> Self {
58        Self {
59            llm,
60            initial_prompt_template: DEFAULT_REFINE_INITIAL_PROMPT.to_string(),
61            refine_prompt_template: DEFAULT_REFINE_PROMPT.to_string(),
62            document_variable_name: "context".to_string(),
63            input_key: "input".to_string(),
64            output_key: "output".to_string(),
65            name: "refine_documents".to_string(),
66            verbose: false,
67        }
68    }
69
70    pub fn with_initial_prompt(mut self, template: impl Into<String>) -> Self {
71        self.initial_prompt_template = template.into();
72        self
73    }
74
75    pub fn with_refine_prompt(mut self, template: impl Into<String>) -> Self {
76        self.refine_prompt_template = template.into();
77        self
78    }
79
80    pub fn with_document_variable(mut self, name: impl Into<String>) -> Self {
81        self.document_variable_name = name.into();
82        self
83    }
84
85    pub fn with_input_key(mut self, key: impl Into<String>) -> Self {
86        self.input_key = key.into();
87        self
88    }
89
90    pub fn with_output_key(mut self, key: impl Into<String>) -> Self {
91        self.output_key = key.into();
92        self
93    }
94
95    pub fn with_name(mut self, name: impl Into<String>) -> Self {
96        self.name = name.into();
97        self
98    }
99
100    pub fn with_verbose(mut self, verbose: bool) -> Self {
101        self.verbose = verbose;
102        self
103    }
104
105    pub fn build_initial_prompt(&self, context: &str, input: &str) -> String {
106        self.initial_prompt_template
107            .replace(&format!("{{{}}}", self.document_variable_name), context)
108            .replace("{input}", input)
109    }
110
111    pub fn build_refine_prompt(&self, context: &str, input: &str, existing_answer: &str) -> String {
112        self.refine_prompt_template
113            .replace(&format!("{{{}}}", self.document_variable_name), context)
114            .replace("{input}", input)
115            .replace("{existing_answer}", existing_answer)
116    }
117
118    /// Invoke with documents and input directly (iterative refinement).
119    pub async fn invoke_with_documents(
120        &self,
121        documents: Vec<Document>,
122        input: &str,
123    ) -> Result<String, ChainError>
124    where
125        <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
126    {
127        if documents.is_empty() {
128            return Err(ChainError::ExecutionError(
129                "Document list is empty".to_string(),
130            ));
131        }
132
133        if self.verbose {
134            println!("\n=== RefineDocumentsChain ===");
135            println!("Document count: {}", documents.len());
136            println!("Input: {}", input);
137        }
138
139        // Step 1: Generate initial answer from the first document
140        let first_context = &documents[0].content;
141        let initial_prompt = self.build_initial_prompt(first_context, input);
142
143        if self.verbose {
144            println!("\n--- Initial processing (document 1) ---");
145        }
146
147        let messages = vec![Message::human(&initial_prompt)];
148        let response =
149            self.llm.invoke(messages, None).await.map_err(|e| {
150                ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
151            })?;
152        let mut answer = response.content;
153
154        if self.verbose {
155            println!("Initial answer: {}", answer);
156        }
157
158        // Subsequent steps: iteratively refine with remaining documents
159        for (i, doc) in documents[1..].iter().enumerate() {
160            if self.verbose {
161                println!("\n--- Refinement step {} (document {}) ---", i + 1, i + 2);
162            }
163
164            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
165
166            let messages = vec![Message::human(&refine_prompt)];
167            let response = self.llm.invoke(messages, None).await.map_err(|e| {
168                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
169            })?;
170            answer = response.content;
171
172            if self.verbose {
173                println!("Refined answer: {}", answer);
174            }
175        }
176
177        if self.verbose {
178            println!("=== RefineDocumentsChain complete ===\n");
179        }
180
181        Ok(answer)
182    }
183}
184
185#[async_trait]
186impl<M: BaseChatModel + Send + Sync + 'static> BaseChain for RefineDocumentsChain<M>
187where
188    <M as Runnable<Vec<Message>, LLMResult>>::Error: std::fmt::Display,
189{
190    fn input_keys(&self) -> Vec<&str> {
191        vec![&self.input_key, "documents"]
192    }
193
194    fn output_keys(&self) -> Vec<&str> {
195        vec![&self.output_key]
196    }
197
198    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
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 output = self.invoke_with_documents(documents, input).await?;
215
216        let mut result = HashMap::new();
217        result.insert(self.output_key.clone(), Value::String(output));
218        Ok(result)
219    }
220
221    /// Stream execution for RefineDocumentsChain.
222    ///
223    /// Runs the initial + all intermediate refine steps via invoke (since
224    /// their output feeds the next step), then streams the final refine step
225    /// token by token via `stream_chat`.
226    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
227        self.validate_inputs(&inputs)?;
228
229        let input = inputs
230            .get(&self.input_key)
231            .and_then(|v| v.as_str())
232            .ok_or_else(|| ChainError::MissingInput(self.input_key.clone()))?;
233
234        let documents: Vec<Document> = inputs
235            .get("documents")
236            .and_then(|v| v.as_array())
237            .map(|arr| {
238                arr.iter()
239                    .filter_map(|v| serde_json::from_value(v.clone()).ok())
240                    .collect()
241            })
242            .ok_or_else(|| ChainError::MissingInput("documents".to_string()))?;
243
244        if documents.is_empty() {
245            return Err(ChainError::ExecutionError(
246                "Document list is empty".to_string(),
247            ));
248        }
249
250        // Step 1: Generate initial answer from the first document
251        let first_context = &documents[0].content;
252        let initial_prompt = self.build_initial_prompt(first_context, input);
253        let messages = vec![Message::human(&initial_prompt)];
254        let response = self.llm.invoke(messages, None).await.map_err(|e| {
255            ChainError::ExecutionError(format!("LLM initial call failed: {}", e))
256        })?;
257        let mut answer = response.content;
258
259        // Step 2: Run intermediate refine steps (all but the last) via invoke
260        let last_idx = documents.len() - 1;
261        for (i, doc) in documents[1..last_idx].iter().enumerate() {
262            let refine_prompt = self.build_refine_prompt(&doc.content, input, &answer);
263            let messages = vec![Message::human(&refine_prompt)];
264            let response = self.llm.invoke(messages, None).await.map_err(|e| {
265                ChainError::ExecutionError(format!("LLM refinement call failed: {}", e))
266            })?;
267            answer = response.content;
268
269            if self.verbose {
270                println!("Refine step {} completed", i + 1);
271            }
272        }
273
274        // Step 3: Stream the final refine step
275        let final_prompt = if last_idx == 0 {
276            // Only one document — stream the initial answer
277            self.build_initial_prompt(&documents[0].content, input)
278        } else {
279            self.build_refine_prompt(&documents[last_idx].content, input, &answer)
280        };
281
282        let messages = vec![Message::human(&final_prompt)];
283        let llm_stream = self
284            .llm
285            .stream_chat(messages, None)
286            .await
287            .map_err(|e| ChainError::StreamError(format!("LLM stream failed: {}", e)))?;
288
289        let stream = llm_stream.map(|result| match result {
290            Ok(token) => Ok(StreamToken {
291                token,
292                is_final: false,
293            }),
294            Err(e) => Err(ChainError::StreamError(format!(
295                "Stream token error: {}",
296                e
297            ))),
298        });
299
300        let final_stream =
301            stream.chain(futures_util::stream::once(async move {
302                Ok(StreamToken {
303                    token: String::new(),
304                    is_final: true,
305                })
306            }));
307
308        Ok(Box::pin(final_stream))
309    }
310
311    fn name(&self) -> &str {
312        &self.name
313    }
314}
lc_chains/document_chains/refine.rs

lc_chains/document_chains/
refine.rs
