prakhar-models-004/Retrieval-Augmented_Generation
0
1import os2import sys3import time4from pathlib import Path5from typing import List, Dict, Any6 7root_dir = Path(__file__).resolve().parent8 9if str(root_dir) not in sys.path:10 sys.path.insert(0, str(root_dir))11 12import gradio as gr13import spaces14 15from app.config import settings16from app.ingestion import ingest_documents17from app.retrieval import ChromaVectorStore, VectorRetriever18from app.generation import RAGGenerator, CitationFormatter19 20 21vector_store = ChromaVectorStore()22retriever = VectorRetriever(vector_store=vector_store)23generator = RAGGenerator()24 25 26def get_db_stats() -> str:27 count = vector_store.get_count()28 return f"**Indexed Chunks in Vector DB**: `{count}`"29 30 31def process_ingestion(files: List[Any], reset_db: bool) -> tuple[str, str]:32 if not files:33 return (34 "⚠️ Please select at least one PDF or TXT file to ingest.",35 get_db_stats()36 )37 38 session_id = f"upload_{int(time.time())}"39 data_dir = Path(f"data/uploads/{session_id}")40 data_dir.mkdir(parents=True, exist_ok=True)41 42 for file_obj in files:43 file_path_str = getattr(file_obj, "name", str(file_obj))44 dest_path = data_dir / Path(file_path_str).name45 46 with open(file_path_str, "rb") as f_in:47 with open(dest_path, "wb") as f_out:48 f_out.write(f_in.read())49 50 result = ingest_documents(51 path=data_dir,52 vector_store=vector_store,53 reset=reset_db54 )55 56 status_msg = f"**Status**: {result.get('status', '').upper()}\n\n"57 status_msg += f"{result.get('message', '')}\n\n"58 status_msg += f"- **Files Processed**: {result.get('files_processed', 0)}\n"59 status_msg += f"- **Pages Processed**: {result.get('pages_processed', 0)}\n"60 status_msg += f"- **Chunks Stored**: {result.get('chunks_stored', 0)}\n"61 status_msg += f"- **Total DB Chunks**: {result.get('total_chunks_in_db', 0)}"62 63 return status_msg, get_db_stats()64 65 66@spaces.GPU(duration=120)67def handle_query(68 query: str,69 top_k: int,70 history: List[Dict[str, str]]71) -> tuple[List[Dict[str, str]], str]:72 73 if not query or not query.strip():74 return history, ""75 76 if vector_store.get_count() == 0:77 warning_msg = (78 "⚠️ The document vector database is empty. "79 "Please upload and ingest a document first."80 )81 82 history.append({83 "role": "user",84 "content": query85 })86 87 history.append({88 "role": "assistant",89 "content": warning_msg90 })91 92 return history, ""93 94 retrieved_chunks = retriever.retrieve(95 query=query,96 top_k=int(top_k)97 )98 99 result = generator.generate(100 query=query,101 retrieved_chunks=retrieved_chunks102 )103 104 clean_answer = result.answer105 106 if "\n\nSources:\n" in clean_answer:107 clean_answer = clean_answer.split("\n\nSources:\n")[0]108 elif "\nSources:\n" in clean_answer:109 clean_answer = clean_answer.split("\nSources:\n")[0]110 111 citations_html = ""112 113 if result.citations:114 citations_html = CitationFormatter.format_citations_badges(115 result.citations116 )117 118 formatted_response = clean_answer119 120 if citations_html:121 formatted_response += (122 "\n\n---\n"123 "**📌 Verified Sources & References**:\n"124 f"{citations_html}"125 )126 127 history.append({128 "role": "user",129 "content": query130 })131 132 history.append({133 "role": "assistant",134 "content": formatted_response135 })136 137 return history, ""138 139 140def handle_clear_db() -> tuple[str, str, List[Dict[str, str]]]:141 vector_store.reset()142 143 return (144 "🧹 Vector database has been completely reset.",145 get_db_stats(),146 []147 )148 149 150custom_css = """151.citation-pill {152 display: inline-flex;153 align-items: center;154 gap: 6px;155 background-color: rgba(56, 189, 248, 0.15);156 color: #38bdf8;157 border: 1px solid rgba(56, 189, 248, 0.35);158 padding: 4px 10px;159 border-radius: 6px;160 font-size: 0.85rem;161 margin: 4px 4px 4px 0;162 font-weight: 500;163}164"""165 166 167with gr.Blocks(168 title="Ask My Docs — RAG Document Assistant"169) as demo:170 171 gr.Markdown(172 """173 # 📚 Ask My Docs — RAG Document Assistant174 175 Upload domain documents (PDF / TXT) to perform semantic vector retrieval and generate grounded answers with source citations.176 """177 )178 179 with gr.Row():180 181 with gr.Column(scale=1):182 183 gr.Markdown("### 📄 Document Ingestion Hub")184 185 file_input = gr.File(186 label="Upload PDF or TXT files",187 file_count="multiple",188 file_types=[".pdf", ".txt", ".md"]189 )190 191 reset_checkbox = gr.Checkbox(192 label="Reset database before ingesting",193 value=False194 )195 196 ingest_btn = gr.Button(197 "🚀 Ingest Documents",198 variant="primary"199 )200 201 ingest_status = gr.Markdown()202 203 db_stats = gr.Markdown(204 value=get_db_stats()205 )206 207 clear_db_btn = gr.Button(208 "🗑️ Clear Vector Database",209 variant="secondary"210 )211 212 with gr.Column(scale=2):213 214 gr.Markdown("### 💬 Interactive Q&A Session")215 216 chatbot = gr.Chatbot(217 label="Chat History",218 height=450219 )220 221 query_input = gr.Textbox(222 label="Ask a detailed question about your documents...",223 placeholder="What is this document about?",224 lines=2225 )226 227 with gr.Row():228 229 top_k_slider = gr.Slider(230 minimum=1,231 maximum=10,232 value=settings.TOP_K,233 step=1,234 label="Top-K Context Chunks"235 )236 237 submit_btn = gr.Button(238 "Submit Question",239 variant="primary"240 )241 242 ingest_btn.click(243 fn=process_ingestion,244 inputs=[file_input, reset_checkbox],245 outputs=[ingest_status, db_stats]246 )247 248 submit_btn.click(249 fn=handle_query,250 inputs=[query_input, top_k_slider, chatbot],251 outputs=[chatbot, query_input]252 )253 254 query_input.submit(255 fn=handle_query,256 inputs=[query_input, top_k_slider, chatbot],257 outputs=[chatbot, query_input]258 )259 260 clear_db_btn.click(261 fn=handle_clear_db,262 inputs=[],263 outputs=[ingest_status, db_stats, chatbot]264 )265 266 267if __name__ == "__main__":268 demo.queue().launch(269 server_name="0.0.0.0",270 server_port=7860,271 theme=gr.themes.Soft(),272 css=custom_css273 )