Spaces:
Running
Running
File size: 5,871 Bytes
e43c82d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """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),
}
|