Jainish210705/Multi-Source-RAG-System
2
1#!/usr/bin/env python32"""3streamlit_rag_app.py4 5A Hybrid Multi-Source RAG system with database support for structured data.6- Structured data (Excel) → DuckDB + Text-to-SQL7- Unstructured data (PDFs, URLs) → Vector store + RAG8 9Run with: streamlit run streamlit_rag_app.py10"""11 12import streamlit as st13import os14import tempfile15from typing import List, Tuple, Dict, Any16import pandas as pd17from dotenv import load_dotenv18import duckdb19import re20 21# LangChain imports22from langchain.document_loaders import PyPDFLoader, WebBaseLoader23from langchain.text_splitter import CharacterTextSplitter24from langchain.schema import Document25from langchain.embeddings import HuggingFaceEmbeddings26from langchain.vectorstores import Chroma27from langchain.chains import RetrievalQA28from langchain.prompts import PromptTemplate29from langchain_groq import ChatGroq30 31# Fallbacks for loading32import requests33from bs4 import BeautifulSoup34 35# Load environment variables36load_dotenv()37 38 39class HybridRAGSystem:40 """Hybrid RAG system that handles both structured (Excel) and unstructured data."""41 42 def __init__(self):43 self.conn = duckdb.connect(":memory:")44 self.vectorstore = None45 self.table_schemas = {}46 self.has_structured_data = False47 self.has_unstructured_data = False48 49 def cleanup(self):50 """Cleanup resources to free memory."""51 try:52 if self.conn:53 self.conn.close()54 except:55 pass56 57 # Clear vector store58 if self.vectorstore:59 try:60 del self.vectorstore61 self.vectorstore = None62 except:63 pass64 65 # Clear table schemas66 self.table_schemas.clear()67 self.has_structured_data = False68 self.has_unstructured_data = False69 70 71 def add_excel_to_db(self, uploaded_file, file_name: str):72 """Add Excel data to DuckDB database."""73 try:74 # Save uploaded file to temporary location75 with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_name.split('.')[-1]}") as tmp_file:76 tmp_file.write(uploaded_file.getvalue())77 tmp_file_path = tmp_file.name78 79 try:80 excel_file = pd.ExcelFile(tmp_file_path)81 sheet_names = excel_file.sheet_names82 83 st.info(f"Found {len(sheet_names)} sheet(s): {', '.join(sheet_names)}")84 85 for sheet_name in sheet_names:86 try:87 df = pd.read_excel(tmp_file_path, sheet_name=sheet_name)88 89 # Clean the data90 df = df.dropna(how='all').dropna(axis=1, how='all')91 92 if len(df) == 0:93 st.warning(f"Sheet '{sheet_name}' is empty")94 continue95 96 # Clean column names for SQL compatibility97 df.columns = [self._clean_column_name(col) for col in df.columns]98 99 # Create table name100 table_name = f"{self._clean_table_name(file_name)}_{self._clean_table_name(sheet_name)}"101 102 # Register DataFrame as a table in DuckDB103 self.conn.register(table_name, df)104 105 # Store schema information106 self.table_schemas[table_name] = {107 'original_file': file_name,108 'sheet_name': sheet_name,109 'columns': list(df.columns),110 'row_count': len(df),111 'column_count': len(df.columns),112 'sample_data': df.head(3).to_dict('records')113 }114 115 self.has_structured_data = True116 st.success(f"Added table '{table_name}' with {len(df)} rows and {len(df.columns)} columns")117 118 except Exception as e:119 st.error(f"Error processing sheet '{sheet_name}': {e}")120 continue121 finally:122 if os.path.exists(tmp_file_path):123 os.unlink(tmp_file_path)124 125 except Exception as e:126 st.error(f"Error loading Excel {file_name}: {e}")127 128 def _clean_column_name(self, col_name: str) -> str:129 """Clean column name for SQL compatibility."""130 # Convert to string and clean131 col_name = str(col_name).strip()132 # Replace spaces and special chars with underscores133 col_name = re.sub(r'[^a-zA-Z0-9_]', '_', col_name)134 # Remove multiple underscores135 col_name = re.sub(r'_+', '_', col_name)136 # Remove leading/trailing underscores137 col_name = col_name.strip('_')138 # Ensure it starts with letter or underscore139 if col_name and col_name[0].isdigit():140 col_name = f"col_{col_name}"141 return col_name or "unnamed_column"142 143 def _clean_table_name(self, name: str) -> str:144 """Clean table name for SQL compatibility."""145 # Remove file extension146 name = os.path.splitext(name)[0]147 # Clean similar to column names148 name = re.sub(r'[^a-zA-Z0-9_]', '_', name)149 name = re.sub(r'_+', '_', name)150 name = name.strip('_')151 return name or "table"152 153 def get_database_schema(self) -> str:154 """Get schema information for all tables."""155 if not self.table_schemas:156 return "No structured data available."157 158 schema_info = "DATABASE SCHEMA:\n\n"159 for table_name, info in self.table_schemas.items():160 schema_info += f"Table: {table_name}\n"161 schema_info += f" Source: {info['original_file']} (Sheet: {info['sheet_name']})\n"162 schema_info += f" Rows: {info['row_count']}, Columns: {info['column_count']}\n"163 schema_info += f" Columns: {', '.join(info['columns'])}\n"164 if info['sample_data']:165 schema_info += f" Sample data: {info['sample_data'][0]}\n"166 schema_info += "\n"167 168 return schema_info169 170 def execute_sql_query(self, sql_query: str) -> Tuple[bool, Any]:171 """Execute SQL query and return results."""172 try:173 # Basic SQL injection protection174 dangerous_keywords = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'ALTER', 'CREATE', 'TRUNCATE']175 sql_upper = sql_query.upper()176 for keyword in dangerous_keywords:177 if keyword in sql_upper:178 return False, f"Query contains potentially dangerous keyword: {keyword}"179 180 result = self.conn.execute(sql_query).fetchall()181 columns = [desc[0] for desc in self.conn.description]182 183 # Convert to DataFrame for better display184 if result:185 df_result = pd.DataFrame(result, columns=columns)186 return True, df_result187 else:188 return True, pd.DataFrame(columns=columns)189 190 except Exception as e:191 return False, f"SQL Error: {str(e)}"192 193 def generate_sql_from_question(self, question: str, api_key: str) -> str:194 """Generate SQL query from natural language question."""195 llm = ChatGroq(196 model="llama-3.3-70b-versatile",197 temperature=0.1,198 max_tokens=500,199 groq_api_key=api_key200 )201 202 schema_info = self.get_database_schema()203 204 prompt = f"""You are a SQL expert. Convert the natural language question to a SQL query based on the database schema provided.205 206{schema_info}207 208IMPORTANT RULES:2091. Only use SELECT statements - no INSERT, UPDATE, DELETE, DROP, etc.2102. Use only the table and column names provided in the schema2113. Return only the SQL query without any explanation2124. Use proper SQL syntax for DuckDB2135. If the question asks about totals, counts, averages, etc., use appropriate aggregate functions2146. For questions about "how many rows", use COUNT(*)215 216Question: {question}217 218SQL Query:"""219 220 try:221 response = llm.invoke(prompt)222 sql_query = response.content.strip()223 224 # Clean up the response - remove markdown formatting if present225 sql_query = sql_query.replace('```sql', '').replace('```', '').strip()226 227 return sql_query228 except Exception as e:229 return f"Error generating SQL: {str(e)}"230 231 232def load_pdf_from_uploaded_file(uploaded_file) -> List[Document]:233 """Load PDF from Streamlit uploaded file."""234 docs = []235 with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:236 tmp_file.write(uploaded_file.read())237 tmp_file_path = tmp_file.name238 239 try:240 loader = PyPDFLoader(tmp_file_path)241 pages = loader.load_and_split()242 243 for i, page in enumerate(pages):244 metadata = {"source": f"{uploaded_file.name}:page_{i+1}", "type": "pdf"}245 docs.append(Document(page_content=page.page_content, metadata=metadata))246 except Exception as e:247 st.error(f"Error loading PDF {uploaded_file.name}: {e}")248 finally:249 if os.path.exists(tmp_file_path):250 os.unlink(tmp_file_path)251 252 return docs253 254 255def load_url(url: str) -> List[Document]:256 """Load web page from URL."""257 docs = []258 try:259 loader = WebBaseLoader(url)260 loaded_docs = loader.load()261 262 for doc in loaded_docs:263 metadata = {"source": url, "type": "url"}264 docs.append(Document(page_content=doc.page_content, metadata=metadata))265 266 except Exception:267 try:268 response = requests.get(url, timeout=15)269 response.raise_for_status()270 soup = BeautifulSoup(response.text, 'html.parser')271 272 text = '\n'.join([p.get_text().strip() for p in soup.find_all('p') if p.get_text().strip()])273 if not text:274 text = soup.get_text()275 276 docs.append(Document(page_content=text, metadata={'source': url, "type": "url"}))277 except Exception as e:278 st.error(f"Failed to fetch {url}: {e}")279 280 return docs281 282 283def chunk_documents(docs: List[Document], chunk_size: int = 1000, chunk_overlap: int = 200, max_chunks: int = 1000) -> List[Document]:284 """Split documents into smaller chunks with optional limit."""285 splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)286 chunked_docs = []287 288 for doc in docs:289 chunks = splitter.split_text(doc.page_content)290 for i, chunk in enumerate(chunks):291 # Stop if we've reached max chunks292 if len(chunked_docs) >= max_chunks:293 st.warning(f"⚠️ Reached maximum chunk limit ({max_chunks}). Some content may be skipped.")294 return chunked_docs295 296 metadata = dict(doc.metadata)297 metadata['chunk'] = i298 chunked_docs.append(Document(page_content=chunk, metadata=metadata))299 300 return chunked_docs301 302 303def create_vectorstore(docs: List[Document]):304 """Create vector store for unstructured data (in-memory only)."""305 if not docs:306 return None307 308 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')309 310 # Use in-memory vectorstore to avoid disk storage issues311 # Setting persist_directory=None ensures it stays in memory312 vectorstore = Chroma.from_documents(313 documents=docs, 314 embedding=embeddings,315 collection_name="rag_collection",316 # Do NOT set persist_directory - keeps everything in RAM317 )318 return vectorstore319 320 321def classify_question(question: str, has_structured: bool, has_unstructured: bool) -> str:322 """Classify if question is about structured data, unstructured data, or both."""323 if not has_structured:324 return "unstructured"325 if not has_unstructured:326 return "structured"327 328 # Keywords that suggest structured data queries329 structured_keywords = [330 'how many rows', 'total rows', 'count', 'total', 'sum', 'average', 'mean', 331 'maximum', 'minimum', 'group by', 'aggregate', 'statistics', 'excel', 332 'sheet', 'table', 'column', 'data', 'number of', 'calculate'333 ]334 335 question_lower = question.lower()336 for keyword in structured_keywords:337 if keyword in question_lower:338 return "structured"339 340 return "unstructured"341 342 343def query_hybrid_system(question: str, rag_system: HybridRAGSystem, api_key: str, top_k: int = 4, similarity_threshold: float = 0.65) -> Tuple[str, List]:344 """Query the hybrid RAG system with Transparent Hybrid behavior."""345 346 # Classify the question347 question_type = classify_question(348 question, 349 rag_system.has_structured_data, 350 rag_system.has_unstructured_data351 )352 353 if question_type == "structured":354 # Handle structured data query (same as before)355 try:356 sql_query = rag_system.generate_sql_from_question(question, api_key)357 358 if sql_query.startswith("Error"):359 return sql_query, []360 361 success, result = rag_system.execute_sql_query(sql_query)362 363 if success:364 llm = ChatGroq(365 model="llama-3.3-70b-versatile",366 temperature=0.3,367 max_tokens=1000,368 groq_api_key=api_key369 )370 371 response_prompt = f"""Convert this SQL query result into a natural, conversational response to answer the user's question.372Original Question: {question}373SQL Query Used: {sql_query}374Query Results: 375{result.to_string() if not result.empty else "No results found"}376Provide a clear, natural language answer that directly addresses the user's question. Include specific numbers and details from the results."""377 378 try:379 response = llm.invoke(response_prompt)380 answer = "According to your Excel data:\n\n" + response.content.strip()381 return answer, []382 except Exception as e:383 return f"Results found but error formatting response: {str(e)}\n\nRaw results:\n{result.to_string()}", []384 else:385 return f"Database query failed: {result}", []386 387 except Exception as e:388 return f"Error processing structured data query: {str(e)}", []389 390 else:391 # Handle unstructured (PDF/URL) with Transparent Hybrid392 if not rag_system.vectorstore:393 return "No unstructured documents available for querying.", []394 395 retriever = rag_system.vectorstore.as_retriever(search_kwargs={"k": top_k})396 docs_with_scores = rag_system.vectorstore.similarity_search_with_score(question, k=top_k)397 398 # Check if any document passes the similarity threshold399 relevant_docs = [doc for doc, score in docs_with_scores if score >= similarity_threshold]400 401 llm = ChatGroq(402 model="llama-3.3-70b-versatile",403 temperature=0.3,404 max_tokens=1500,405 groq_api_key=api_key406 )407 408 if relevant_docs:409 # Use RAG answer410 context = "\n\n".join([doc.page_content for doc in relevant_docs])411 sources = relevant_docs412 413 prompt = f"""You are an intelligent assistant.414According to the user's uploaded documents, provide a helpful answer.415 416Context from documents:417{context}418 419Question: {question}420 421Answer clearly, based ONLY on the above documents. If possible, reference details from them."""422 423 try:424 response = llm.invoke(prompt)425 return "According to your uploaded documents:\n\n" + response.content.strip(), sources426 except Exception as e:427 return f"Error generating document-based answer: {e}", sources428 429 else:430 # Fallback to general LLM431 fallback_prompt = f"""The user's question is unrelated to their uploaded documents.432 433Question: {question}434 435Provide a clear, general knowledge answer, but explicitly say it's NOT from their documents."""436 437 try:438 response = llm.invoke(fallback_prompt)439 return "This information was not found in your uploaded documents. From general knowledge:\n\n" + response.content.strip(), []440 except Exception as e:441 return f"Error generating fallback answer: {e}", []442 443 444def main():445 st.set_page_config(446 page_title="Multi-Source RAG System",447 layout="wide"448 )449 450 st.title("Multi-Source RAG System")451 st.markdown("**Intelligent document analysis with specialized handling for Excel data and text documents!**")452 453 # Initialize session state - do this ONCE454 if 'rag_system' not in st.session_state:455 st.session_state.rag_system = HybridRAGSystem()456 457 # Sidebar configuration458 with st.sidebar:459 st.header("Configuration")460 461 # API Key management462 env_api_key = os.getenv("GROQ_API_KEY") or os.getenv("GRQO_API_KEY")463 464 if env_api_key:465 st.success("API key loaded from environment")466 api_key = env_api_key467 masked_key = env_api_key[:8] + "..." + env_api_key[-4:] if len(env_api_key) > 12 else "***"468 st.text(f"Using key: {masked_key}")469 else:470 st.warning("API key not found in environment")471 api_key = st.text_input(472 "Enter Groq API Key", 473 type="password", 474 help="Enter your Groq API key for LLM functionality."475 )476 477 st.header("System Status")478 if st.session_state.rag_system.has_structured_data:479 st.success("Excel data loaded in database")480 st.info(f"{len(st.session_state.rag_system.table_schemas)} table(s) available")481 else:482 st.info("No Excel data loaded")483 484 if st.session_state.rag_system.has_unstructured_data:485 st.success("Text documents loaded in vector store")486 else:487 st.info("No text documents loaded")488 489 # Reset button to clear memory490 if st.button("Clear All Data", help="Remove all processed data and free memory"):491 st.session_state.rag_system.cleanup()492 st.session_state.rag_system = HybridRAGSystem()493 st.success("All data cleared!")494 st.rerun()495 496 st.header("🔧 Document Processing")497 chunk_size = st.slider("Chunk Size (Text docs)", 500, 2000, 1000, step=100)498 chunk_overlap = st.slider("Chunk Overlap (Text docs)", 50, 500, 200, step=50)499 top_k = st.slider("Top K Results (Text search)", 1, 10, 4)500 501 # Database schema viewer502 if st.session_state.rag_system.table_schemas:503 with st.expander("📋 View Database Schema"):504 st.text(st.session_state.rag_system.get_database_schema())505 506 # Main content - DEFINE VARIABLES AT TOP LEVEL507 # Excel files - will go to database508 st.subheader("Excel Files")509 uploaded_excels = st.file_uploader(510 "Upload Excel files for structured querying",511 type=['xlsx', 'xls'],512 accept_multiple_files=True,513 help="Excel data will be stored in database for SQL-like queries"514 )515 516 # PDF files - will go to vector store517 st.subheader("PDF Files")518 uploaded_pdfs = st.file_uploader(519 "Upload PDF files for text analysis",520 type=['pdf'],521 accept_multiple_files=True,522 help="PDFs will be processed for semantic text search"523 )524 525 # URLs - will go to vector store526 st.subheader("URLs")527 url_input = st.text_area(528 "Enter URLs (one per line)",529 placeholder="https://example.com\nhttps://another-site.com",530 help="Web content will be processed for semantic text search"531 )532 533 # Process documents button534 if st.button("Process All Documents", type="primary"):535 if not api_key:536 st.error("Please provide a Groq API key.")537 else:538 # Process Excel files539 if uploaded_excels:540 with st.spinner("Processing Excel files into database..."):541 for excel in uploaded_excels:542 st.session_state.rag_system.add_excel_to_db(excel, excel.name)543 544 # Process unstructured documents545 all_docs = []546 547 # Process PDFs548 if uploaded_pdfs:549 with st.spinner("Processing PDF files..."):550 for pdf in uploaded_pdfs:551 docs = load_pdf_from_uploaded_file(pdf)552 all_docs.extend(docs)553 st.success(f"Loaded {len(docs)} pages from {pdf.name}")554 555 # Process URLs556 if url_input.strip():557 urls = [url.strip() for url in url_input.strip().split('\n') if url.strip()]558 with st.spinner("Fetching URLs..."):559 for url in urls:560 docs = load_url(url)561 all_docs.extend(docs)562 st.success(f"Loaded content from {url}")563 564 # Create vector store for unstructured data565 if all_docs:566 with st.spinner("Creating vector store for text documents..."):567 chunked_docs = chunk_documents(all_docs, chunk_size, chunk_overlap)568 st.session_state.rag_system.vectorstore = create_vectorstore(chunked_docs)569 st.session_state.rag_system.has_unstructured_data = True570 st.success(f"Processed {len(all_docs)} documents into {len(chunked_docs)} chunks!")571 572 st.markdown("---")573 574 # Query section575 st.header("Ask Questions")576 577 # Check if system is ready578 system_ready = (st.session_state.rag_system.has_structured_data or 579 st.session_state.rag_system.has_unstructured_data)580 581 if system_ready:582 st.success("System ready for queries!")583 584 # Question input585 question = st.text_input(586 "Enter your question:",587 placeholder="How many rows are in the Excel sheet? OR What does the document say about...?",588 help="Ask about Excel data (counts, totals, analysis) or document content"589 )590 591 if st.button("Get Answer", type="primary"):592 if not question.strip():593 st.warning("Please enter a question.")594 elif not api_key:595 st.error("Please provide a Groq API key.")596 else:597 with st.spinner("Analyzing question and generating answer..."):598 answer, source_docs = query_hybrid_system(599 question,600 st.session_state.rag_system,601 api_key,602 top_k603 )604 605 st.subheader("Answer:")606 st.write(answer)607 608 # Show sources for unstructured data609 if source_docs:610 st.subheader("Sources:")611 for i, doc in enumerate(source_docs, 1):612 source = doc.metadata.get('source', 'unknown')613 doc_type = doc.metadata.get('type', 'unknown')614 st.write(f"{i}. {source} ({doc_type})")615 616 with st.expander(f"View excerpt {i}"):617 st.write(doc.page_content[:500] + "..." if len(doc.page_content) > 500 else doc.page_content)618 else:619 # Check if files are uploaded but not processed620 has_uploads = uploaded_pdfs or (url_input and url_input.strip()) or uploaded_excels621 if has_uploads:622 st.warning("Documents uploaded but not processed yet. Please click '🔄 Process All Documents' button first!")623 else:624 st.info("Please upload documents first!")625 626 # Example questions627 if system_ready:628 st.subheader("Example Questions")629 630 example_questions = []631 if st.session_state.rag_system.has_structured_data:632 example_questions.extend([633 "How many rows are in the Excel sheet?",634 "What is the total/sum of [column name]?",635 "Show me the average of [column name]",636 "Group the data by [column name]"637 ])638 639 if st.session_state.rag_system.has_unstructured_data:640 example_questions.extend([641 "Summarize the main points from the documents",642 "What are the key findings mentioned?",643 "Explain the methodology described"644 ])645 646 for eq in example_questions[:6]: # Limit to 6 examples647 st.button(eq, key=f"example_{hash(eq)}")648 649 # Footer650 st.markdown("---")651 st.markdown("**Hybrid Approach:**")652 st.markdown("- **Excel data** → Database queries for precise calculations and counts")653 st.markdown("- **PDF/Web content** → Semantic search for contextual understanding") 654 st.markdown("- **Intelligent routing** → System automatically chooses the best approach")655 st.markdown("---")656 657if __name__ == "__main__":658 main()