Spaces:
Running
Running
| """Complete end-to-end RAG Pipeline orchestrator. | |
| Coordinates all 5 core stages of Retrieval-Augmented Generation: | |
| 1. Ingestion: Load raw text, PDFs, or demo articles into Document objects. | |
| 2. Chunking: Split documents into smaller, overlapping TextChunk passages. | |
| 3. Indexing: Index chunks into both Dense Vector Store and BM25 Keyword Search. | |
| 4. Hybrid Retrieval: Retrieve and rank chunks using Reciprocal Rank Fusion (RRF). | |
| 5. Generation: Augment prompt with retrieved chunks and generate answer with LLM. | |
| """ | |
| import time | |
| from pathlib import Path | |
| from config.settings import get_settings | |
| from rag_engine.documents import Document, load_text_file, load_pdf_file, save_documents_jsonl, load_documents_jsonl | |
| from rag_engine.chunking import chunk_documents | |
| from rag_engine.vector_store import VectorStore | |
| from rag_engine.bm25_search import BM25SearchIndex | |
| from rag_engine.hybrid_search import HybridRetriever | |
| from rag_engine.generator import generate_answer | |
| from rag_engine.eval import evaluate_response | |
| class RAGPipeline: | |
| """The central RAG engine coordinating ingestion, search, and generation.""" | |
| def __init__(self, settings=None): | |
| self.settings = settings or get_settings() | |
| self.documents = [] | |
| self.chunks = [] | |
| self.vector_store = VectorStore() | |
| self.bm25_index = BM25SearchIndex() | |
| self.retriever = HybridRetriever( | |
| vector_store=self.vector_store, | |
| bm25_index=self.bm25_index, | |
| dense_weight=self.settings.dense_weight, | |
| bm25_weight=self.settings.bm25_weight, | |
| ) | |
| def add_document(self, text, title=None, doc_id=None, source="upload"): | |
| """Add a single raw text document and index it.""" | |
| doc = Document(text=text, title=title, doc_id=doc_id, source=source) | |
| self.documents.append(doc) | |
| new_chunks = chunk_documents( | |
| [doc], | |
| chunk_size=self.settings.chunk_size, | |
| chunk_overlap=self.settings.chunk_overlap, | |
| ) | |
| self.chunks.extend(new_chunks) | |
| self.vector_store.add_chunks(new_chunks) | |
| self.bm25_index.build(self.chunks) | |
| save_documents_jsonl(self.documents, self.settings.documents_file) | |
| return doc | |
| def add_pdf(self, filepath): | |
| """Extract text from a PDF file and index it.""" | |
| doc = load_pdf_file(filepath) | |
| return self.add_document(text=doc.text, title=doc.title, doc_id=doc.doc_id, source=doc.source) | |
| def load_demo_corpus(self): | |
| """Load bundled educational demo articles.""" | |
| demo_dir = Path(__file__).resolve().parents[2] / "tests" / "fixtures" / "demo_corpus" | |
| if not demo_dir.exists(): | |
| return [] | |
| loaded = [] | |
| for txt_file in sorted(demo_dir.glob("*.txt")): | |
| doc = load_text_file(txt_file) | |
| self.documents.append(doc) | |
| loaded.append(doc) | |
| if loaded: | |
| self.chunks = chunk_documents( | |
| self.documents, | |
| chunk_size=self.settings.chunk_size, | |
| chunk_overlap=self.settings.chunk_overlap, | |
| ) | |
| self.vector_store.clear() | |
| self.vector_store.add_chunks(self.chunks) | |
| self.bm25_index.build(self.chunks) | |
| save_documents_jsonl(self.documents, self.settings.documents_file) | |
| return loaded | |
| def query(self, question, top_k=None): | |
| """Execute the end-to-end RAG workflow for a user question. | |
| Steps: | |
| 1. Query dense vector store and BM25 keyword index. | |
| 2. Merge results using Reciprocal Rank Fusion. | |
| 3. Format retrieved context into LLM prompt. | |
| 4. Synthesize final answer with source citations. | |
| 5. Evaluate faithfulness and relevancy. | |
| Returns: | |
| Dict containing: | |
| - 'question': input question | |
| - 'answer': generated answer | |
| - 'sources': cited document titles | |
| - 'retrieved_chunks': top ranked chunks | |
| - 'dense_results': vector search results | |
| - 'bm25_results': keyword search results | |
| - 'evaluation': faithfulness & relevancy metrics | |
| - 'latency_ms': elapsed execution time | |
| """ | |
| start_time = time.perf_counter() | |
| top_k = top_k or self.settings.final_top_k | |
| # Step 1: Hybrid Retrieval | |
| retrieval_output = self.retriever.retrieve( | |
| query=question, | |
| top_k=top_k, | |
| dense_k=self.settings.dense_top_k, | |
| bm25_k=self.settings.bm25_top_k, | |
| ) | |
| retrieved_chunks = retrieval_output["results"] | |
| # Step 2: Generation | |
| gen_output = generate_answer( | |
| question=question, | |
| retrieved_items=retrieved_chunks, | |
| settings=self.settings, | |
| ) | |
| # Step 3: Evaluation | |
| evaluation = evaluate_response( | |
| question=question, | |
| answer=gen_output["answer"], | |
| retrieved_items=retrieved_chunks, | |
| ) | |
| elapsed_ms = round((time.perf_counter() - start_time) * 1000, 1) | |
| return { | |
| "question": question, | |
| "answer": gen_output["answer"], | |
| "sources": gen_output["sources"], | |
| "used_llm": gen_output["used_llm"], | |
| "retrieved_chunks": retrieved_chunks, | |
| "dense_results": retrieval_output["dense_results"], | |
| "bm25_results": retrieval_output["bm25_results"], | |
| "evaluation": evaluation, | |
| "latency_ms": elapsed_ms, | |
| } | |
| def get_stats(self): | |
| """Return summary statistics of the indexed knowledge base.""" | |
| return { | |
| "document_count": len(self.documents), | |
| "chunk_count": len(self.chunks), | |
| "vector_store_size": len(self.vector_store), | |
| "bm25_index_size": len(self.bm25_index), | |
| } | |