#!/usr/bin/env python3 """ # RAG System for Hugging Face Spaces A simplified Retrieval-Augmented Generation (RAG) system using: - **FAISS** for vector search and similarity matching - **BM25** for keyword-based sparse retrieval - **Hybrid Search** combining both dense and sparse methods - **Streamlit** for modern, interactive web interface - **Qwen 2.5 1.5B** for intelligent response generation ## Features - 🔍 **Multi-Method Retrieval**: Hybrid, dense, and sparse search options - 📄 **PDF Processing**: Automatic document loading and chunking - 💬 **Real-time Chat**: Interactive conversation interface - ⚡ **Parallel Loading**: Concurrent document processing - 📊 **Performance Metrics**: Response times and confidence scores - 🎯 **Smart Fallbacks**: Graceful handling of model loading failures ## Architecture The system follows a modular architecture: 1. **Document Processing**: PDF extraction and chunking 2. **Vector Storage**: FAISS index for embeddings 3. **Search Engine**: BM25 for keyword matching 4. **Response Generation**: LLM-based answer synthesis 5. **Web Interface**: Streamlit for user interaction ## Usage 1. Upload PDF documents or use pre-loaded ones 2. Choose retrieval method (hybrid/dense/sparse) 3. Ask questions in natural language 4. View answers with source citations and confidence scores """ import streamlit as st import os import tempfile from pathlib import Path import time from typing import List, Dict, Optional import json import glob from concurrent.futures import ThreadPoolExecutor, as_completed from loguru import logger # Import our simplified components from rag_system import SimpleRAGSystem from pdf_processor import SimplePDFProcessor from hf_spaces_config import get_hf_config, is_hf_spaces from guard_rails import GuardRailConfig # ============================================================================= # PAGE CONFIGURATION # ============================================================================= # Configure Streamlit page settings for optimal user experience st.set_page_config( page_title="RAG System - Hugging Face", page_icon="🤖", layout="wide", # Use full width for better content display initial_sidebar_state="expanded", # Show sidebar by default ) # ============================================================================= # SESSION STATE INITIALIZATION # ============================================================================= # Initialize Streamlit session state for persistent data across interactions if "rag_system" not in st.session_state: st.session_state.rag_system = None # Main RAG system instance if "documents_loaded" not in st.session_state: st.session_state.documents_loaded = False # Document loading status if "chat_history" not in st.session_state: st.session_state.chat_history = [] # Conversation history if "initializing" not in st.session_state: st.session_state.initializing = False # Initialization status # ============================================================================= # UTILITY FUNCTIONS # ============================================================================= def display_environment_info(): """ Display information about the current deployment environment """ if is_hf_spaces(): st.sidebar.markdown("### 🌐 Environment") st.sidebar.info("**Hugging Face Spaces**") # Get HF Spaces configuration details try: hf_config = get_hf_config() st.sidebar.markdown("**Configuration:**") st.sidebar.text( f"• Cache: {hf_config.cache_dirs.get('transformers_cache', 'N/A')}" ) st.sidebar.text( f"• Vector Store: {hf_config.cache_dirs.get('vector_store', 'N/A')}" ) # Show resource limits resource_limits = hf_config.get_resource_limits() st.sidebar.markdown("**Resource Limits:**") st.sidebar.text(f"• Memory: {resource_limits['max_memory_usage']*100:.0f}%") st.sidebar.text(f"• CPU: {resource_limits['max_cpu_usage']*100:.0f}%") st.sidebar.text( f"• Concurrent: {resource_limits['max_concurrent_requests']}" ) except Exception as e: st.sidebar.warning(f"Config error: {e}") else: st.sidebar.markdown("### 💻 Environment") st.sidebar.info("**Local Development**") def load_single_document(rag_system, pdf_path): """ Load a single document into the RAG system Args: rag_system: The RAG system instance pdf_path: Path to the PDF file Returns: tuple: (filename, success_status, error_message) """ try: filename = os.path.basename(pdf_path) success = rag_system.add_document(pdf_path, filename) return filename, success, None except Exception as e: return os.path.basename(pdf_path), False, str(e) def initialize_rag_system(): """ Initialize the RAG system with automatic document loading This function: 1. Creates the RAG system instance 2. Automatically loads all available PDF documents 3. Uses parallel processing for faster loading 4. Provides real-time feedback on loading progress """ if st.session_state.rag_system is None and not st.session_state.initializing: st.session_state.initializing = True st.write("🚀 Starting RAG system initialization...") # Check deployment environment if is_hf_spaces(): st.info("🌐 Running in Hugging Face Spaces environment") st.write("📁 Setting up HF Spaces optimized configuration...") else: st.info("💻 Running in local development environment") st.write("📁 Using local development configuration...") with st.spinner("Initializing RAG system..."): try: # Get HF Spaces configuration hf_config = get_hf_config() model_config = hf_config.get_model_config() guard_config = GuardRailConfig(**hf_config.get_guard_rail_config()) # Create RAG system instance with HF Spaces optimized settings st.session_state.rag_system = SimpleRAGSystem( embedding_model=model_config["embedding_model"], generative_model=model_config["generative_model"], chunk_sizes=model_config["chunk_sizes"], vector_store_path=model_config["vector_store_path"], enable_guard_rails=model_config["enable_guard_rails"], guard_rail_config=guard_config, ) st.write("✅ RAG system created successfully") # Auto-load all available PDF documents in parallel pdf_files = glob.glob("/app/*.pdf") st.write(f"📁 Found {len(pdf_files)} PDF files") if pdf_files: loaded_count = 0 failed_count = 0 with st.spinner( f"Loading {len(pdf_files)} PDF documents in parallel..." ): # Use ThreadPoolExecutor for parallel loading # This significantly speeds up document processing with ThreadPoolExecutor(max_workers=4) as executor: # Submit all document loading tasks future_to_pdf = { executor.submit( load_single_document, st.session_state.rag_system, pdf_path, ): pdf_path for pdf_path in pdf_files } # Process completed tasks and provide real-time feedback for future in as_completed(future_to_pdf): filename, success, error = future.result() if success: loaded_count += 1 st.write(f"✅ Loaded: {filename}") logger.info(f"✅ Loaded: {filename}") else: failed_count += 1 st.write(f"⚠️ Failed: {filename} - {error}") logger.warning( f"⚠️ Failed to load {filename}: {error}" ) # Update system status based on loading results if loaded_count > 0: st.session_state.documents_loaded = True st.success( f"✅ Successfully loaded {loaded_count} PDF documents!" ) if failed_count > 0: st.warning(f"⚠️ Failed to load {failed_count} documents") else: st.warning("⚠️ No documents could be loaded") # Still allow querying even if no documents loaded st.session_state.documents_loaded = True else: st.info("📚 No PDF documents found in the container") # Still allow querying even if no documents found st.session_state.documents_loaded = True st.success("✅ RAG system initialized!") except Exception as e: st.error(f"❌ Failed to initialize RAG system: {e}") logger.error(f"RAG system initialization failed: {e}") # Reset initialization flag on error st.session_state.initializing = False raise finally: # Always reset initialization flag st.session_state.initializing = False def upload_document(uploaded_file): """ Upload and process a document through the web interface Args: uploaded_file: Streamlit uploaded file object """ if uploaded_file is not None: try: # Create temporary file for processing with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file: tmp_file.write(uploaded_file.getvalue()) tmp_path = tmp_file.name # Process the document with progress feedback with st.spinner(f"Processing {uploaded_file.name}..."): success = st.session_state.rag_system.add_document( tmp_path, uploaded_file.name ) if success: st.success(f"✅ {uploaded_file.name} processed successfully!") st.session_state.documents_loaded = True # Clean up temporary file os.unlink(tmp_path) else: st.error(f"❌ Failed to process {uploaded_file.name}") os.unlink(tmp_path) except Exception as e: st.error(f"❌ Error processing document: {str(e)}") def query_rag( query: str, method: str = "hybrid", top_k: int = 5, user_id: str = "anonymous" ): """ Query the RAG system with detailed logging and error handling Args: query: User's question method: Retrieval method (hybrid/dense/sparse) top_k: Number of results to retrieve user_id: User identifier for guard rail tracking Returns: tuple: (response_object, response_time) """ try: st.write(f"🔍 Starting query: {query}") st.write(f"🔍 Method: {method}, top_k: {top_k}") if st.session_state.rag_system is None: st.error("❌ RAG system is not initialized") return None, "RAG system not initialized" st.write(f"✅ RAG system is available") start_time = time.time() st.write(f"🔍 Calling rag_system.query with guard rails...") response = st.session_state.rag_system.query(query, method, top_k, user_id) response_time = time.time() - start_time st.write(f"✅ Response received in {response_time:.2f}s") st.write(f"✅ Response type: {type(response)}") if response: st.write(f"✅ Response answer: {response.answer[:100]}...") return response, response_time except Exception as e: st.error(f"❌ Error during query: {str(e)}") logger.error(f"Query error: {e}") import traceback st.error(f"❌ Full error: {traceback.format_exc()}") return None, f"Error: {str(e)}" def display_search_results(results: List[Dict]): """ Display search results with detailed information and metrics Args: results: List of search result dictionaries """ if not results: st.info("No search results found.") return # Display each search result with comprehensive information for i, result in enumerate(results, 1): st.markdown(f"---") st.markdown(f"**Result {i}** - Score: {result.score:.3f}") st.write(f"**Source:** {result.filename}") st.write(f"**Method:** {result.search_method}") st.write(f"**Text:** {result.text[:500]}...") # Show detailed scores for hybrid search if result.dense_score and result.sparse_score: col1, col2 = st.columns(2) with col1: st.metric("Dense Score", f"{result.dense_score:.3f}") with col2: st.metric("Sparse Score", f"{result.sparse_score:.3f}") # ============================================================================= # MAIN APPLICATION # ============================================================================= def main(): """ Main application function that orchestrates the entire RAG system interface This function: 1. Sets up the user interface 2. Initializes the RAG system 3. Handles document uploads 4. Manages the chat interface 5. Displays results and metrics """ st.write("🚀 App starting...") # Display environment information in sidebar display_environment_info() st.title("🤖 RAG System - Hugging Face Spaces") st.markdown("A simplified RAG system using FAISS + BM25 + Qwen 2.5 1.5B") # Initialize RAG system initialize_rag_system() # ============================================================================= # SIDEBAR CONFIGURATION # ============================================================================= with st.sidebar: st.header("📁 Document Upload") # File uploader for PDF documents uploaded_file = st.file_uploader( "Upload PDF Document", type=["pdf"], help="Upload a PDF document to add to the knowledge base", ) if uploaded_file: upload_document(uploaded_file) st.divider() st.header("⚙️ Settings") # Retrieval method selection method = st.selectbox( "Retrieval Method", ["hybrid", "dense", "sparse"], help="Choose the retrieval method: hybrid (combines dense and sparse), dense (vector similarity), or sparse (keyword matching)", ) # Number of results slider top_k = st.slider( "Number of Results", min_value=1, max_value=10, value=5, help="Number of top results to retrieve and use for answer generation", ) st.divider() # System information display if st.session_state.rag_system: stats = st.session_state.rag_system.get_stats() st.header("📊 System Info") st.write(f"**Documents:** {stats['total_documents']}") st.write(f"**Chunks:** {stats['total_chunks']}") st.write(f"**Vector Size:** {stats['vector_size']}") st.write(f"**Model:** {stats['model_name']}") # ============================================================================= # MAIN CONTENT AREA # ============================================================================= # Initialize RAG system if not already done if not st.session_state.rag_system: if st.session_state.initializing: st.info("🔄 RAG system is initializing... Please wait.") return else: initialize_rag_system() return # Show system info and allow querying immediately after initialization stats = st.session_state.rag_system.get_stats() documents_available = stats["total_documents"] > 0 if not documents_available: st.info( "📚 No documents loaded yet, but you can still ask questions. The system will respond based on its general knowledge." ) # ============================================================================= # CHAT INTERFACE # ============================================================================= st.header("💬 Ask Questions About Your Documents") # Chat input for user questions query = st.chat_input("Ask a question about the loaded documents...") if query: st.write(f"📝 Processing query: {query}") # Add user message to chat history st.session_state.chat_history.append({"role": "user", "content": query}) # Get response from RAG system response, response_time = query_rag(query, method, top_k) st.write(f"📊 Response type: {type(response)}") st.write(f"📊 Response time: {response_time}") if response: st.write("✅ Got valid response, adding to chat history") # Add assistant response to chat history with metadata st.session_state.chat_history.append( { "role": "assistant", "content": response.answer, "search_results": response.search_results, "method_used": response.method_used, "confidence": response.confidence, "response_time": response_time, } ) else: st.write("❌ No valid response received") st.session_state.chat_history.append( {"role": "assistant", "content": f"Error: {response_time}"} ) # ============================================================================= # CHAT HISTORY DISPLAY # ============================================================================= # Display conversation history with detailed information for message in st.session_state.chat_history: if message["role"] == "user": with st.chat_message("user"): st.write(message["content"]) else: with st.chat_message("assistant"): st.write(message["content"]) # Show additional information for assistant messages if "search_results" in message: st.markdown("**🔍 Search Results:**") display_search_results(message["search_results"]) # Display performance metrics col1, col2, col3 = st.columns(3) with col1: st.metric("Method", message["method_used"]) with col2: st.metric("Confidence", f"{message['confidence']:.3f}") with col3: st.metric("Response Time", f"{message['response_time']:.2f}s") # ============================================================================= # UTILITY CONTROLS # ============================================================================= # Clear chat history button if st.session_state.chat_history: if st.button("🗑️ Clear Chat History"): st.session_state.chat_history = [] st.rerun() # ============================================================================= # APPLICATION ENTRY POINT # ============================================================================= if __name__ == "__main__": main()