lc_rag/
contextual.rs

1// lc-rag/src/contextual.rs
2//! Contextual Retrieval: index-time context injection for chunks (0.21.0 S4.1).
3//!
4//! Chunks lose their context when split out of a document — an ambiguous chunk
5//! ("revenue grew 3%") may embed without the subject it refers to and never be
6//! recalled for the right query. Contextual Retrieval (Anthropic, ~49% fewer
7//! retrieval failures) asks a small LLM to write 1-2 situating sentences per
8//! chunk at index time and embeds `context + original content`. The context is
9//! also stored in metadata, so it stays auditable and searchable.
10//!
11//! Design:
12//! - pure transform: takes already-split documents, returns enhanced copies —
13//!   the caller inserts it before `index_documents`; retrieval flow unchanged;
14//! - bounded concurrency ([`tokio::sync::Semaphore`]) to cap index cost;
15//! - **fail-open**: a failed context generation logs a warning and keeps the
16//!   original content — indexing must never block on the enhancement;
17//! - idempotent: chunks already enhanced (metadata marker present) are skipped,
18//!   so re-running over an indexed corpus is a no-op.
19
20use lc_core::language_models::BaseChatModel;
21use lc_prompts::PromptTemplate;
22use lc_providers::ProviderError;
23use lc_schema::Message;
24use lc_vector_stores::Document;
25use std::collections::HashMap;
26use std::sync::Arc;
27use tokio::sync::Semaphore;
28
29/// Metadata key under which the generated context is stored (auditable).
30pub const CONTEXTUAL_METADATA_KEY: &str = "contextual_context";
31
32/// Contextual Retrieval error type.
33#[derive(Debug)]
34#[non_exhaustive]
35pub enum ContextualError {
36    /// LLM call error.
37    LLMError(String),
38}
39
40impl std::fmt::Display for ContextualError {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            ContextualError::LLMError(msg) => write!(f, "LLM error: {}", msg),
44        }
45    }
46}
47
48impl std::error::Error for ContextualError {}
49
50/// Contextual Retrieval configuration.
51pub struct ContextualConfig {
52    /// Prompt template; `{chunk}` is replaced with the chunk text.
53    pub prompt_template: String,
54    /// Max LLM calls in flight (index-cost cap).
55    pub max_concurrency: usize,
56}
57
58impl Default for ContextualConfig {
59    fn default() -> Self {
60        Self {
61            prompt_template: DEFAULT_CONTEXTUAL_PROMPT.to_string(),
62            max_concurrency: 4,
63        }
64    }
65}
66
67impl ContextualConfig {
68    /// Creates a config with defaults.
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    /// Sets the prompt template (`{chunk}` placeholder required).
74    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
75        self.prompt_template = prompt.into();
76        self
77    }
78
79    /// Sets the max in-flight context-generation calls.
80    pub fn with_max_concurrency(mut self, max_concurrency: usize) -> Self {
81        self.max_concurrency = max_concurrency.max(1);
82        self
83    }
84}
85
86const DEFAULT_CONTEXTUAL_PROMPT: &str = r#"You are preparing document chunks for a retrieval index. Write 1-2 short sentences that situate the chunk within its broader document: what the document is about and what part this chunk plays. Answer with the context sentences only — no preamble, no quotes.
87
88Chunk:
89{chunk}
90
91Context:"#;
92
93/// Contextual Retrieval enhancer (index-time transform).
94///
95/// `L` is any [`BaseChatModel`]; a cheap small model is recommended (one call
96/// per chunk). Enhancement failures are logged and degrade to the original
97/// content — see the module docs.
98pub struct ContextualEnhancer {
99    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
100    config: ContextualConfig,
101    /// Bounds in-flight LLM calls across all `enhance_documents` calls on
102    /// clones of this enhancer (the Arc is shared through clones).
103    semaphore: Arc<Semaphore>,
104}
105
106impl std::fmt::Debug for ContextualEnhancer {
107    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        f.debug_struct("ContextualEnhancer")
109            .field("max_concurrency", &self.config.max_concurrency)
110            .finish()
111    }
112}
113
114impl ContextualEnhancer {
115    /// Creates an enhancer from any [`BaseChatModel`].
116    pub fn new<L>(llm: L) -> Self
117    where
118        L: BaseChatModel + Send + Sync + 'static,
119        L::Error: Into<ProviderError>,
120    {
121        Self::new_arc(lc_providers::wrap_chat_model(llm))
122    }
123
124    /// Builds from an already-wrapped `Arc<dyn BaseChatModel<Error = ProviderError>>`.
125    pub fn new_arc(llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>) -> Self {
126        let semaphore = Arc::new(Semaphore::new(4));
127        Self {
128            llm,
129            config: ContextualConfig::default(),
130            semaphore,
131        }
132    }
133
134    /// Sets the configuration.
135    pub fn with_config(mut self, config: ContextualConfig) -> Self {
136        self.semaphore = Arc::new(Semaphore::new(config.max_concurrency));
137        self.config = config;
138        self
139    }
140
141    /// Sets the max in-flight LLM calls.
142    pub fn with_max_concurrency(mut self, max_concurrency: usize) -> Self {
143        self.config.max_concurrency = max_concurrency.max(1);
144        self.semaphore = Arc::new(Semaphore::new(max_concurrency.max(1)));
145        self
146    }
147
148    /// Generates the situating context for one chunk (raw LLM access).
149    pub async fn generate_context(&self, chunk: &str) -> Result<String, ContextualError> {
150        let template = PromptTemplate::new(&self.config.prompt_template);
151        let mut vars = HashMap::new();
152        vars.insert("chunk", chunk);
153        let prompt = template
154            .format(&vars)
155            .unwrap_or_else(|_| self.config.prompt_template.clone());
156
157        let response = self
158            .llm
159            .invoke(vec![Message::human(prompt)], None)
160            .await
161            .map_err(|e| ContextualError::LLMError(e.to_string()))?;
162        Ok(response.content.trim().to_string())
163    }
164
165    /// Enhances one document: returns a copy whose content is
166    /// `context + "\n" + original` and whose metadata carries the context under
167    /// [`CONTEXTUAL_METADATA_KEY`].
168    ///
169    /// Errors surface to the caller (use [`Self::enhance_documents`] for the
170    /// fail-open batch path).
171    pub async fn enhance_document(&self, doc: &Document) -> Result<Document, ContextualError> {
172        let context = self.generate_context(&doc.content).await?;
173        let mut enhanced = Document::new(format!("{}\n{}", context, doc.content))
174            .with_metadata(CONTEXTUAL_METADATA_KEY, context.clone());
175        // Preserve id and all existing metadata.
176        if let Some(id) = &doc.id {
177            enhanced = enhanced.with_id(id.clone());
178        }
179        for (k, v) in &doc.metadata {
180            enhanced.metadata.insert(k.clone(), v.clone());
181        }
182        Ok(enhanced)
183    }
184
185    /// Batch enhancement with bounded concurrency and fail-open semantics.
186    ///
187    /// Each chunk is enhanced concurrently (up to `max_concurrency` in flight);
188    /// a chunk whose context generation fails keeps its original content (a
189    /// warning is logged) — indexing never blocks on the enhancement.
190    /// Idempotent: chunks already carrying [`CONTEXTUAL_METADATA_KEY`] are
191    /// passed through untouched, so re-running over an indexed corpus is a
192    /// no-op.
193    pub async fn enhance_documents(&self, docs: &[Document]) -> Vec<Document> {
194        let mut handles = Vec::with_capacity(docs.len());
195        for doc in docs {
196            let doc = doc.clone();
197            let llm = self.llm.clone();
198            let prompt_template = self.config.prompt_template.clone();
199            let permit = self.semaphore.clone();
200            handles.push(tokio::spawn(async move {
201                if doc.metadata.contains_key(CONTEXTUAL_METADATA_KEY) {
202                    return doc; // idempotent skip
203                }
204                let _permit = permit.acquire().await;
205                let template = PromptTemplate::new(&prompt_template);
206                let mut vars = HashMap::new();
207                vars.insert("chunk", doc.content.as_str());
208                let prompt = template
209                    .format(&vars)
210                    .unwrap_or_else(|_| prompt_template.clone());
211                match llm.invoke(vec![Message::human(prompt)], None).await {
212                    Ok(response) => {
213                        let context = response.content.trim().to_string();
214                        if context.is_empty() {
215                            log::warn!("contextual retrieval: empty context generated, keeping original content");
216                            doc
217                        } else {
218                            let mut enhanced =
219                                Document::new(format!("{}\n{}", context, doc.content))
220                                    .with_metadata(CONTEXTUAL_METADATA_KEY, context);
221                            if let Some(id) = &doc.id {
222                                enhanced = enhanced.with_id(id.clone());
223                            }
224                            for (k, v) in &doc.metadata {
225                                enhanced.metadata.insert(k.clone(), v.clone());
226                            }
227                            enhanced
228                        }
229                    }
230                    Err(e) => {
231                        log::warn!(
232                            "contextual retrieval: context generation failed ({}), keeping original content",
233                            e
234                        );
235                        doc
236                    }
237                }
238            }));
239        }
240        let mut out = Vec::with_capacity(handles.len());
241        for handle in handles {
242            out.push(handle.await.unwrap_or_else(|e| {
243                log::warn!(
244                    "contextual retrieval: enhancement task panicked ({}), keeping placeholder",
245                    e
246                );
247                Document::new("")
248            }));
249        }
250        out
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use futures_util::Stream;
258    use lc_core::language_models::{BaseLanguageModel, LLMResult, StreamChunk};
259    use lc_core::runnables::{Runnable, RunnableConfig};
260    use std::pin::Pin;
261    use std::sync::atomic::{AtomicUsize, Ordering};
262
263    /// Mock chat model returning a fixed context sentence per call.
264    #[derive(Clone)]
265    struct MockChat {
266        reply: String,
267        calls: Arc<AtomicUsize>,
268    }
269
270    impl MockChat {
271        fn new(reply: &str) -> Self {
272            Self {
273                reply: reply.to_string(),
274                calls: Arc::new(AtomicUsize::new(0)),
275            }
276        }
277    }
278
279    #[derive(Debug)]
280    struct MockChatError(String);
281
282    impl std::fmt::Display for MockChatError {
283        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
284            write!(f, "mock chat error: {}", self.0)
285        }
286    }
287    impl std::error::Error for MockChatError {}
288
289    impl From<MockChatError> for ProviderError {
290        fn from(e: MockChatError) -> Self {
291            ProviderError::Config(e.to_string())
292        }
293    }
294
295    #[async_trait::async_trait]
296    impl Runnable<Vec<Message>, LLMResult> for MockChat {
297        type Error = MockChatError;
298        async fn invoke(
299            &self,
300            input: Vec<Message>,
301            _config: Option<RunnableConfig>,
302        ) -> Result<LLMResult, Self::Error> {
303            self.calls.fetch_add(1, Ordering::SeqCst);
304            let prompt = input.first().map(|m| m.content.clone()).unwrap_or_default();
305            Ok(LLMResult {
306                content: format!("{} [from: {}]", self.reply, prompt),
307                model: "mock".to_string(),
308                token_usage: None,
309                tool_calls: None,
310                thinking_content: None,
311            })
312        }
313    }
314
315    #[async_trait::async_trait]
316    impl BaseLanguageModel<Vec<Message>, LLMResult> for MockChat {
317        fn model_name(&self) -> &str {
318            "mock"
319        }
320        fn get_num_tokens(&self, t: &str) -> usize {
321            t.len()
322        }
323        fn with_temperature(self, _: f32) -> Self {
324            self
325        }
326        fn with_max_tokens(self, _: usize) -> Self {
327            self
328        }
329    }
330
331    #[async_trait::async_trait]
332    impl BaseChatModel for MockChat {
333        async fn chat(
334            &self,
335            messages: Vec<Message>,
336            config: Option<RunnableConfig>,
337        ) -> Result<LLMResult, Self::Error> {
338            <Self as Runnable<Vec<Message>, LLMResult>>::invoke(self, messages, config).await
339        }
340        async fn stream_chat(
341            &self,
342            _messages: Vec<Message>,
343            _config: Option<RunnableConfig>,
344        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
345        {
346            Err(MockChatError("no stream".to_string()))
347        }
348    }
349
350    /// Failing mock for the fail-open path.
351    struct FailingChat;
352
353    #[async_trait::async_trait]
354    impl Runnable<Vec<Message>, LLMResult> for FailingChat {
355        type Error = MockChatError;
356        async fn invoke(
357            &self,
358            _input: Vec<Message>,
359            _config: Option<RunnableConfig>,
360        ) -> Result<LLMResult, Self::Error> {
361            Err(MockChatError("llm down".to_string()))
362        }
363    }
364
365    #[async_trait::async_trait]
366    impl BaseLanguageModel<Vec<Message>, LLMResult> for FailingChat {
367        fn model_name(&self) -> &str {
368            "failing"
369        }
370        fn get_num_tokens(&self, t: &str) -> usize {
371            t.len()
372        }
373        fn with_temperature(self, _: f32) -> Self {
374            self
375        }
376        fn with_max_tokens(self, _: usize) -> Self {
377            self
378        }
379    }
380
381    #[async_trait::async_trait]
382    impl BaseChatModel for FailingChat {
383        async fn chat(
384            &self,
385            _messages: Vec<Message>,
386            _config: Option<RunnableConfig>,
387        ) -> Result<LLMResult, Self::Error> {
388            Err(MockChatError("llm down".to_string()))
389        }
390        async fn stream_chat(
391            &self,
392            _messages: Vec<Message>,
393            _config: Option<RunnableConfig>,
394        ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>, Self::Error>
395        {
396            Err(MockChatError("no stream".to_string()))
397        }
398    }
399
400    #[test]
401    fn config_defaults_and_builders() {
402        let config = ContextualConfig::default();
403        assert_eq!(config.max_concurrency, 4);
404        assert!(config.prompt_template.contains("{chunk}"));
405
406        let config = ContextualConfig::new()
407            .with_prompt("Contextualize: {chunk}")
408            .with_max_concurrency(0); // clamped to 1
409        assert_eq!(config.max_concurrency, 1);
410        assert!(config.prompt_template.contains("{chunk}"));
411    }
412
413    /// Single-document enhancement: content is prefixed with the context, the
414    /// context is stored in metadata, and the prompt carried the chunk text.
415    #[tokio::test]
416    async fn enhance_document_prefixes_context_and_stores_metadata() {
417        let mock = MockChat::new("This chunk is part of the revenue report.");
418        let enhancer = ContextualEnhancer::new(mock);
419        let doc = Document::new("Revenue grew 3% year over year.");
420
421        let enhanced = enhancer.enhance_document(&doc).await.unwrap();
422
423        assert!(
424            enhanced
425                .content
426                .starts_with("This chunk is part of the revenue report."),
427            "context should prefix the original content"
428        );
429        assert!(enhanced
430            .content
431            .ends_with("Revenue grew 3% year over year."));
432        let content = &enhanced.content;
433        let expected_context_len = content.len() - doc.content.len() - "\n".len();
434        assert_eq!(
435            enhanced
436                .metadata
437                .get(CONTEXTUAL_METADATA_KEY)
438                .and_then(|v| v.as_str()),
439            Some(&content[..expected_context_len]),
440            "metadata stores exactly the generated context (the content prefix)"
441        );
442    }
443
444    /// The prompt sent to the LLM contains the chunk text ({chunk} replaced):
445    /// the mock echoes the prompt into its reply, so the stored context must
446    /// contain the original chunk text.
447    #[tokio::test]
448    async fn prompt_sent_to_llm_contains_chunk() {
449        let mock = MockChat::new("ctx");
450        let enhancer = ContextualEnhancer::new(mock.clone());
451        let doc = Document::new("Revenue grew 3% year over year.");
452        let _ = enhancer.generate_context(&doc.content).await.unwrap();
453
454        let document = enhancer.enhance_document(&doc).await.unwrap();
455        let context = document
456            .metadata
457            .get(CONTEXTUAL_METADATA_KEY)
458            .and_then(|v| v.as_str())
459            .unwrap();
460        assert!(
461            context.contains("Revenue grew 3% year over year."),
462            "prompt should carry the chunk text, got: {}",
463            context
464        );
465    }
466
467    /// Fail-open: LLM failure keeps the original content and adds no metadata.
468    #[tokio::test]
469    async fn batch_fails_open_on_llm_error() {
470        let enhancer = ContextualEnhancer::new(FailingChat);
471        let docs = vec![Document::new("chunk a"), Document::new("chunk b")];
472
473        let out = enhancer.enhance_documents(&docs).await;
474        assert_eq!(out.len(), 2);
475        for (original, enhanced) in docs.iter().zip(out.iter()) {
476            assert_eq!(enhanced.content, original.content, "original kept");
477            assert!(
478                !enhanced.metadata.contains_key(CONTEXTUAL_METADATA_KEY),
479                "no context metadata on failure"
480            );
481        }
482    }
483
484    /// Batch enhancement prefixes each chunk and reports the call count.
485    #[tokio::test]
486    async fn batch_enhances_each_chunk() {
487        let mock = MockChat::new("ctx");
488        let enhancer = ContextualEnhancer::new(mock.clone());
489        let docs = vec![Document::new("a"), Document::new("b"), Document::new("c")];
490
491        let out = enhancer.enhance_documents(&docs).await;
492        assert_eq!(out.len(), 3);
493        for (original, enhanced) in docs.iter().zip(out.iter()) {
494            assert!(enhanced.content.starts_with("ctx"));
495            assert!(enhanced.content.ends_with(original.content.as_str()));
496            assert!(enhanced.metadata.contains_key(CONTEXTUAL_METADATA_KEY));
497        }
498        assert_eq!(
499            mock.calls.load(Ordering::SeqCst),
500            3,
501            "one LLM call per chunk"
502        );
503    }
504
505    /// Idempotency: already-enhanced chunks are skipped (no extra LLM calls).
506    #[tokio::test]
507    async fn batch_is_idempotent() {
508        let mock = MockChat::new("ctx");
509        let enhancer = ContextualEnhancer::new(mock.clone());
510        let docs = vec![Document::new("a"), Document::new("b")];
511        let enhanced = enhancer.enhance_documents(&docs).await;
512        assert_eq!(mock.calls.load(Ordering::SeqCst), 2);
513
514        // Re-run over the enhanced corpus: no additional calls, content unchanged.
515        let rerun = enhancer.enhance_documents(&enhanced).await;
516        assert_eq!(
517            mock.calls.load(Ordering::SeqCst),
518            2,
519            "no new calls on rerun"
520        );
521        for (first, second) in enhanced.iter().zip(rerun.iter()) {
522            assert_eq!(first.content, second.content);
523        }
524    }
525
526    /// Bounded concurrency: max_concurrency=1 serializes LLM calls.
527    #[tokio::test]
528    async fn concurrency_is_bounded() {
529        #[derive(Default)]
530        struct CountingState {
531            in_flight: std::sync::Mutex<usize>,
532            max_seen: AtomicUsize,
533        }
534
535        #[derive(Clone)]
536        struct CountingChat {
537            state: Arc<CountingState>,
538        }
539
540        #[async_trait::async_trait]
541        impl Runnable<Vec<Message>, LLMResult> for CountingChat {
542            type Error = MockChatError;
543            async fn invoke(
544                &self,
545                _input: Vec<Message>,
546                _config: Option<RunnableConfig>,
547            ) -> Result<LLMResult, Self::Error> {
548                let now = {
549                    let mut g = self.state.in_flight.lock().unwrap();
550                    *g += 1;
551                    let n = *g;
552                    if n > self.state.max_seen.load(Ordering::SeqCst) {
553                        self.state.max_seen.store(n, Ordering::SeqCst);
554                    }
555                    n
556                };
557                // Hold the slot briefly so overlaps are observable.
558                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
559                {
560                    let mut g = self.state.in_flight.lock().unwrap();
561                    *g -= 1;
562                }
563                Ok(LLMResult {
564                    content: format!("ctx-{}", now),
565                    model: "counting".to_string(),
566                    token_usage: None,
567                    tool_calls: None,
568                    thinking_content: None,
569                })
570            }
571        }
572
573        #[async_trait::async_trait]
574        impl BaseLanguageModel<Vec<Message>, LLMResult> for CountingChat {
575            fn model_name(&self) -> &str {
576                "counting"
577            }
578            fn get_num_tokens(&self, t: &str) -> usize {
579                t.len()
580            }
581            fn with_temperature(self, _: f32) -> Self {
582                self
583            }
584            fn with_max_tokens(self, _: usize) -> Self {
585                self
586            }
587        }
588
589        #[async_trait::async_trait]
590        impl BaseChatModel for CountingChat {
591            async fn chat(
592                &self,
593                messages: Vec<Message>,
594                config: Option<RunnableConfig>,
595            ) -> Result<LLMResult, Self::Error> {
596                <Self as Runnable<Vec<Message>, LLMResult>>::invoke(self, messages, config).await
597            }
598            async fn stream_chat(
599                &self,
600                _messages: Vec<Message>,
601                _config: Option<RunnableConfig>,
602            ) -> Result<
603                Pin<Box<dyn Stream<Item = Result<StreamChunk, Self::Error>> + Send>>,
604                Self::Error,
605            > {
606                Err(MockChatError("no stream".to_string()))
607            }
608        }
609
610        let state = Arc::new(CountingState::default());
611        let enhancer = ContextualEnhancer::new(CountingChat {
612            state: state.clone(),
613        })
614        .with_max_concurrency(1);
615        let docs: Vec<Document> = (0..4)
616            .map(|i| Document::new(format!("chunk {}", i)))
617            .collect();
618        let out = enhancer.enhance_documents(&docs).await;
619        assert_eq!(out.len(), 4);
620        assert!(
621            state.max_seen.load(Ordering::SeqCst) <= 1,
622            "max_concurrency=1 must serialize calls"
623        );
624    }
625}
lc_rag/contextual.rs

lc_rag/
contextual.rs
