Update app.py

app.py

@@ -1,5 +1,70 @@
 import streamlit as st
 from search_utils import SemanticSearch
+import logging
+import time
+import os
+import sys
+import psutil  # Added missing import
+from urllib.parse import urlparse
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger("SemanticSearchApp")
+
+# Security validation functions
+def is_valid_url(url):
+    """Validate URL format and safety"""
+    try:
+        result = urlparse(url)
+        if not all([result.scheme, result.netloc]):
+            return False
+        # Add additional security checks here
+        return True
+    except:
+        return False
+
+def sanitize_query(query):
+    """Sanitize user input to prevent injection attacks"""
+    return re.sub(r'[^\w\s-]', '', query)[:256]
+
+# Diagnostics integration
+try:
+    from diagnostics import diagnose_parquet_files
+    diagnostics_available = True
+except ImportError:
+    diagnostics_available = False
+    logger.warning("Diagnostics module not available")
+
+def add_diagnostics_ui(search_system):
+    """Enhanced diagnostics UI with system checks"""
+    with st.sidebar.expander("🔧 Diagnostics", expanded=False):
+        if st.button("Run Full System Check"):
+            with st.spinner("Performing comprehensive system check..."):
+                col1, col2 = st.columns(2)
+                
+                # Metadata check
+                with col1:
+                    if diagnose_parquet_files("metadata_shards"):
+                        st.success("✅ Metadata shards valid")
+                    else:
+                        st.error("❌ Metadata issues detected")
+                
+                # Index check
+                with col2:
+                    if len(search_system.index_shards) > 0:
+                        st.success(f"✅ {len(search_system.index_shards)} FAISS shards loaded")
+                    else:
+                        st.error("❌ No FAISS shards found")
+                
+                # Resource check
+                st.metric("Memory Usage", f"{psutil.Process().memory_info().rss // 1024 ** 2} MB")
+                st.metric("CPU Utilization", f"{psutil.cpu_percent()}%")
 
 def main():
     st.set_page_config(
@@ -8,64 +73,162 @@ def main():
         layout="wide"
     )
 
-    # Initialize search system first
-    @st.cache_resource
+    # Initialize search system with enhanced caching
+    @st.cache_resource(ttl=3600, show_spinner="Initializing search engine...")
     def init_search_system():
-        system = SemanticSearch()
-        system.initialize_system()
-        return system
+        try:
+            system = SemanticSearch()
+            system.initialize_system()
+            logger.info("Search system initialized successfully")
+            return system
+        except Exception as e:
+            logger.error(f"System initialization failed: {str(e)}")
+            st.error("Critical system initialization error. Check logs.")
+            st.stop()
 
-    # Custom CSS moved outside cached function
+    # Custom CSS with enhanced visual design
     st.markdown("""
     <style>
     div[data-testid="stExpander"] div[role="button"] p {
         font-size: 1.2rem;
         font-weight: bold;
+        color: #1e88e5;
     }
     a.source-link {
-        color: #1e88e5 !important;
+        color: #1a73e8 !important;
         text-decoration: none !important;
-        border-bottom: 1px dotted #1e88e5;
+        border-bottom: 2px solid transparent;
+        transition: all 0.3s ease;
     }
     a.source-link:hover {
-        opacity: 0.8;
-        border-bottom-style: solid;
+        border-bottom-color: #1a73e8;
+        opacity: 0.9;
+    }
+    .similarity-badge {
+        padding: 0.2em 0.5em;
+        border-radius: 4px;
+        background: #e3f2fd;
+        color: #1e88e5;
+        font-weight: 500;
     }
     </style>
     """, unsafe_allow_html=True)
 
-    search_system = init_search_system()
+    try:
+        search_system = init_search_system()
+    except Exception as e:
+        st.error(f"Failed to initialize search system: {str(e)}")
+        st.stop()
 
     # Main UI components
     st.title("🔍 Semantic Search Engine")
-    query = st.text_input("Enter your search query:", placeholder="Search documents...")
-
+    
+    # Search input with sanitization
+    query = st.text_input("Enter your search query:", 
+                        placeholder="Search documents...",
+                        max_chars=200)
+    
     if query:
-        with st.spinner("🔍 Searching through documents..."):
-            results = search_system.search(query, 5)
-            
-            if not results.empty:
-                st.subheader("Top Results")
-                for _, row in results.iterrows():
-                    with st.expander(f"{row['title']} (Similarity: {row['similarity']:.1%})"):
-                        st.markdown(f"**Summary**: {row['summary']}")
-                        st.markdown(f"<a class='source-link' href='{row['source']}' target='_blank'>View Source</a>", 
-                                  unsafe_allow_html=True)
-            else:
-                st.warning("No matching documents found")
-
-        # System status sidebar
-        with st.sidebar:
-            st.subheader("System Status")
-            st.metric("Total Documents", f"{search_system.metadata_mgr.total_docs:,}")
-            st.metric("FAISS Shards", len(search_system.index_shards))
-            st.metric("Metadata Shards", len(search_system.metadata_mgr.shard_map))
-
-    # Sidebar controls outside main query block
+        try:
+            # Sanitize and validate query
+            clean_query = sanitize_query(query)
+            if not clean_query:
+                st.warning("Please enter a valid search query")
+                st.stop()
+                
+            with st.spinner("🔍 Searching through documents..."):
+                start_time = time.time()
+                results = search_system.search(clean_query, 5)
+                search_duration = time.time() - start_time
+                
+                if not results.empty:
+                    st.subheader(f"Top Results ({search_duration:.2f}s)")
+                    
+                    # Visualize results with enhanced formatting
+                    for _, row in results.iterrows():
+                        with st.expander(f"{row['title']}"):
+                            # Similarity visualization
+                            col1, col2 = st.columns([3, 1])
+                            with col1:
+                                st.markdown(f"**Summary**: {row['summary']}")
+                            with col2:
+                                st.markdown(
+                                    f"<div class='similarity-badge'>"
+                                    f"Confidence: {row['similarity']:.1%}"
+                                    f"</div>", 
+                                    unsafe_allow_html=True
+                                )
+                                st.progress(float(row['similarity']))
+                            
+                            # Safe URL handling
+                            if is_valid_url(row['source']):
+                                st.markdown(
+                                    f"<a class='source-link' href='{row['source']}' "
+                                    f"target='_blank' rel='noopener noreferrer'>"
+                                    f"🌐 View Source</a>", 
+                                    unsafe_allow_html=True
+                                )
+                            else:
+                                st.warning("Invalid source URL")
+                else:
+                    st.warning("No matching documents found")
+                    st.info("Try these tips:")
+                    st.markdown("""
+                    - Use more specific keywords
+                    - Check your spelling
+                    - Avoid special characters
+                    """)
+
+        except Exception as e:
+            logger.error(f"Search failed: {str(e)}")
+            st.error("Search operation failed. Please try again.")
+
+    # System monitoring sidebar
     with st.sidebar:
-        if st.button("Clear Cache"):
-            st.cache_resource.clear()
-            st.rerun()
+        st.subheader("📊 System Status")
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.metric("Total Documents", 
+                    f"{search_system.metadata_mgr.total_docs:,}",
+                    help="Total indexed documents in system")
+            
+        with col2:
+            st.metric("FAISS Shards", 
+                    len(search_system.index_shards),
+                    help="Number of loaded vector index shards")
+        
+        st.metric("Active Memory", 
+                f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
+                help="Current memory usage by the application")
+
+        # Diagnostics section
+        if diagnostics_available:
+            add_diagnostics_ui(search_system)
+        else:
+            st.warning("Diagnostics module not available")
+
+        # Health check with error handling
+        if st.button("🩺 Run Health Check"):
+            try:
+                system_stats = {
+                    "shards_loaded": len(search_system.index_shards),
+                    "metadata_records": search_system.metadata_mgr.total_docs,
+                    "memory_usage": f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
+                    "active_threads": threading.active_count(),
+                    "system_load": f"{os.getloadavg()[0]:.2f}"
+                }
+                st.json(system_stats)
+            except Exception as e:
+                st.error(f"Health check failed: {str(e)}")
+
+        # Cache management
+        if st.button("♻️ Clear Cache"):
+            try:
+                st.cache_resource.clear()
+                st.rerun()
+            except Exception as e:
+                st.error(f"Cache clearance failed: {str(e)}")
 
 if __name__ == "__main__":
     main()