Add data clearing service and vector store management

- Introduced `DataClearingService` for managing the clearing of vector store and chat history data.
- Implemented `clear_all_documents` method in `VectorStoreManager` to remove all documents from the vector store.
- Enhanced UI with a button to clear all data, integrating the new service for user interaction.
- Updated logging to provide detailed feedback during data clearing operations.

src/rag/vector_store.py

@@ -225,6 +225,37 @@ class VectorStoreManager:
         except Exception as e:
             logger.error(f"Error searching with metadata filter: {e}")
             return []
+    
+    def clear_all_documents(self) -> bool:
+        """
+        Clear all documents from the vector store collection.
+        
+        Returns:
+            True if successful, False otherwise
+        """
+        try:
+            vector_store = self.get_vector_store()
+            
+            # Get all document IDs first
+            collection = vector_store._collection
+            all_docs = collection.get()
+            
+            if not all_docs or not all_docs.get('ids'):
+                logger.info("No documents found in vector store to clear")
+                return True
+            
+            # Delete all documents using their IDs
+            collection.delete(ids=all_docs['ids'])
+            
+            # Reset the vector store instance to ensure clean state
+            self._vector_store = None
+            
+            logger.info(f"Successfully cleared {len(all_docs['ids'])} documents from vector store")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error clearing all documents: {e}")
+            return False
 
 # Global vector store manager instance
 vector_store_manager = VectorStoreManager()

src/services/data_clearing_service.py

@@ -0,0 +1,263 @@
+"""
+Data clearing service for both local and Hugging Face Space environments.
+Provides functionality to clear vector store and chat history data.
+"""
+
+import os
+import shutil
+from pathlib import Path
+from typing import Dict, Any, Tuple, List
+from src.core.config import config
+from src.core.logging_config import get_logger
+from src.rag.vector_store import vector_store_manager
+
+logger = get_logger(__name__)
+
+
+class DataClearingService:
+    """Service for clearing all RAG-related data across different environments."""
+    
+    def __init__(self):
+        """Initialize the data clearing service."""
+        self.is_hf_space = bool(os.getenv("SPACE_ID"))
+        logger.info(f"DataClearingService initialized (HF Space: {self.is_hf_space})")
+    
+    def get_data_paths(self) -> Tuple[str, str]:
+        """
+        Get the correct data paths for current environment.
+        
+        Returns:
+            Tuple of (vector_store_path, chat_history_path)
+        """
+        vector_store_path = config.rag.vector_store_path
+        chat_history_path = config.rag.chat_history_path
+        
+        logger.info(f"Data paths - Vector store: {vector_store_path}, Chat history: {chat_history_path}")
+        return vector_store_path, chat_history_path
+    
+    def clear_vector_store(self) -> Tuple[bool, str, Dict[str, Any]]:
+        """
+        Clear all documents from the vector store.
+        
+        Returns:
+            Tuple of (success, message, stats)
+        """
+        try:
+            # Get initial document count
+            collection_info = vector_store_manager.get_collection_info()
+            initial_count = collection_info.get("document_count", 0)
+            
+            if initial_count == 0:
+                return True, "Vector store is already empty", {"cleared_documents": 0}
+            
+            # Clear the collection using the vector store manager's method
+            success = vector_store_manager.clear_all_documents()
+            
+            if not success:
+                return False, "Failed to clear vector store", {"error": "clear_all_documents returned False"}
+            
+            logger.info(f"Cleared {initial_count} documents from vector store")
+            
+            return True, f"Successfully cleared {initial_count} documents from vector store", {
+                "cleared_documents": initial_count,
+                "collection_name": collection_info.get("collection_name", "unknown")
+            }
+            
+        except Exception as e:
+            error_msg = f"Error clearing vector store: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg, {"error": str(e)}
+    
+    def clear_chat_history(self) -> Tuple[bool, str, Dict[str, Any]]:
+        """
+        Clear all chat history files.
+        
+        Returns:
+            Tuple of (success, message, stats)
+        """
+        try:
+            _, chat_history_path = self.get_data_paths()
+            chat_dir = Path(chat_history_path)
+            
+            if not chat_dir.exists():
+                return True, "Chat history directory doesn't exist", {"cleared_files": 0}
+            
+            # Count files before deletion
+            files_to_clear = list(chat_dir.rglob("*"))
+            file_count = len([f for f in files_to_clear if f.is_file()])
+            
+            if file_count == 0:
+                return True, "Chat history is already empty", {"cleared_files": 0}
+            
+            # Clear all contents of the chat history directory
+            for item in chat_dir.iterdir():
+                if item.is_file():
+                    item.unlink()
+                    logger.debug(f"Removed file: {item}")
+                elif item.is_dir():
+                    shutil.rmtree(item)
+                    logger.debug(f"Removed directory: {item}")
+            
+            logger.info(f"Cleared {file_count} files from chat history")
+            
+            return True, f"Successfully cleared {file_count} files from chat history", {
+                "cleared_files": file_count,
+                "chat_history_path": str(chat_dir)
+            }
+            
+        except Exception as e:
+            error_msg = f"Error clearing chat history: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg, {"error": str(e)}
+    
+    def clear_directory_contents(self, directory_path: str) -> Tuple[bool, str, int]:
+        """
+        Clear all contents of a specific directory.
+        
+        Args:
+            directory_path: Path to directory to clear
+            
+        Returns:
+            Tuple of (success, message, items_cleared)
+        """
+        try:
+            dir_path = Path(directory_path)
+            
+            if not dir_path.exists():
+                return True, f"Directory doesn't exist: {directory_path}", 0
+            
+            items_cleared = 0
+            for item in dir_path.iterdir():
+                try:
+                    if item.is_file():
+                        item.unlink()
+                        items_cleared += 1
+                        logger.debug(f"Removed file: {item}")
+                    elif item.is_dir():
+                        shutil.rmtree(item)
+                        items_cleared += 1
+                        logger.debug(f"Removed directory: {item}")
+                except Exception as e:
+                    logger.warning(f"Failed to remove {item}: {e}")
+            
+            return True, f"Cleared {items_cleared} items from {directory_path}", items_cleared
+            
+        except Exception as e:
+            error_msg = f"Error clearing directory {directory_path}: {str(e)}"
+            logger.error(error_msg)
+            return False, error_msg, 0
+    
+    def clear_all_data(self) -> Tuple[bool, str, Dict[str, Any]]:
+        """
+        Clear all RAG-related data (vector store + chat history).
+        
+        Returns:
+            Tuple of (success, message, combined_stats)
+        """
+        logger.info("Starting complete data clearing operation")
+        
+        combined_stats = {
+            "vector_store": {},
+            "chat_history": {},
+            "total_cleared_documents": 0,
+            "total_cleared_files": 0,
+            "environment": "hf_space" if self.is_hf_space else "local",
+            "errors": []
+        }
+        
+        # Clear vector store
+        vs_success, vs_message, vs_stats = self.clear_vector_store()
+        combined_stats["vector_store"] = {
+            "success": vs_success,
+            "message": vs_message,
+            **vs_stats
+        }
+        
+        if not vs_success:
+            combined_stats["errors"].append(f"Vector store: {vs_message}")
+        else:
+            combined_stats["total_cleared_documents"] = vs_stats.get("cleared_documents", 0)
+        
+        # Clear chat history
+        ch_success, ch_message, ch_stats = self.clear_chat_history()
+        combined_stats["chat_history"] = {
+            "success": ch_success,
+            "message": ch_message,
+            **ch_stats
+        }
+        
+        if not ch_success:
+            combined_stats["errors"].append(f"Chat history: {ch_message}")
+        else:
+            combined_stats["total_cleared_files"] = ch_stats.get("cleared_files", 0)
+        
+        # Overall success
+        overall_success = vs_success and ch_success
+        
+        if overall_success:
+            total_items = combined_stats["total_cleared_documents"] + combined_stats["total_cleared_files"]
+            if total_items == 0:
+                overall_message = "All data was already clear"
+            else:
+                overall_message = f"Successfully cleared all data: {combined_stats['total_cleared_documents']} documents, {combined_stats['total_cleared_files']} files"
+        else:
+            overall_message = f"Data clearing completed with errors: {'; '.join(combined_stats['errors'])}"
+        
+        logger.info(f"Data clearing operation completed: {overall_message}")
+        
+        return overall_success, overall_message, combined_stats
+    
+    def get_data_status(self) -> Dict[str, Any]:
+        """
+        Get current status of data directories and vector store.
+        
+        Returns:
+            Dictionary with data status information
+        """
+        try:
+            vector_store_path, chat_history_path = self.get_data_paths()
+            
+            # Vector store status
+            collection_info = vector_store_manager.get_collection_info()
+            vs_document_count = collection_info.get("document_count", 0)
+            
+            # Chat history status
+            chat_dir = Path(chat_history_path)
+            ch_file_count = 0
+            if chat_dir.exists():
+                ch_file_count = len([f for f in chat_dir.rglob("*") if f.is_file()])
+            
+            # Directory status
+            vs_dir = Path(vector_store_path)
+            vs_exists = vs_dir.exists()
+            ch_exists = chat_dir.exists()
+            
+            status = {
+                "environment": "hf_space" if self.is_hf_space else "local",
+                "vector_store": {
+                    "path": vector_store_path,
+                    "exists": vs_exists,
+                    "document_count": vs_document_count,
+                    "collection_name": collection_info.get("collection_name", "unknown")
+                },
+                "chat_history": {
+                    "path": chat_history_path,
+                    "exists": ch_exists,
+                    "file_count": ch_file_count
+                },
+                "total_data_items": vs_document_count + ch_file_count,
+                "has_data": vs_document_count > 0 or ch_file_count > 0
+            }
+            
+            return status
+            
+        except Exception as e:
+            logger.error(f"Error getting data status: {e}")
+            return {
+                "error": str(e),
+                "environment": "hf_space" if self.is_hf_space else "local"
+            }
+
+
+# Global data clearing service instance
+data_clearing_service = DataClearingService()

src/ui/ui.py

@@ -15,6 +15,7 @@ from src.core.exceptions import (
 )
 from src.core.logging_config import get_logger
 from src.rag import rag_chat_service, document_ingestion_service
+from src.services.data_clearing_service import data_clearing_service
 
 # Use centralized logging
 logger = get_logger(__name__)
@@ -252,6 +253,48 @@ def start_new_chat_session():
         logger.error(error_msg)
         return [], f"❌ {error_msg}"
 
+def handle_clear_all_data():
+    """Handle clearing all RAG data (vector store + chat history)."""
+    try:
+        # Clear all data using the data clearing service
+        success, message, stats = data_clearing_service.clear_all_data()
+        
+        if success:
+            # Reset chat session after clearing data
+            session_id = rag_chat_service.start_new_session()
+            
+            # Get updated status
+            updated_status = get_chat_status()
+            
+            # Create success message with stats
+            if stats.get("total_cleared_documents", 0) > 0 or stats.get("total_cleared_files", 0) > 0:
+                clear_msg = f"✅ {message}"
+                session_msg = f"🆕 Started new session: {session_id}"
+                combined_msg = f'{clear_msg}<br/><div class="session-info">{session_msg}</div>'
+            else:
+                combined_msg = f'ℹ️ {message}<br/><div class="session-info">🆕 Started new session: {session_id}</div>'
+            
+            logger.info(f"Data cleared successfully: {message}")
+            
+            return [], combined_msg, updated_status
+        else:
+            error_msg = f"❌ {message}"
+            logger.error(f"Data clearing failed: {message}")
+            
+            # Still get updated status even on error
+            updated_status = get_chat_status()
+            
+            return None, f'<div class="session-info">{error_msg}</div>', updated_status
+            
+    except Exception as e:
+        error_msg = f"Error clearing data: {str(e)}"
+        logger.error(error_msg)
+        
+        # Get current status
+        current_status = get_chat_status()
+        
+        return None, f'<div class="session-info">❌ {error_msg}</div>', current_status
+
 def get_chat_status():
     """Get current chat system status."""
     try:
@@ -261,6 +304,9 @@ def get_chat_status():
         # Check usage stats
         usage_stats = rag_chat_service.get_usage_stats()
         
+        # Get data status for additional context
+        data_status = data_clearing_service.get_data_status()
+        
         # Modern status card design with better styling
         status_html = f"""
         <div class="status-card">
@@ -273,20 +319,20 @@ def get_chat_status():
             
             <div class="status-grid">
                 <div class="status-item">
-                    <div class="status-label">Documents Processed</div>
-                    <div class="status-value">{ingestion_status.get('processed_documents', 0)}</div>
+                    <div class="status-label">Vector Store Docs</div>
+                    <div class="status-value">{data_status.get('vector_store', {}).get('document_count', 0)}</div>
                 </div>
                 <div class="status-item">
-                    <div class="status-label">Vector Store</div>
-                    <div class="status-value">{ingestion_status.get('total_documents_in_store', 0)} docs</div>
+                    <div class="status-label">Chat History Files</div>
+                    <div class="status-value">{data_status.get('chat_history', {}).get('file_count', 0)}</div>
                 </div>
                 <div class="status-item">
                     <div class="status-label">Session Usage</div>
                     <div class="status-value">{usage_stats.get('session_messages', 0)}/{usage_stats.get('session_limit', 50)}</div>
                 </div>
                 <div class="status-item">
-                    <div class="status-label">Hourly Usage</div>
-                    <div class="status-value">{usage_stats.get('hourly_messages', 0)}/{usage_stats.get('hourly_limit', 100)}</div>
+                    <div class="status-label">Environment</div>
+                    <div class="status-value">{'HF Space' if data_status.get('environment') == 'hf_space' else 'Local'}</div>
                 </div>
             </div>
             
@@ -556,6 +602,16 @@ def create_ui():
             transform: translateY(-1px);
         }
         
+        .btn-clear-data {
+            background: #dc3545;
+            color: white;
+        }
+        
+        .btn-clear-data:hover {
+            background: #c82333;
+            transform: translateY(-1px);
+        }
+        
         /* Chat interface styling */
         .chat-main-container {
             background: #ffffff;
@@ -819,6 +875,7 @@ def create_ui():
                     with gr.Row(elem_classes=["control-buttons"]):
                         refresh_status_btn = gr.Button("🔄 Refresh Status", elem_classes=["control-btn", "btn-refresh"])
                         new_session_btn = gr.Button("🆕 New Session", elem_classes=["control-btn", "btn-new-session"])
+                        clear_data_btn = gr.Button("🗑️ Clear All Data", elem_classes=["control-btn", "btn-clear-data"], variant="stop")
                     
                     # Main chat interface
                     with gr.Column(elem_classes=["chat-main-container"]):
@@ -885,6 +942,13 @@ def create_ui():
                     inputs=[],
                     outputs=[status_display]
                 )
+                
+                # Clear all data handler
+                clear_data_btn.click(
+                    fn=handle_clear_all_data,
+                    inputs=[],
+                    outputs=[chatbot, session_info, status_display]
+                )
 
     return demo