Spaces:
Running
Running
File size: 5,029 Bytes
a1588ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 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.") |