Update app.py

app.py

@@ -1,46 +1,40 @@
 import streamlit as st
-from search_utils import SemanticSearch
+from search_utils import OptimizedSemanticSearch
 import logging
 import time
 import os
 import sys
-import psutil  # Added missing import
-from urllib.parse import urlparse
+import psutil
+from urllib.parse import urlparse, quote
 import threading
 import re
 
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.StreamHandler()
-    ]
+    handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger("SemanticSearchApp")
 
 # Security validation functions
-def is_valid_url(url):
-    """Validate URL format and safety"""
+def is_valid_url(url: str) -> bool:
+    """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 all([result.scheme, result.netloc])
+    except Exception:
         return False
 
-def sanitize_query(query):
-    """Sanitize user input to prevent injection attacks"""
+def sanitize_query(query: str) -> str:
+    """Sanitize user input to prevent injection attacks."""
     try:
-        # Remove non-alphanumeric characters except spaces and hyphens
+        # Remove non-alphanumeric characters except spaces and hyphens.
         clean_query = re.sub(r'[^\w\s-]', '', query)
-        return clean_query[:256]  # Truncate to prevent long queries
+        return clean_query[:256]  # Truncate to prevent overly long queries.
     except Exception as e:
         logger.error(f"Query sanitization failed: {str(e)}")
-        return query[:256]  # Fallback truncation
+        return query[:256]  # Fallback truncation.
 
 # Diagnostics integration
 try:
@@ -50,29 +44,24 @@ except ImportError:
     diagnostics_available = False
     logger.warning("Diagnostics module not available")
 
-def add_diagnostics_ui(search_system):
-    """Enhanced diagnostics UI with proper directory checks"""
+def add_diagnostics_ui(search_system: OptimizedSemanticSearch):
+    """Enhanced diagnostics UI with proper directory checks."""
     with st.sidebar.expander("🔧 Diagnostics", expanded=False):
         if st.button("Run Full System Check"):
             with st.spinner("Performing comprehensive system check..."):
-                # Create columns for organized display
                 col1, col2 = st.columns(2)
                 
-                # Get actual paths from the search system
-                metadata_dir = search_system.metadata_mgr.shard_dir
-                faiss_dir = search_system.shard_dir  # From SemanticSearch class
+                # Get actual paths from the search system.
+                metadata_dir = search_system.metadata_mgr.metadata_dir
+                faiss_dir = Path("compressed_shards")  # Assuming FAISS index dir, if applicable.
 
                 with col1:
-                    # Metadata directory check
                     st.subheader("📂 Metadata Validation")
                     if metadata_dir.exists():
-                        # Check directory structure
-                        dir_status = any(metadata_dir.glob("*.parquet"))
+                        parquet_files = list(metadata_dir.glob("*.parquet"))
                         st.write(f"Directory: `{metadata_dir}`")
-                        st.write(f"Parquet Files Found: {'✅' if dir_status else '❌'}")
-                        
-                        # Check individual files
-                        if diagnose_parquet_files(str(metadata_dir)):
+                        st.write(f"Parquet Files Found: {'✅' if parquet_files else '❌'}")
+                        if diagnostics_available and diagnose_parquet_files(str(metadata_dir)):
                             st.success("✅ Metadata shards valid")
                         else:
                             st.error("❌ Metadata issues detected")
@@ -80,33 +69,24 @@ def add_diagnostics_ui(search_system):
                         st.error("Metadata directory not found")
 
                 with col2:
-                    # FAISS index check
                     st.subheader("📚 FAISS Validation")
-                    if faiss_dir.exists():
-                        index_files = list(faiss_dir.glob("*.index"))
-                        st.write(f"Directory: `{faiss_dir}`")
-                        st.write(f"Index Files Found: {len(index_files)}")
-                        
-                        if len(search_system.index_shards) > 0:
-                            st.success(f"✅ {len(search_system.index_shards)} FAISS shards loaded")
-                            st.write(f"Total Vectors: {sum(s.ntotal for s in search_system.index_shards):,}")
-                        else:
-                            st.error("❌ No FAISS shards loaded")
+                    # For the FAISS index, we rely on our search system.
+                    if hasattr(search_system, 'index') and search_system.index is not None:
+                        st.success(f"✅ FAISS index loaded with {search_system.index.ntotal} vectors")
                     else:
-                        st.error("FAISS directory not found")
+                        st.error("❌ No FAISS index loaded")
 
-                # System resource check
                 st.subheader("💻 System Resources")
                 col_res1, col_res2 = st.columns(2)
                 with col_res1:
                     st.metric("Memory Usage", 
-                            f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
-                            help="Current process memory usage")
-                
+                              f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
+                              help="Current process memory usage")
                 with col_res2:
                     st.metric("CPU Utilization", 
-                            f"{psutil.cpu_percent()}%",
-                            help="Total system CPU usage")
+                              f"{psutil.cpu_percent()}%",
+                              help="Total system CPU usage")
+
 def main():
     st.set_page_config(
         page_title="Semantic Search Engine",
@@ -114,12 +94,11 @@ def main():
         layout="wide"
     )
 
-    # Initialize search system with enhanced caching
+    # Initialize the optimized search system with caching.
     @st.cache_resource(ttl=3600, show_spinner="Initializing search engine...")
     def init_search_system():
         try:
-            system = SemanticSearch()
-            system.initialize_system()
+            system = OptimizedSemanticSearch()
             logger.info("Search system initialized successfully")
             return system
         except Exception as e:
@@ -127,7 +106,7 @@ def main():
             st.error("Critical system initialization error. Check logs.")
             st.stop()
 
-    # Custom CSS with enhanced visual design
+    # Custom CSS for enhanced visual design.
     st.markdown("""
     <style>
     div[data-testid="stExpander"] div[role="button"] p {
@@ -161,17 +140,15 @@ def main():
         st.error(f"Failed to initialize search system: {str(e)}")
         st.stop()
 
-    # Main UI components
     st.title("🔍 Semantic Search Engine")
     
-    # Search input with sanitization
+    # Search input with sanitization.
     query = st.text_input("Enter your search query:", 
-                        placeholder="Search documents...",
-                        max_chars=200)
+                          placeholder="Search documents...",
+                          max_chars=200)
     
     if query:
         try:
-            # Sanitize and validate query
             clean_query = sanitize_query(query)
             if not clean_query:
                 st.warning("Please enter a valid search query")
@@ -179,29 +156,29 @@ def main():
                 
             with st.spinner("🔍 Searching through documents..."):
                 start_time = time.time()
-                results = search_system.search(clean_query, 5)
+                results = search_system.search(clean_query, top_k=5)
                 search_duration = time.time() - start_time
                 
-                if not results.empty:
-                    st.subheader(f"Top Results ({search_duration:.2f}s)")
+                if results:
+                    st.subheader(f"Top Results (Search completed in {search_duration:.2f}s)")
                     
-                    # Visualize results with enhanced formatting
-                    for _, row in results.iterrows():
-                        with st.expander(f"{row['title']}"):
-                            # Similarity visualization
+                    for res in results:
+                        with st.expander(f"{res.get('title', 'Untitled')}"):
                             col1, col2 = st.columns([3, 1])
                             with col1:
-                                st.markdown(f"**Summary**: {row['summary']}")
+                                st.markdown(f"**Summary:** {res.get('summary', '')}")
                             with col2:
                                 st.markdown(
                                     f"<div class='similarity-badge'>"
-                                    f"Confidence: {row['similarity']:.1%}"
+                                    f"Confidence: {res.get('similarity', 0):.1%}"
                                     f"</div>", 
                                     unsafe_allow_html=True
                                 )
-                                st.progress(float(row['similarity']))
-                                if row['source']:
-                                    st.markdown(row['source'], unsafe_allow_html=True)                            
+                                st.progress(float(res.get("similarity", 0)))
+                            if res.get("source"):
+                                # Validate URL before displaying.
+                                if is_valid_url(res["source"]):
+                                    st.markdown(f"<a class='source-link' href='{res['source']}' target='_blank'>Source Link</a>", unsafe_allow_html=True)
                                 else:
                                     st.warning("Invalid source URL")
                 else:
@@ -212,41 +189,36 @@ def main():
                     - 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
+    # Sidebar: System monitoring and diagnostics.
     with st.sidebar:
         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")
-            
+                      f"{search_system.metadata_mgr.total_docs:,}",
+                      help="Total indexed documents in the system")
         with col2:
-            st.metric("FAISS Shards", 
-                    len(search_system.index_shards),
-                    help="Number of loaded vector index shards")
-        
+            if hasattr(search_system, 'index'):
+                st.metric("FAISS Vectors", 
+                          f"{search_system.index.ntotal:,}",
+                          help="Number of vectors in the FAISS index")
         st.metric("Active Memory", 
-                f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
-                help="Current memory usage by the application")
+                  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),
+                    "shards_loaded": 1 if hasattr(search_system, 'index') else 0,
                     "metadata_records": search_system.metadata_mgr.total_docs,
                     "memory_usage": f"{psutil.Process().memory_info().rss // 1024 ** 2} MB",
                     "active_threads": threading.active_count(),
@@ -256,13 +228,12 @@ def main():
             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()
+                st.experimental_rerun()
             except Exception as e:
                 st.error(f"Cache clearance failed: {str(e)}")
 
 if __name__ == "__main__":
-    main()
+    main()