Fix HF Spaces initialization - improve ChromaDB connection and collection handling

app.py

@@ -67,29 +67,63 @@ class RAGChatbot:
         Initialize ChromaDB client and embedding model for retrieval.
         """
         try:
-            # Check if we're in Hugging Face Spaces environment
-            if os.path.exists('chroma.sqlite3'):
-                # We're likely in HF Spaces - use current directory
-                self.db_path = '.'
-            
-            # Initialize ChromaDB client
-            self.chroma_client = chromadb.PersistentClient(
-                path=self.db_path,
-                settings=Settings(anonymized_telemetry=False)
-            )
+            # Detect environment and set appropriate database path
+            current_dir = os.getcwd()
+            
+            # Check for database files in different locations
+            if os.path.exists(os.path.join(current_dir, 'chroma.sqlite3')):
+                self.db_path = current_dir
+                logger.info(f"Using database in current directory: {current_dir}")
+            elif os.path.exists(os.path.join(self.db_path, 'chroma.sqlite3')):
+                logger.info(f"Using database in specified path: {self.db_path}")
+            else:
+                logger.warning("No database file found, will attempt to rebuild from chunks")
             
-            # Get or create collection
+            # Initialize ChromaDB client with error handling
             try:
-                self.collection = self.chroma_client.get_collection(
-                    name=self.collection_name
+                self.chroma_client = chromadb.PersistentClient(
+                    path=self.db_path,
+                    settings=Settings(
+                        anonymized_telemetry=False,
+                        allow_reset=True,
+                        is_persistent=True
+                    )
                 )
-            except Exception:
-                # If collection doesn't exist, try to recreate it from chunks
+                logger.info(f"ChromaDB client initialized at: {self.db_path}")
+            except Exception as client_error:
+                logger.error(f"ChromaDB client initialization failed: {client_error}")
+                # Try with default settings
+                self.chroma_client = chromadb.PersistentClient(path=self.db_path)
+            
+            # Get or create collection with robust error handling
+            collection_found = False
+            try:
+                # First, list all collections to see what's available
+                collections = self.chroma_client.list_collections()
+                collection_names = [col.name for col in collections]
+                logger.info(f"Available collections: {collection_names}")
+                
+                if self.collection_name in collection_names:
+                    self.collection = self.chroma_client.get_collection(name=self.collection_name)
+                    collection_found = True
+                    logger.info(f"Successfully loaded collection: {self.collection_name}")
+                else:
+                    logger.warning(f"Collection '{self.collection_name}' not found in {collection_names}")
+                    
+            except Exception as col_error:
+                logger.error(f"Error accessing collections: {col_error}")
+            
+            # If collection not found, rebuild from chunks
+            if not collection_found:
                 if os.path.exists('chunks.json'):
-                    st.warning("Database collection not found. Rebuilding from chunks...")
+                    logger.info("Attempting to rebuild collection from chunks.json")
+                    if 'streamlit' in sys.modules:
+                        st.warning("🔄 Database collection not found. Rebuilding from chunks...")
                     self._rebuild_collection_from_chunks()
                 else:
-                    raise Exception("Neither database collection nor chunks.json found. Please build the database first.")
+                    error_msg = f"Collection '{self.collection_name}' not found and no chunks.json available for rebuilding"
+                    logger.error(error_msg)
+                    raise Exception(error_msg)
             
             # Load embedding model (same as used for building the database)
             self.embedding_model = SentenceTransformer(self.embedding_model_name)
@@ -110,18 +144,31 @@ class RAGChatbot:
         This is useful for Hugging Face Spaces deployment.
         """
         try:
-            st.info("🔄 Rebuilding database collection from chunks...")
+            logger.info("Starting collection rebuild from chunks.json")
+            if 'streamlit' in sys.modules:
+                st.info("🔄 Rebuilding database collection from chunks...")
             
-            # Load chunks
-            with open('chunks.json', 'r', encoding='utf-8') as f:
+            # Load chunks with error handling
+            chunks_path = 'chunks.json'
+            if not os.path.exists(chunks_path):
+                raise FileNotFoundError(f"chunks.json not found at {chunks_path}")
+                
+            with open(chunks_path, 'r', encoding='utf-8') as f:
                 chunks = json.load(f)
             
-            # Create collection
-            try:
-                self.chroma_client.delete_collection(name=self.collection_name)
-            except:
-                pass  # Collection might not exist
+            logger.info(f"Loaded {len(chunks)} chunks from {chunks_path}")
             
+            # Safely create collection
+            try:
+                # Try to delete existing collection first
+                existing_collections = [col.name for col in self.chroma_client.list_collections()]
+                if self.collection_name in existing_collections:
+                    logger.info(f"Deleting existing collection: {self.collection_name}")
+                    self.chroma_client.delete_collection(name=self.collection_name)
+            except Exception as del_error:
+                logger.warning(f"Could not delete existing collection: {del_error}")
+            
+            # Create new collection
             self.collection = self.chroma_client.create_collection(
                 name=self.collection_name,
                 metadata={"description": "Scikit-learn documentation embeddings"}
@@ -617,5 +664,5 @@ def main():
         """)
 
 
-if __name__ == "__main__":
-    main()
+# This ensures the app runs properly on Hugging Face Spaces
+main()