Spaces:
Running
Running
| """ | |
| Main Gradio interface for the Professional RAG Assistant. | |
| """ | |
| import gradio as gr | |
| import asyncio | |
| import threading | |
| import time | |
| import json | |
| import sys | |
| import signal | |
| import logging | |
| from typing import Any, Dict, List, Optional, Tuple, Union | |
| from pathlib import Path | |
| from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError | |
| from .themes import get_theme, get_custom_css | |
| from .components import ( | |
| create_header, create_file_upload_section, create_search_interface, | |
| create_results_display, create_document_management, create_system_status, | |
| create_analytics_dashboard, format_document_list, format_search_results, | |
| create_analytics_charts, format_system_overview, create_error_display, | |
| create_success_display, create_loading_display | |
| ) | |
| from .utils import ( | |
| save_uploaded_files, validate_file_types, cleanup_temp_files, | |
| generate_session_id, format_file_size, parse_search_filters | |
| ) | |
| sys.path.append(str(Path(__file__).parent.parent)) | |
| from src.rag_system import RAGSystem | |
| from src.error_handler import RAGError | |
| class RAGInterface: | |
| """Main interface class for the RAG system.""" | |
| def __init__(self, config_path: str = None): | |
| """Initialize the RAG interface.""" | |
| self.rag_system: Optional[RAGSystem] = None | |
| self.config_path = config_path or "config.yaml" | |
| self.active_sessions: Dict[str, Dict[str, Any]] = {} | |
| self._initialization_lock = threading.Lock() | |
| self._initialized = False | |
| # Setup logging | |
| self.logger = logging.getLogger(__name__) | |
| # Initialize RAG system | |
| self._initialize_rag_system() | |
| def _initialize_rag_system(self) -> None: | |
| """Initialize the RAG system.""" | |
| try: | |
| with self._initialization_lock: | |
| if not self._initialized: | |
| print("Initializing RAG system...") | |
| self.rag_system = RAGSystem(config_path=self.config_path) | |
| # Warm up the system | |
| warmup_result = self.rag_system.warmup() | |
| if warmup_result.get("success"): | |
| print("RAG system initialized and warmed up successfully") | |
| self._initialized = True | |
| else: | |
| print(f"RAG system warmup failed: {warmup_result.get('error', {}).get('message')}") | |
| except Exception as e: | |
| print(f"Failed to initialize RAG system: {e}") | |
| self.rag_system = None | |
| def get_system_status(self) -> str: | |
| """Get current system status HTML.""" | |
| if not self.rag_system or not self._initialized: | |
| return create_error_display("System not initialized. Please check configuration and restart.") | |
| try: | |
| stats_response = self.rag_system.get_system_stats() | |
| if not stats_response.get("success"): | |
| return create_error_display(f"Failed to get system status: {stats_response.get('error', {}).get('message')}") | |
| stats = stats_response["data"] | |
| status_info = stats.get("status", {}) | |
| if status_info.get("ready"): | |
| status_message = f"System ready - {status_info.get('documents_indexed', 0)} documents indexed" | |
| return create_success_display(status_message) | |
| else: | |
| return create_error_display("System not ready") | |
| except Exception as e: | |
| return create_error_display(f"Error getting system status: {str(e)}") | |
| def process_documents( | |
| self, | |
| files: List[gr.File], | |
| session_id: str, | |
| progress=gr.Progress() | |
| ) -> Tuple[str, str, bool, str]: | |
| """Process uploaded documents.""" | |
| if not files: | |
| return ( | |
| create_error_display("No files uploaded"), | |
| create_error_display("Please select files to upload"), | |
| False, # upload button disabled | |
| "No documents uploaded yet." | |
| ) | |
| if not self.rag_system or not self._initialized: | |
| return ( | |
| create_error_display("System not initialized"), | |
| create_error_display("Please restart the application"), | |
| False, | |
| "No documents uploaded yet." | |
| ) | |
| try: | |
| self.logger.info(f"Starting document upload process - {len(files)} files received") | |
| # Validate file types | |
| allowed_extensions = [".pdf", ".docx", ".txt"] | |
| valid_files, validation_errors = validate_file_types(files, allowed_extensions) | |
| if validation_errors: | |
| self.logger.warning(f"File validation errors: {validation_errors}") | |
| error_html = create_error_display("\\n".join(validation_errors)) | |
| return error_html, error_html, len(valid_files) > 0, self.get_document_list() | |
| self.logger.info(f"File validation passed - {len(valid_files)} valid files") | |
| # Save uploaded files | |
| progress(0.1, desc="Saving uploaded files...") | |
| self.logger.info("Saving uploaded files to temporary directory...") | |
| saved_files = save_uploaded_files(valid_files) | |
| for file_path, original_name in saved_files: | |
| file_size = Path(file_path).stat().st_size / (1024 * 1024) # Size in MB | |
| self.logger.info(f"Saved file: {original_name} ({file_size:.2f} MB) -> {file_path}") | |
| if not saved_files: | |
| return ( | |
| create_error_display("No valid files to process"), | |
| create_error_display("Please check your files and try again"), | |
| False, | |
| self.get_document_list() | |
| ) | |
| # Process each file with timeout | |
| processed_count = 0 | |
| total_files = len(saved_files) | |
| results = [] | |
| timeout_seconds = 600 # 10 minutes | |
| def process_single_file(file_path, original_name, session_id): | |
| """Process a single file - to be run with timeout.""" | |
| self.logger.info(f"Processing file: {original_name}") | |
| start_time = time.time() | |
| result = self.rag_system.add_document( | |
| file_path=file_path, | |
| filename=original_name, | |
| user_session=session_id | |
| ) | |
| processing_time = time.time() - start_time | |
| self.logger.info(f"File processing completed: {original_name} (took {processing_time:.2f}s)") | |
| return result | |
| self.logger.info(f"Starting processing of {total_files} files with {timeout_seconds//60}-minute timeout per file") | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| for i, (file_path, original_name) in enumerate(saved_files): | |
| progress((i + 1) / total_files * 0.8 + 0.1, desc=f"Processing {original_name}...") | |
| self.logger.info(f"Processing file {i+1}/{total_files}: {original_name}") | |
| try: | |
| # Submit the task with timeout | |
| future = executor.submit(process_single_file, file_path, original_name, session_id) | |
| result = future.result(timeout=timeout_seconds) | |
| if result.get("success"): | |
| processed_count += 1 | |
| chunks_created = result['data']['chunks_created'] | |
| # Log detailed success info | |
| self.logger.info(f"SUCCESS: {original_name} - {chunks_created} chunks created") | |
| # Log sample chunk info if available | |
| if 'sample_chunks' in result['data']: | |
| sample_chunks = result['data']['sample_chunks'] | |
| self.logger.info(f"Sample chunks from {original_name}:") | |
| for idx, chunk in enumerate(sample_chunks[:3]): # Show first 3 chunks | |
| chunk_preview = chunk['content'][:100] + "..." if len(chunk['content']) > 100 else chunk['content'] | |
| self.logger.info(f" Chunk {idx}: {chunk_preview}") | |
| results.append(f"✅ {original_name}: {chunks_created} chunks created") | |
| else: | |
| error_msg = result.get("error", {}).get("message", "Unknown error") | |
| self.logger.error(f"FAILED: {original_name} - {error_msg}") | |
| results.append(f"❌ {original_name}: {error_msg}") | |
| except FutureTimeoutError: | |
| self.logger.error(f"TIMEOUT: {original_name} exceeded {timeout_seconds//60} minute limit") | |
| results.append(f"⏰ {original_name}: Processing timed out after {timeout_seconds//60} minutes") | |
| future.cancel() # Cancel the task if possible | |
| except Exception as e: | |
| self.logger.error(f"EXCEPTION: {original_name} - {str(e)}") | |
| results.append(f"❌ {original_name}: {str(e)}") | |
| progress(1.0, desc="Cleaning up...") | |
| self.logger.info("Cleaning up temporary files...") | |
| # Clean up temporary files | |
| cleanup_temp_files([fp for fp, _ in saved_files]) | |
| # Log final summary | |
| total_processing_time = time.time() - time.time() # This will be updated properly | |
| self.logger.info(f"Document upload process completed:") | |
| self.logger.info(f" - Total files: {total_files}") | |
| self.logger.info(f" - Successfully processed: {processed_count}") | |
| self.logger.info(f" - Failed: {total_files - processed_count}") | |
| self.logger.info(f" - Success rate: {(processed_count/total_files*100):.1f}%") | |
| # Create result message | |
| if processed_count == total_files: | |
| self.logger.info(f"✅ ALL DOCUMENTS PROCESSED SUCCESSFULLY ({processed_count}/{total_files})") | |
| status_html = create_success_display( | |
| f"Successfully processed {processed_count} documents:\\n" + "\\n".join(results) | |
| ) | |
| upload_status = create_success_display(f"All {processed_count} documents processed successfully!") | |
| elif processed_count > 0: | |
| self.logger.warning(f"⚠️ PARTIAL SUCCESS ({processed_count}/{total_files} documents processed)") | |
| status_html = f""" | |
| <div style='background: #fef3c7; border: 1px solid #f59e0b; border-radius: 8px; padding: 1rem; margin: 1rem 0;'> | |
| <div style='font-weight: 600; color: #92400e; margin-bottom: 0.5rem;'> | |
| ⚠️ Partially successful ({processed_count}/{total_files} files processed) | |
| </div> | |
| <div style='color: #78350f; font-size: 0.9rem;'>{"<br>".join(results)}</div> | |
| </div> | |
| """ | |
| upload_status = status_html | |
| else: | |
| self.logger.error(f"❌ NO DOCUMENTS PROCESSED SUCCESSFULLY (0/{total_files})") | |
| status_html = create_error_display( | |
| f"Failed to process any documents:\\n" + "\\n".join(results) | |
| ) | |
| upload_status = create_error_display("Document processing failed") | |
| return ( | |
| status_html, | |
| upload_status, | |
| gr.update(interactive=True), # Enable search button | |
| self.get_document_list() | |
| ) | |
| except Exception as e: | |
| # Clean up on error | |
| try: | |
| if 'saved_files' in locals(): | |
| cleanup_temp_files([fp for fp, _ in saved_files]) | |
| except: | |
| pass | |
| error_message = f"Document processing failed: {str(e)}" | |
| error_html = create_error_display(error_message) | |
| return error_html, error_html, gr.update(interactive=False), self.get_document_list() | |
| def perform_search( | |
| self, | |
| query: str, | |
| search_mode: str, | |
| num_results: int, | |
| enable_reranking: bool, | |
| metadata_filters: str, | |
| session_id: str | |
| ) -> Tuple[str, str, str]: | |
| """Perform search and return results.""" | |
| if not self.rag_system or not self._initialized: | |
| error_html = create_error_display("System not initialized") | |
| return error_html, "{}", "" | |
| if not query or not query.strip(): | |
| error_html = create_error_display("Please enter a search query") | |
| return error_html, "{}", "" | |
| try: | |
| # Parse metadata filters | |
| filters = parse_search_filters(metadata_filters) if metadata_filters else None | |
| # Perform search | |
| result = self.rag_system.search( | |
| query=query.strip(), | |
| k=num_results, | |
| search_mode=search_mode, | |
| enable_reranking=enable_reranking, | |
| metadata_filter=filters, | |
| user_session=session_id | |
| ) | |
| if not result.get("success"): | |
| error_msg = result.get("error", {}).get("message", "Search failed") | |
| error_html = create_error_display(error_msg) | |
| return error_html, "{}", "" | |
| # Format results | |
| search_data = result["data"] | |
| results = search_data.get("results", []) | |
| search_time = search_data.get("search_time", 0) | |
| # Create HTML display | |
| results_html, stats_html = format_search_results(results, search_time, query) | |
| # Create JSON data for detailed view | |
| json_data = { | |
| "query": query, | |
| "search_mode": search_mode, | |
| "results_count": len(results), | |
| "search_time": search_time, | |
| "results": results[:5], # Limit JSON display | |
| "query_suggestions": search_data.get("query_suggestions", []) | |
| } | |
| return results_html, json.dumps(json_data, indent=2), stats_html | |
| except Exception as e: | |
| error_html = create_error_display(f"Search failed: {str(e)}") | |
| return error_html, "{}", "" | |
| def get_document_list(self) -> str: | |
| """Get formatted document list.""" | |
| if not self.rag_system or not self._initialized: | |
| return "<div style='text-align: center; color: #6b7280; padding: 1rem;'>System not initialized</div>" | |
| try: | |
| result = self.rag_system.get_document_list() | |
| if result.get("success"): | |
| documents = result["data"]["documents"] | |
| return format_document_list(documents) | |
| else: | |
| return create_error_display("Failed to load document list") | |
| except Exception as e: | |
| return create_error_display(f"Error loading documents: {str(e)}") | |
| def clear_documents(self) -> Tuple[str, str]: | |
| """Clear all documents.""" | |
| if not self.rag_system or not self._initialized: | |
| error_html = create_error_display("System not initialized") | |
| return error_html, error_html | |
| try: | |
| result = self.rag_system.clear_all_documents() | |
| if result.get("success"): | |
| success_msg = f"Cleared {result['data']['documents_removed']} documents" | |
| success_html = create_success_display(success_msg) | |
| return success_html, self.get_document_list() | |
| else: | |
| error_msg = result.get("error", {}).get("message", "Failed to clear documents") | |
| error_html = create_error_display(error_msg) | |
| return error_html, self.get_document_list() | |
| except Exception as e: | |
| error_html = create_error_display(f"Error clearing documents: {str(e)}") | |
| return error_html, self.get_document_list() | |
| def get_analytics_data(self) -> Tuple[str, gr.Plot, gr.Plot, List[List[str]]]: | |
| """Get analytics dashboard data.""" | |
| if not self.rag_system or not self._initialized: | |
| return ( | |
| create_error_display("System not initialized"), | |
| gr.Plot(), | |
| gr.Plot(), | |
| [] | |
| ) | |
| try: | |
| result = self.rag_system.get_analytics_dashboard() | |
| if not result.get("success"): | |
| error_html = create_error_display("Failed to load analytics data") | |
| return error_html, gr.Plot(), gr.Plot(), [] | |
| analytics_data = result["data"] | |
| # Format system overview | |
| overview_html = format_system_overview(analytics_data) | |
| # Create charts | |
| query_chart, modes_chart = create_analytics_charts(analytics_data) | |
| # Create activity table data | |
| activity_data = [] | |
| system_data = analytics_data.get("system", {}) | |
| activity_data.append([ | |
| "System Started", | |
| "System Initialization", | |
| f"Uptime: {system_data.get('uptime_hours', 0):.1f} hours", | |
| "✅ Active" | |
| ]) | |
| if system_data.get("total_queries", 0) > 0: | |
| activity_data.append([ | |
| "Recent", | |
| "Search Queries", | |
| f"{system_data.get('total_queries')} total queries", | |
| "📊 Active" | |
| ]) | |
| if system_data.get("total_documents_processed", 0) > 0: | |
| activity_data.append([ | |
| "Recent", | |
| "Document Processing", | |
| f"{system_data.get('total_documents_processed')} documents processed", | |
| "📄 Complete" | |
| ]) | |
| return overview_html, query_chart, modes_chart, activity_data | |
| except Exception as e: | |
| error_html = create_error_display(f"Error loading analytics: {str(e)}") | |
| return error_html, gr.Plot(), gr.Plot(), [] | |
| def create_interface(self) -> gr.Blocks: | |
| """Create the main Gradio interface.""" | |
| theme = get_theme() | |
| css = get_custom_css() | |
| with gr.Blocks( | |
| theme=theme, | |
| css=css, | |
| title="Professional RAG Assistant", | |
| analytics_enabled=False | |
| ) as interface: | |
| # Session state | |
| session_id_state = gr.State(value=generate_session_id()) | |
| # Header | |
| create_header() | |
| # System status | |
| system_status = create_system_status() | |
| # Main tabs | |
| with gr.Tabs() as main_tabs: | |
| # Document Upload Tab | |
| with gr.Tab("📁 Document Upload", id="upload"): | |
| gr.Markdown("Upload your documents to build the knowledge base. Supports PDF, DOCX, and TXT files.") | |
| file_upload, upload_status, upload_button = create_file_upload_section() | |
| with gr.Accordion("Upload Settings", open=False): | |
| gr.Markdown(""" | |
| **Supported formats:** PDF, DOCX, TXT | |
| **Maximum file size:** 50MB per file | |
| **Processing:** Documents are split into chunks and indexed for search | |
| """) | |
| # Search Tab | |
| with gr.Tab("🔍 Search", id="search"): | |
| gr.Markdown("Search your uploaded documents using advanced AI-powered retrieval.") | |
| with gr.Row(): | |
| with gr.Column(scale=4): | |
| search_components = create_search_interface() | |
| search_query, search_controls, search_button = search_components[:3] | |
| search_mode, num_results, enable_reranking = search_components[3:] | |
| with gr.Column(scale=1): | |
| with gr.Accordion("Advanced Options", open=False): | |
| metadata_filters = gr.Textbox( | |
| label="Metadata Filters", | |
| placeholder='{"source": "document.pdf"}', | |
| lines=3, | |
| info="JSON or key:value,key2:value2 format" | |
| ) | |
| # Results display | |
| results_html, results_json, search_stats = create_results_display() | |
| with gr.Accordion("Detailed Results (JSON)", open=False): | |
| results_json | |
| # Document Management Tab | |
| with gr.Tab("📚 Documents", id="documents"): | |
| gr.Markdown("Manage your uploaded documents and view indexing status.") | |
| document_list, refresh_docs_btn, clear_docs_btn = create_document_management() | |
| # Analytics Tab | |
| with gr.Tab("📊 Analytics", id="analytics"): | |
| gr.Markdown("View system performance and usage analytics.") | |
| analytics_components = create_analytics_dashboard() | |
| system_overview, query_chart, search_modes_chart, activity_table = analytics_components | |
| with gr.Row(): | |
| with gr.Column(): | |
| query_chart | |
| with gr.Column(): | |
| search_modes_chart | |
| with gr.Accordion("Recent Activity", open=False): | |
| activity_table | |
| refresh_analytics_btn = gr.Button("Refresh Analytics", variant="secondary") | |
| # Event handlers | |
| # File upload events | |
| file_upload.change( | |
| fn=lambda files: ( | |
| create_success_display(f"✅ {len(files)} file(s) selected! Click the green '🚀 Process Documents' button below to continue.") if files and len(files) > 0 else create_loading_display("No files selected"), | |
| gr.update(interactive=files is not None and len(files) > 0) | |
| ), | |
| inputs=[file_upload], | |
| outputs=[upload_status, upload_button], | |
| show_progress=False | |
| ) | |
| upload_button.click( | |
| fn=self.process_documents, | |
| inputs=[file_upload, session_id_state], | |
| outputs=[upload_status, system_status, search_button, document_list], | |
| show_progress=True | |
| ) | |
| # Search events | |
| search_query.change( | |
| fn=lambda query: gr.update(interactive=len(query.strip()) > 0 if query else False), | |
| inputs=[search_query], | |
| outputs=[search_button], | |
| show_progress=False | |
| ) | |
| search_button.click( | |
| fn=lambda: create_loading_display("Searching..."), | |
| inputs=[], | |
| outputs=[results_html], | |
| show_progress=False | |
| ).then( | |
| fn=self.perform_search, | |
| inputs=[ | |
| search_query, search_mode, num_results, | |
| enable_reranking, metadata_filters, session_id_state | |
| ], | |
| outputs=[results_html, results_json, search_stats], | |
| show_progress=True | |
| ) | |
| # Document management events | |
| refresh_docs_btn.click( | |
| fn=self.get_document_list, | |
| inputs=[], | |
| outputs=[document_list], | |
| show_progress=False | |
| ) | |
| clear_docs_btn.click( | |
| fn=self.clear_documents, | |
| inputs=[], | |
| outputs=[system_status, document_list], | |
| show_progress=True | |
| ) | |
| # Analytics events | |
| refresh_analytics_btn.click( | |
| fn=self.get_analytics_data, | |
| inputs=[], | |
| outputs=[system_overview, query_chart, search_modes_chart, activity_table], | |
| show_progress=True | |
| ) | |
| # Initialize interface | |
| interface.load( | |
| fn=lambda: ( | |
| self.get_system_status(), | |
| gr.update(interactive=False), # Upload button disabled initially | |
| gr.update(interactive=False), # Search button disabled initially | |
| self.get_document_list(), | |
| *self.get_analytics_data() | |
| ), | |
| inputs=[], | |
| outputs=[ | |
| system_status, upload_button, search_button, document_list, | |
| system_overview, query_chart, search_modes_chart, activity_table | |
| ], | |
| show_progress=False | |
| ) | |
| return interface | |
| def create_interface(config_path: str = None) -> gr.Blocks: | |
| """Create and return the RAG interface.""" | |
| rag_interface = RAGInterface(config_path) | |
| return rag_interface.create_interface() | |
| if __name__ == "__main__": | |
| # For testing | |
| interface = create_interface() | |
| interface.launch(debug=True) |