lc_rag/
pipeline.rs

1// lc-rag/src/pipeline.rs
2//! RAGPipeline & RAGPipelineBuilder — a complete RAG pipeline in one line
3//!
4//! Provides a fluent Builder API that assembles LLM + Embeddings + VectorStore + Retriever
5//! into a complete RAG pipeline.
6//!
7//! # Example
8//!
9//! ```ignore
10//! let rag = RAGPipelineBuilder::new()
11//!     .llm(OpenAIChat::new(OpenAIConfig::new("sk-...")))
12//!     .embeddings(OpenAIEmbeddings::new(config)?)
13//!     .vector_store(InMemoryVectorStore::new())
14//!     .build()?;
15//!
16//! rag.index_documents(docs).await?;
17//! let answer = rag.query("What is RustB?").await?;
18//! ```
19
20use lc_core::language_models::BaseChatModel;
21use lc_embeddings::Embeddings;
22use lc_providers::ProviderError;
23use lc_schema::Message;
24use lc_vector_stores::{Document, VectorStore, VectorStoreError};
25
26use crate::retriever::{RetrieverError, RetrieverTrait, SimilarityRetriever};
27
28use std::sync::Arc;
29
30/// RAG Pipeline — chunking + embedding + storage + retrieval + generation
31///
32/// Assembles an LLM with a `RetrieverTrait` implementation (BM25, vector similarity,
33/// hybrid retrieval, etc.) into a complete RAG pipeline, exposing three core methods:
34/// `index_documents()`, `query()`, `query_with_sources()`.
35///
36/// P0-2: The retrieval path converges on `Arc<dyn RetrieverTrait>` instead of depending
37/// directly on `Embeddings + VectorStore`, so any retriever can be swapped in seamlessly.
38pub struct RAGPipeline {
39    llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
40    retriever: Arc<dyn RetrieverTrait>,
41    /// Number of documents to retrieve
42    retrieve_k: usize,
43    /// System prompt
44    system_prompt: String,
45}
46
47impl std::fmt::Debug for RAGPipeline {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        f.debug_struct("RAGPipeline")
50            .field("model_name", &self.llm.model_name())
51            .field("retrieve_k", &self.retrieve_k)
52            .field("system_prompt", &self.system_prompt)
53            .finish()
54    }
55}
56
57impl RAGPipeline {
58    /// Indexes documents
59    ///
60    /// P0-2: Delegates to `RetrieverTrait::add_documents` (embedding + storage are handled
61    /// internally by the retriever).
62    pub async fn index_documents(&self, documents: Vec<Document>) -> Result<(), RetrieverError> {
63        self.retriever.add_documents(documents).await
64    }
65
66    /// Query: retrieve + generate an answer
67    ///
68    /// 1. Embed the question
69    /// 2. Retrieve similar documents from the VectorStore
70    /// 3. Use the retrieved results as context and let the LLM generate the answer
71    pub async fn query(&self, question: &str) -> Result<String, RetrieverError> {
72        let result = self.query_with_sources(question).await?;
73        Ok(result.answer)
74    }
75
76    /// Queries and returns the source documents
77    ///
78    /// Returns the generated answer and the list of retrieved source documents.
79    pub async fn query_with_sources(
80        &self,
81        question: &str,
82    ) -> Result<RAGQueryResult, RetrieverError> {
83        // 1. Retrieve relevant documents (P0-2: delegated to RetrieverTrait)
84        let search_results = self
85            .retriever
86            .retrieve_with_scores(question, self.retrieve_k)
87            .await?;
88
89        let sources: Vec<Document> = search_results.iter().map(|r| r.document.clone()).collect();
90
91        // 3. Build the context
92        let context = if sources.is_empty() {
93            "No relevant documents found.".to_string()
94        } else {
95            sources
96                .iter()
97                .enumerate()
98                .map(|(i, doc)| format!("[{}] {}", i + 1, doc.page_content()))
99                .collect::<Vec<_>>()
100                .join("\n\n")
101        };
102
103        // 4. Generate the answer
104        let messages = vec![
105            Message::system(format!(
106                "{}\n\nUse the following context to answer the question. If the context doesn't contain the answer, say so.",
107                self.system_prompt
108            )),
109            Message::human(format!("Context:\n{}\n\nQuestion: {}", context, question)),
110        ];
111
112        let llm_result = self
113            .llm
114            .chat(messages, None)
115            .await
116            .map_err(|e| RetrieverError::EmbeddingError(format!("LLM call failed: {}", e)))?;
117
118        Ok(RAGQueryResult {
119            answer: llm_result.content,
120            sources,
121        })
122    }
123}
124
125/// RAG query result
126#[derive(Debug, Clone)]
127pub struct RAGQueryResult {
128    /// The generated answer
129    pub answer: String,
130    /// The retrieved source documents
131    pub sources: Vec<Document>,
132}
133
134// ---------------------------------------------------------------------------
135// RAGPipelineBuilder
136// ---------------------------------------------------------------------------
137
138/// RAG Pipeline Builder — fluent API for creating a RAG pipeline
139///
140/// # Example
141///
142/// ```ignore
143/// let rag = RAGPipelineBuilder::new()
144///     .llm(OpenAIChat::new(OpenAIConfig::new("sk-...")))
145///     .embeddings(OpenAIEmbeddings::new(config)?)
146///     .vector_store(InMemoryVectorStore::new())
147///     .build()?;
148/// ```
149pub struct RAGPipelineBuilder {
150    llm: Option<Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>>,
151    embeddings: Option<Arc<dyn Embeddings + Send + Sync>>,
152    vector_store: Option<Arc<dyn VectorStore + Send + Sync>>,
153    /// P0-2: An explicitly passed-in retriever (takes priority); when absent, one is built
154    /// from embeddings + vector_store
155    retriever: Option<Arc<dyn RetrieverTrait>>,
156    retrieve_k: usize,
157    system_prompt: Option<String>,
158}
159
160impl RAGPipelineBuilder {
161    /// Creates a new RAGPipelineBuilder
162    pub fn new() -> Self {
163        Self {
164            llm: None,
165            embeddings: None,
166            vector_store: None,
167            retriever: None,
168            retrieve_k: 4,
169            system_prompt: None,
170        }
171    }
172
173    /// Sets the LLM (any type implementing `BaseChatModel`)
174    pub fn llm<L>(mut self, llm: L) -> Self
175    where
176        L: BaseChatModel + Send + Sync + 'static,
177        L::Error: Into<ProviderError>,
178    {
179        self.llm = Some(lc_providers::wrap_chat_model(llm));
180        self
181    }
182
183    /// Sets the LLM (from an already-wrapped `Arc<dyn BaseChatModel>`)
184    pub fn llm_from_arc(
185        mut self,
186        llm: Arc<dyn BaseChatModel<Error = ProviderError> + Send + Sync>,
187    ) -> Self {
188        self.llm = Some(llm);
189        self
190    }
191
192    /// Sets the LLM (from an `LLMClient`)
193    pub fn llm_client(mut self, client: lc_providers::LLMClient) -> Self {
194        let provider_arc = client.into_inner();
195        self.llm = Some(provider_arc);
196        self
197    }
198
199    /// Sets the Embeddings
200    pub fn embeddings<E: Embeddings + Send + Sync + 'static>(mut self, embeddings: E) -> Self {
201        self.embeddings = Some(Arc::new(embeddings));
202        self
203    }
204
205    /// Sets the VectorStore
206    pub fn vector_store<V: VectorStore + Send + Sync + 'static>(mut self, store: V) -> Self {
207        self.vector_store = Some(Arc::new(store));
208        self
209    }
210
211    /// Sets a custom retriever (any type implementing `RetrieverTrait`,
212    /// such as BM25, UnifiedHybridIndex, etc.)
213    ///
214    /// P0-2: An explicit retriever takes priority over the similarity retriever built from
215    /// `.embeddings() + .vector_store()`.
216    pub fn retriever<R>(mut self, retriever: R) -> Self
217    where
218        R: RetrieverTrait + Send + Sync + 'static,
219    {
220        self.retriever = Some(Arc::new(retriever));
221        self
222    }
223
224    /// Sets the retriever (from an already-wrapped `Arc<dyn RetrieverTrait>`)
225    pub fn retriever_from_arc(mut self, retriever: Arc<dyn RetrieverTrait>) -> Self {
226        self.retriever = Some(retriever);
227        self
228    }
229
230    /// Sets the number of documents to retrieve
231    pub fn retrieve_k(mut self, k: usize) -> Self {
232        self.retrieve_k = k;
233        self
234    }
235
236    /// Sets the system prompt
237    pub fn system(mut self, prompt: impl Into<String>) -> Self {
238        self.system_prompt = Some(prompt.into());
239        self
240    }
241
242    /// Builds the RAGPipeline
243    ///
244    /// # Errors
245    ///
246    /// Returns an error if the LLM, Embeddings, or VectorStore is missing.
247    pub fn build(self) -> Result<RAGPipeline, RetrieverError> {
248        let llm = self.llm.ok_or_else(|| {
249            RetrieverError::EmbeddingError(
250                "RAGPipelineBuilder: LLM is required. Call .llm() first.".into(),
251            )
252        })?;
253
254        // P0-2: Prefer the explicit retriever; otherwise fall back to building a
255        // SimilarityRetriever from embeddings + vector_store for backward compatibility.
256        let retriever = match self.retriever {
257            Some(r) => r,
258            None => {
259                let embeddings = self.embeddings.ok_or_else(|| {
260                    RetrieverError::EmbeddingError(
261                        "RAGPipelineBuilder: Embeddings is required (or use .retriever()). Call .embeddings() first."
262                            .into(),
263                    )
264                })?;
265
266                let vector_store = self.vector_store.ok_or_else(|| {
267                    RetrieverError::StoreError(VectorStoreError::StorageError(
268                        "RAGPipelineBuilder: VectorStore is required (or use .retriever()). Call .vector_store() first."
269                            .into(),
270                    ))
271                })?;
272
273                Arc::new(SimilarityRetriever::new(vector_store, embeddings))
274            }
275        };
276
277        Ok(RAGPipeline {
278            llm,
279            retriever,
280            retrieve_k: self.retrieve_k,
281            system_prompt: self.system_prompt.unwrap_or_else(|| {
282                "You are a helpful assistant that answers questions based on the provided context.".to_string()
283            }),
284        })
285    }
286}
287
288impl Default for RAGPipelineBuilder {
289    fn default() -> Self {
290        Self::new()
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297    use lc_embeddings::MockEmbeddings;
298    use lc_providers::{OpenAIChat, OpenAIConfig};
299    use lc_vector_stores::InMemoryVectorStore;
300
301    #[test]
302    fn test_builder_missing_llm() {
303        let result = RAGPipelineBuilder::new()
304            .embeddings(MockEmbeddings::new(3))
305            .vector_store(InMemoryVectorStore::new())
306            .build();
307
308        assert!(result.is_err());
309        assert!(result.unwrap_err().to_string().contains("LLM is required"));
310    }
311
312    #[test]
313    fn test_builder_missing_embeddings() {
314        let config = OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1");
315        let result = RAGPipelineBuilder::new()
316            .llm(OpenAIChat::new(config))
317            .vector_store(InMemoryVectorStore::new())
318            .build();
319
320        assert!(result.is_err());
321        assert!(result
322            .unwrap_err()
323            .to_string()
324            .contains("Embeddings is required"));
325    }
326
327    #[test]
328    fn test_builder_missing_vector_store() {
329        let config = OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1");
330        let result = RAGPipelineBuilder::new()
331            .llm(OpenAIChat::new(config))
332            .embeddings(MockEmbeddings::new(3))
333            .build();
334
335        assert!(result.is_err());
336        assert!(result
337            .unwrap_err()
338            .to_string()
339            .contains("VectorStore is required"));
340    }
341
342    #[test]
343    fn test_builder_success() {
344        let config = OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1");
345        let result = RAGPipelineBuilder::new()
346            .llm(OpenAIChat::new(config))
347            .embeddings(MockEmbeddings::new(3))
348            .vector_store(InMemoryVectorStore::new())
349            .system("You are a test assistant.")
350            .retrieve_k(5)
351            .build();
352
353        assert!(result.is_ok());
354        let pipeline = result.unwrap();
355        assert_eq!(pipeline.retrieve_k, 5);
356        assert_eq!(pipeline.system_prompt, "You are a test assistant.");
357    }
358
359    #[test]
360    fn test_builder_default() {
361        let builder = RAGPipelineBuilder::default();
362        assert_eq!(builder.retrieve_k, 4);
363        assert!(builder.llm.is_none());
364    }
365
366    /// P0-2: Supports injecting a custom `RetrieverTrait` implementation (BM25) via
367    /// `.retriever()`; `index_documents` delegates to it without needing Embeddings/VectorStore.
368    #[tokio::test]
369    async fn test_builder_with_custom_retriever() {
370        use crate::bm25::BM25Retriever;
371
372        let config = OpenAIConfig::new("test_key").with_base_url("http://localhost:8080/v1");
373        let rag = RAGPipelineBuilder::new()
374            .llm(OpenAIChat::new(config))
375            .retriever(BM25Retriever::new())
376            .build()
377            .expect("build should succeed");
378
379        rag.index_documents(vec![Document::new("Rust is a systems language")])
380            .await
381            .expect("index_documents should delegate to BM25 retriever successfully");
382    }
383}
lc_rag/pipeline.rs

lc_rag/
pipeline.rs
