Fix RAG processing crashes with multiprocessing and memory optimizations

- Enhanced multiprocessing controls in vector_store.py with OMP/MKL thread limits
- Added early environment setup in app.py to prevent initialization conflicts
- Reduced batch sizes and improved memory management with garbage collection
- Added comprehensive test suite to verify connection error fixes
- Disabled worker threads and multiprocessing pools for Gradio stability

test_connection_fix.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+Test RAG connection error fix
+Tests the specific multiprocessing and connection timeout issues
+"""
+
+import os
+import tempfile
+import warnings
+
+# Set environment variables before any imports
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+os.environ['OMP_NUM_THREADS'] = '1'
+os.environ['MKL_NUM_THREADS'] = '1'
+
+# Suppress warnings for cleaner output
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def test_connection_fix():
+    """Test the connection error fix specifically"""
+    print("Testing RAG connection error fix...")
+    
+    try:
+        # Test conditional import
+        try:
+            from rag_tool import RAGTool
+            has_rag = True
+            print("✓ RAG dependencies available")
+        except ImportError:
+            print("✗ RAG dependencies not available")
+            return False
+        
+        # Create a test document
+        test_content = """This is a test document for connection error testing.
+        It contains multiple sentences to test the embedding process.
+        The document should be processed without connection errors.
+        This tests multiprocessing fixes and memory management."""
+        
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
+            f.write(test_content)
+            test_file = f.name
+        
+        try:
+            print("✓ Test document created")
+            
+            # Initialize RAG tool with environment variables already set
+            print("Initializing RAG tool with connection fixes...")
+            rag_tool = RAGTool()
+            print("✓ RAG tool initialized successfully")
+            
+            # Process document - this was causing the connection error
+            print("Processing document (this was causing connection errors)...")
+            result = rag_tool.process_uploaded_files([test_file])
+            
+            if result['success']:
+                print(f"✓ Document processed successfully: {result['message']}")
+                print(f"  - Chunks created: {result.get('index_stats', {}).get('total_chunks', 'unknown')}")
+                
+                # Test search to ensure embeddings work
+                context = rag_tool.get_relevant_context("test document", max_chunks=1)
+                print(f"✓ Search test successful, context length: {len(context)}")
+                
+                return True
+            else:
+                print(f"✗ Document processing failed: {result['message']}")
+                return False
+                
+        finally:
+            # Clean up
+            if os.path.exists(test_file):
+                os.unlink(test_file)
+                print("✓ Test file cleaned up")
+        
+    except Exception as e:
+        print(f"✗ Test failed with error: {e}")
+        return False
+
+def test_gradio_integration():
+    """Test integration with Gradio interface"""
+    print("\nTesting Gradio integration...")
+    
+    try:
+        import gradio as gr
+        
+        # Create a minimal Gradio interface similar to the main app
+        def test_process_documents(files):
+            """Minimal version of process_documents for testing"""
+            if not files:
+                return "No files uploaded"
+            
+            try:
+                from rag_tool import RAGTool
+                rag_tool = RAGTool()
+                
+                # Simulate file processing
+                file_paths = [f.name if hasattr(f, 'name') else str(f) for f in files]
+                result = rag_tool.process_uploaded_files(file_paths)
+                
+                if result['success']:
+                    return f"✓ Success: {result['message']}"
+                else:
+                    return f"✗ Failed: {result['message']}"
+                    
+            except Exception as e:
+                return f"✗ Error: {str(e)}"
+        
+        # Create interface without launching
+        with gr.Blocks() as interface:
+            file_input = gr.File(file_count="multiple", label="Test Documents")
+            output = gr.Textbox(label="Result")
+            process_btn = gr.Button("Process")
+            
+            process_btn.click(
+                test_process_documents,
+                inputs=[file_input],
+                outputs=[output]
+            )
+        
+        print("✓ Gradio interface created successfully")
+        print("  Interface can be launched without connection errors")
+        return True
+        
+    except Exception as e:
+        print(f"✗ Gradio integration test failed: {e}")
+        return False
+
+if __name__ == "__main__":
+    success = test_connection_fix()
+    if success:
+        success = test_gradio_integration()
+    
+    if success:
+        print("\n🎉 All connection error fixes are working!")
+        print("The RAG processing should now work without connection timeouts.")
+    else:
+        print("\n❌ Some tests failed. Check the error messages above.")

vector_store.py

@@ -50,6 +50,8 @@ class VectorStore:
             # Set environment variables to prevent multiprocessing issues
             import os
             os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+            os.environ['OMP_NUM_THREADS'] = '1'
+            os.environ['MKL_NUM_THREADS'] = '1'
             
             # Initialize with specific settings to avoid multiprocessing issues
             self.embedding_model = SentenceTransformer(
@@ -57,13 +59,20 @@ class VectorStore:
                 device='cpu',  # Force CPU to avoid GPU/multiprocessing conflicts
                 cache_folder=None,  # Use default cache
                 # Additional parameters to reduce memory usage
-                use_auth_token=False
+                use_auth_token=False,
+                trust_remote_code=False  # Security best practice
             )
             
             # Disable multiprocessing for stability in web apps
             if hasattr(self.embedding_model, 'pool'):
                 self.embedding_model.pool = None
             
+            # Additional stability measures for Gradio environment
+            if hasattr(self.embedding_model, '_modules'):
+                for module in self.embedding_model._modules.values():
+                    if hasattr(module, 'num_workers'):
+                        module.num_workers = 0
+            
             # Update dimension based on model
             self.dimension = self.embedding_model.get_sentence_embedding_dimension()
             print(f"Model loaded successfully, dimension: {self.dimension}")
@@ -79,7 +88,7 @@ class VectorStore:
             else:
                 raise RuntimeError(f"Could not load embedding model '{self.embedding_model_name}': {e}")
     
-    def create_embeddings(self, texts: List[str], batch_size: int = 16) -> np.ndarray:
+    def create_embeddings(self, texts: List[str], batch_size: int = 8) -> np.ndarray:
         """Create embeddings for a list of texts"""
         if not self.embedding_model:
             self._initialize_model()
@@ -99,13 +108,13 @@ class VectorStore:
                     show_progress_bar=False,
                     device='cpu',  # Force CPU to avoid GPU conflicts
                     normalize_embeddings=False,  # We'll normalize later with FAISS
-                    batch_size=batch_size  # Explicit batch size
+                    batch_size=min(batch_size, 4)  # Extra safety on batch size
                 )
                 embeddings.append(batch_embeddings)
                 
-                # Clear any caches to free memory
-                if hasattr(self.embedding_model, 'clear_cache'):
-                    self.embedding_model.clear_cache()
+                # Import gc for garbage collection
+                import gc
+                gc.collect()  # Force garbage collection between batches
                     
         except Exception as e:
             # Log the error and provide a helpful message