import os import time import pandas as pd import gradio as gr import config as cfg from cache import LRUCache from data_loader import load_documents, validate_document from metadata_extractor import extract_metadata from chunker import chunk_documents from embedder import Embedder from duplicate_detector import DuplicateDetector from vector_store import VectorStore from bm25_retriever import BM25Retriever from dense_retriever import DenseRetriever from hybrid_retriever import HybridRetriever from reranker import Reranker from context_manager import ContextManager from llm_handler import LLMHandler from evaluator import Evaluator from query_rewriter import QueryRewriter class RagPipeline: def __init__(self): self.chunks = [] self.embedder = None self.vector_store = None self.bm25 = None self.dense = None self.hybrid = None self.reranker = None self.context_manager = None self.llm = None self.evaluator = None self.rewriter = None self.index_built = False self.answer_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) self.retrieval_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) self.embedding_cache = LRUCache(max_size=cfg.CACHE_MAX_SIZE, ttl_seconds=cfg.CACHE_TTL_SECONDS) def build_index(self, progress=None): if progress: progress(0, desc="Loading documents...") raw = load_documents(cfg.DATA_PATH) if not validate_document(raw): raise ValueError("Invalid document format") if progress: progress(0.15, desc="Extracting metadata...") entries = extract_metadata(raw) if progress: progress(0.30, desc="Chunking documents...") self.chunks = chunk_documents( entries, chunk_size=cfg.CHUNK_SIZE, overlap=cfg.CHUNK_OVERLAP, ) if progress: progress(0.40, desc="Removing duplicates...") dedup = DuplicateDetector( threshold=cfg.DUP_SIM_THRESHOLD, num_perm=cfg.DUP_NUM_PERM, ) self.chunks = dedup.filter_duplicates(self.chunks) if progress: progress(0.50, desc="Loading embedding model...") self.embedder = Embedder( model_name=cfg.EMBEDDING_MODEL, device=cfg.EMBEDDING_DEVICE, ) if progress: progress(0.60, desc="Embedding chunks...") for c in self.chunks: c["search_text"] = " | ".join(p for p in (c.get("title", ""), c.get("summary", ""), c["text"]) if p) texts = [c["search_text"] for c in self.chunks] embeddings = self.embedder.embed(texts) if progress: progress(0.70, desc="Indexing vector store...") self.vector_store = VectorStore( persist_dir=cfg.CHROMA_DB_PATH, embedding_dim=cfg.EMBEDDING_DIM, ) self.vector_store.create_collection( overwrite=True, ef_construction=cfg.HNSW_EF_CONSTRUCTION, m=cfg.HNSW_M, ef_search=cfg.HNSW_EF_SEARCH, ) ids = [c["chunk_id"] for c in self.chunks] metadatas = [] for c in self.chunks: meta = {} for k, v in c.items(): if k in ("text", "chunk_id"): continue if isinstance(v, list): meta[k] = " | ".join(str(x) for x in v if x) else: meta[k] = str(v) metadatas.append(meta) self.vector_store.add(ids, embeddings, texts, metadatas) if progress: progress(0.85, desc="Building BM25 index...") self.bm25 = BM25Retriever() self.bm25.fit(self.chunks) self.dense = DenseRetriever(self.vector_store) self.hybrid = HybridRetriever(self.bm25, self.dense, alpha=cfg.HYBRID_ALPHA, fusion_method=cfg.HYBRID_FUSION_METHOD) if progress: progress(0.90, desc="Loading reranker...") self.reranker = Reranker( model_name=cfg.RERANKER_MODEL, device=cfg.RERANKER_DEVICE, ) self.context_manager = ContextManager(max_tokens=cfg.MAX_CONTEXT_TOKENS) if cfg.GROQ_API_KEY: if progress: progress(0.95, desc="Initializing LLM...") self.llm = LLMHandler( api_key=cfg.GROQ_API_KEY, model=cfg.GROQ_MODEL, temperature=cfg.GROQ_TEMPERATURE, max_tokens=cfg.GROQ_MAX_TOKENS, timeout=cfg.GROQ_TIMEOUT, ) self.rewriter = QueryRewriter(llm=self.llm) self.evaluator = Evaluator( self.hybrid, self.reranker, self.llm, self.context_manager, self.embedder ) self.answer_cache.clear() self.retrieval_cache.clear() self.embedding_cache.clear() self.index_built = True def _retrieve(self, search_query: str, alpha: float, top_k: int): cache_key = (search_query, alpha) cached = self.retrieval_cache.get(cache_key) if cached: return cached cached_emb = self.embedding_cache.get(search_query) if cached_emb is not None: query_emb = cached_emb else: query_emb = self.embedder.embed_query(search_query) self.embedding_cache.put(search_query, query_emb) self.hybrid.alpha = alpha results = self.hybrid.search(search_query, query_emb, n_results=cfg.HYBRID_TOP_K) reranked = self.reranker.rerank(search_query, results[:cfg.RERANK_CANDIDATES], top_k=top_k) self.retrieval_cache.put(cache_key, reranked) return reranked def _build_sources_df(self, reranked): sources = [] for i, r in enumerate(reranked): sources.append({ "Rank": i + 1, "Score": round(r.get("rerank_score", r.get("score", 0)), 4), "Title": r["metadata"].get("title", ""), "Doc ID": r["id"], "Preview": r["text"][:200] + ("..." if len(r["text"]) > 200 else ""), }) return pd.DataFrame(sources) def query(self, question: str, alpha: float, top_k: int, use_llm: bool, use_reformulation: bool = False): if not self.index_built: return "Index not built yet. Click 'Build Index' first.", pd.DataFrame() cache_key = (question, alpha, top_k, use_llm, use_reformulation) cached = self.answer_cache.get(cache_key) if cached: return cached["answer"], cached["sources"] start = time.perf_counter() search_query = question if use_reformulation and self.rewriter: search_query = self.rewriter.rewrite(question) reranked = self._retrieve(search_query, alpha, top_k) retrieval_time = (time.perf_counter() - start) * 1000 if use_llm and self.llm: prompt = self.context_manager.assemble_prompt(question, reranked) answer = self.llm.generate(prompt) gen_time = (time.perf_counter() - start) * 1000 - retrieval_time answer_text = answer or "(LLM returned no response)" else: answer_text = "(LLM disabled — retrieval only)" gen_time = 0 total_time = (time.perf_counter() - start) * 1000 df = self._build_sources_df(reranked) result = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" self.answer_cache.put(cache_key, {"answer": result, "sources": df}) return result, df def query_stream(self, question: str, alpha: float, top_k: int, use_llm: bool, use_reformulation: bool = False): if not question.strip(): yield "Please enter a question.", pd.DataFrame() return if not self.index_built: yield "Index not built yet. Click 'Build Index' first.", pd.DataFrame() return cache_key = (question, alpha, top_k, use_llm, use_reformulation) cached = self.answer_cache.get(cache_key) if cached: yield cached["answer"], cached["sources"] return start = time.perf_counter() search_query = question if use_reformulation and self.rewriter: search_query = self.rewriter.rewrite(question) reranked = self._retrieve(search_query, alpha, top_k) retrieval_time = (time.perf_counter() - start) * 1000 df = self._build_sources_df(reranked) if use_llm and self.llm: yield "**Retrieval complete. Generating answer...**", df prompt = self.context_manager.assemble_prompt(question, reranked) full_answer = "" for chunk in self.llm.generate_stream(prompt): if chunk: full_answer += chunk yield f"{full_answer}\n\n---\n*Generating...*", df gen_time = (time.perf_counter() - start) * 1000 - retrieval_time total_time = (time.perf_counter() - start) * 1000 answer_text = full_answer or "(LLM returned no response)" else: answer_text = "(LLM disabled — retrieval only)" gen_time = 0 total_time = (time.perf_counter() - start) * 1000 final = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" self.answer_cache.put(cache_key, {"answer": final, "sources": df}) yield final, df def run_evaluation(self, num_questions: int): if not self.index_built: return "Index not built yet.", pd.DataFrame() results = self.evaluator.full_evaluation(self.chunks, num_questions) rows = [] for k, v in results["retrieval"].items(): val_str = f"{v:.4f}" if isinstance(v, float) else str(v) rows.append({"Metric": k, "Value": val_str}) rows.append({ "Metric": "Num Eval Questions", "Value": results["num_questions"], }) gen = results["generation"] rows.append({ "Metric": "Avg End-to-End Latency (ms)", "Value": gen.get("Avg End-to-End Latency (ms)", "N/A"), }) for m in ("Avg Faithfulness", "Avg Answer Accuracy"): if m in gen: v = gen[m] rows.append({"Metric": m, "Value": f"{v:.4f}" if isinstance(v, float) else str(v)}) df = pd.DataFrame(rows) details = "" for a in gen.get("answers", []): faith = a.get("faithfulness", "N/A") acc = a.get("accuracy", "N/A") faith_str = f"{faith:.3f}" if isinstance(faith, float) else str(faith) acc_str = f"{acc:.3f}" if isinstance(acc, float) else str(acc) details += f"**Q:** {a['question']}\n**A:** {a['answer'][:300]}...\n**Latency:** {a.get('latency_ms', 'N/A')}ms | **Faithfulness:** {faith_str} | **Accuracy:** {acc_str}\n\n---\n\n" return details, df def get_chunk_preview(self, max_rows: int = 50): if not self.chunks: return pd.DataFrame() data = [] for c in self.chunks[:max_rows]: data.append({ "Chunk ID": c.get("chunk_id", ""), "Title": c.get("title", ""), "Level": c.get("level", ""), "Tokens": c.get("num_tokens", 0), "Preview": c["text"][:150] + ("..." if len(c["text"]) > 150 else ""), }) return pd.DataFrame(data) pipeline = RagPipeline() def build_index_fn(progress=gr.Progress()): if pipeline.index_built: return "Index already built. Reload the page to rebuild." try: pipeline.build_index(progress=progress) count = pipeline.vector_store.count() return f"Index built successfully — {count} chunks indexed." except Exception as e: return f"Error: {e}" _stream_state = {} def retrieve_fn(question, alpha, top_k, use_llm, use_reformulation): if not question.strip(): return "Please enter a question.", pd.DataFrame() if not pipeline.index_built: return "Index not built yet. Click 'Build Index' first.", pd.DataFrame() cache_key = (question, alpha, top_k, use_llm, use_reformulation) cached = pipeline.answer_cache.get(cache_key) if cached: return cached["answer"], cached["sources"] search_query = question if use_reformulation and pipeline.rewriter: search_query = pipeline.rewriter.rewrite(question) retrieve_start = time.perf_counter() reranked = pipeline._retrieve(search_query, alpha, top_k) retrieval_time_ms = (time.perf_counter() - retrieve_start) * 1000 df = pipeline._build_sources_df(reranked) _stream_state["query"] = question _stream_state["search_query"] = search_query _stream_state["reranked"] = reranked _stream_state["alpha"] = alpha _stream_state["top_k"] = top_k _stream_state["use_llm"] = use_llm _stream_state["use_reformulation"] = use_reformulation _stream_state["retrieval_time_ms"] = retrieval_time_ms _stream_state["gen_start"] = time.perf_counter() if use_llm and pipeline.llm: return "**Retrieval complete. Generating answer...**", df return "(LLM disabled — retrieval only)", df def stream_fn(): if "query" not in _stream_state or "gen_start" not in _stream_state: return use_llm = _stream_state.get("use_llm", False) if not use_llm or not pipeline.llm: return question = _stream_state.pop("query") reranked = _stream_state.pop("reranked") gen_start = _stream_state.pop("gen_start") retrieval_time_ms = _stream_state.pop("retrieval_time_ms") prompt = pipeline.context_manager.assemble_prompt(question, reranked) full_answer = "" for chunk in pipeline.llm.generate_stream(prompt): if chunk: full_answer += chunk yield f"{full_answer}\n\n---\n*Generating...*" gen_time = (time.perf_counter() - gen_start) * 1000 total_time = retrieval_time_ms + gen_time answer_text = full_answer or "(LLM returned no response)" final = f"{answer_text}\n\n---\n*Retrieval: {retrieval_time_ms:.1f}ms | Generation: {gen_time:.1f}ms | Total: {total_time:.1f}ms*" cache_key = (question, _stream_state["alpha"], _stream_state["top_k"], use_llm, _stream_state["use_reformulation"]) pipeline.answer_cache.put(cache_key, {"answer": final, "sources": pipeline._build_sources_df(reranked)}) yield final def eval_fn(num_q): return pipeline.run_evaluation(int(num_q)) def browse_fn(max_rows): return pipeline.get_chunk_preview(int(max_rows)) with gr.Blocks( title="Production RAG Pipeline", ) as demo: gr.Markdown( "# Production RAG Pipeline\n" "Hybrid retrieval (BM25 + Dense) → Reranking → LLM generation via Groq. " f"Embedding: `{cfg.EMBEDDING_MODEL}` | Reranker: `{cfg.RERANKER_MODEL}`" ) with gr.Tab("Ask"): with gr.Row(): with gr.Column(scale=3): query_input = gr.Textbox( label="Your Question", placeholder="e.g., What is the DeepSeek-V4 architecture?", lines=3, ) with gr.Row(): submit_btn = gr.Button("Ask", variant="primary") clear_btn = gr.Button("Clear") with gr.Column(scale=1): alpha_slider = gr.Slider( minimum=0.0, maximum=1.0, value=cfg.HYBRID_ALPHA, step=0.1, label="Hybrid α (0=BM25, 1=Dense)", ) top_k_dropdown = gr.Dropdown( choices=[3, 5, 10, 15, 20], value=5, label="Final Top-K Results", ) use_llm_checkbox = gr.Checkbox( label="Use LLM (Groq)", value=bool(cfg.GROQ_API_KEY), interactive=bool(cfg.GROQ_API_KEY), ) reformulate_checkbox = gr.Checkbox( label="Reformulate query (better retrieval)", value=False, interactive=True, ) answer_output = gr.Markdown(label="Answer", value="*Ask a question to get started.*") sources_output = gr.Dataframe(label="Retrieved Sources", wrap=True) click_event = submit_btn.click( fn=retrieve_fn, inputs=[query_input, alpha_slider, top_k_dropdown, use_llm_checkbox, reformulate_checkbox], outputs=[answer_output, sources_output], ) click_event.then( fn=stream_fn, outputs=[answer_output], ) clear_btn.click( fn=lambda: ("", pd.DataFrame()), outputs=[query_input, sources_output], ) with gr.Tab("Evaluate"): gr.Markdown("Run evaluation to measure retrieval accuracy and end-to-end latency.") with gr.Row(): num_q_slider = gr.Slider( minimum=5, maximum=50, value=cfg.EVAL_NUM_QUESTIONS, step=5, label="Number of Questions" ) eval_btn = gr.Button("Run Evaluation", variant="primary") eval_details = gr.Markdown(label="Sample Answers") eval_metrics = gr.Dataframe(label="Evaluation Metrics") eval_btn.click( fn=eval_fn, inputs=[num_q_slider], outputs=[eval_details, eval_metrics], ) with gr.Tab("Browse"): gr.Markdown("Browse the indexed document chunks.") with gr.Row(): max_rows_slider = gr.Slider( minimum=10, maximum=200, value=50, step=10, label="Max Rows" ) browse_btn = gr.Button("Refresh", variant="secondary") browse_output = gr.Dataframe(label="Chunks", wrap=True) browse_btn.click(fn=browse_fn, inputs=[max_rows_slider], outputs=[browse_output]) with gr.Tab("Configuration"): config_rows = [ ["Embedding Model", cfg.EMBEDDING_MODEL], ["Reranker Model", cfg.RERANKER_MODEL], ["Chunk Size", str(cfg.CHUNK_SIZE)], ["Chunk Overlap", str(cfg.CHUNK_OVERLAP)], ["HNSW ef_construction", str(cfg.HNSW_EF_CONSTRUCTION)], ["HNSW M", str(cfg.HNSW_M)], ["HNSW ef_search", str(cfg.HNSW_EF_SEARCH)], ["BM25 Top-K", str(cfg.BM25_TOP_K)], ["Dense Top-K", str(cfg.DENSE_TOP_K)], ["Hybrid Top-K", str(cfg.HYBRID_TOP_K)], ["Rerank Top-K", str(cfg.RERANK_TOP_K)], ["Final Top-K", str(cfg.FINAL_TOP_K)], ["Max Context Tokens", str(cfg.MAX_CONTEXT_TOKENS)], ["Dup Threshold", str(cfg.DUP_SIM_THRESHOLD)], ["Groq Model", cfg.GROQ_MODEL], ["Groq API Key Set", "Yes" if cfg.GROQ_API_KEY else "No"], ["Data File", cfg.DATA_PATH], ["Chroma DB Path", cfg.CHROMA_DB_PATH], ] config_df = pd.DataFrame(config_rows, columns=["Parameter", "Value"]) gr.Dataframe(value=config_df, label="Current Configuration") with gr.Tab("Build"): gr.Markdown( "Build or rebuild the index. This loads the document, chunks, embeds, " "and builds both the vector and BM25 indexes." ) build_btn = gr.Button("Build Index", variant="primary") build_output = gr.Markdown("*Index not built yet.*") build_btn.click(fn=build_index_fn, outputs=[build_output]) demo.queue() demo.launch( server_port=cfg.GRADIO_PORT, share=cfg.GRADIO_SHARE, theme=gr.themes.Soft(), css="footer {display:none !important}", )