431 lines
16 KiB
Python
431 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
OCR PDF Upload and Search Validation Script with Authentication
|
|
Handles login and authentication for LightRAG web UI
|
|
"""
|
|
|
|
import requests
|
|
import json
|
|
import time
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# Configuration
|
|
BASE_URL = "http://localhost:3015"
|
|
WEBUI_URL = f"{BASE_URL}/webui/"
|
|
API_URL = f"{BASE_URL}/api"
|
|
OCR_PDF_PATH = "ocr.pdf"
|
|
TEST_QUERY = "document processing"
|
|
MAX_WAIT_TIME = 300
|
|
POLL_INTERVAL = 10
|
|
|
|
# Authentication
|
|
USERNAME = "jleu3482"
|
|
PASSWORD = "jleu1212"
|
|
|
|
class OCRWorkflowValidator:
|
|
def __init__(self):
|
|
self.session = requests.Session()
|
|
self.doc_id = None
|
|
self.auth_token = None
|
|
|
|
def log_step(self, message, status="INFO"):
|
|
"""Log step with timestamp"""
|
|
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
print(f"[{timestamp}] [{status}] {message}")
|
|
|
|
def authenticate(self):
|
|
"""Authenticate with the LightRAG server"""
|
|
self.log_step("Authenticating with LightRAG server...")
|
|
|
|
# Try multiple authentication methods
|
|
auth_methods = [
|
|
self._authenticate_via_login_api,
|
|
self._authenticate_via_basic_auth,
|
|
self._authenticate_via_webui_login
|
|
]
|
|
|
|
for auth_method in auth_methods:
|
|
try:
|
|
if auth_method():
|
|
self.log_step("✓ Authentication successful")
|
|
return True
|
|
except Exception as e:
|
|
self.log_step(f"Authentication method failed: {e}", "WARNING")
|
|
continue
|
|
|
|
self.log_step("✗ All authentication methods failed", "ERROR")
|
|
return False
|
|
|
|
def _authenticate_via_login_api(self):
|
|
"""Authenticate via /auth/login endpoint"""
|
|
login_payload = {
|
|
"username": USERNAME,
|
|
"password": PASSWORD
|
|
}
|
|
|
|
login_endpoints = [
|
|
f"{API_URL}/auth/login",
|
|
f"{BASE_URL}/auth/login",
|
|
f"{WEBUI_URL}auth/login"
|
|
]
|
|
|
|
for endpoint in login_endpoints:
|
|
try:
|
|
self.log_step(f"Trying login endpoint: {endpoint}")
|
|
response = self.session.post(
|
|
endpoint,
|
|
json=login_payload,
|
|
timeout=10
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
result = response.json()
|
|
# Store token if provided
|
|
if 'access_token' in result:
|
|
self.auth_token = result['access_token']
|
|
self.session.headers.update({'Authorization': f'Bearer {self.auth_token}'})
|
|
elif 'token' in result:
|
|
self.auth_token = result['token']
|
|
self.session.headers.update({'Authorization': f'Bearer {self.auth_token}'})
|
|
return True
|
|
except requests.exceptions.RequestException:
|
|
continue
|
|
|
|
return False
|
|
|
|
def _authenticate_via_basic_auth(self):
|
|
"""Use basic authentication"""
|
|
from requests.auth import HTTPBasicAuth
|
|
self.session.auth = HTTPBasicAuth(USERNAME, PASSWORD)
|
|
|
|
# Test if basic auth works
|
|
try:
|
|
response = self.session.get(f"{BASE_URL}/", timeout=5)
|
|
if response.status_code == 200:
|
|
return True
|
|
except:
|
|
pass
|
|
|
|
self.session.auth = None
|
|
return False
|
|
|
|
def _authenticate_via_webui_login(self):
|
|
"""Try web UI login form"""
|
|
# Get CSRF token if needed
|
|
try:
|
|
login_page = self.session.get(f"{WEBUI_URL}login", timeout=5)
|
|
if login_page.status_code == 200:
|
|
# For now, just set basic auth for webui endpoints
|
|
self.session.auth = HTTPBasicAuth(USERNAME, PASSWORD)
|
|
return True
|
|
except:
|
|
pass
|
|
|
|
return False
|
|
|
|
def check_server_status(self):
|
|
"""Check if server is responding"""
|
|
self.log_step("Checking server status...")
|
|
|
|
endpoints_to_try = [
|
|
"/",
|
|
"/health",
|
|
"/api/health"
|
|
]
|
|
|
|
for endpoint in endpoints_to_try:
|
|
try:
|
|
response = self.session.get(f"{BASE_URL}{endpoint}", timeout=5)
|
|
self.log_step(f"✓ Endpoint {endpoint}: {response.status_code}")
|
|
if response.status_code == 200:
|
|
return True
|
|
except requests.exceptions.RequestException as e:
|
|
self.log_step(f"✗ Endpoint {endpoint}: {e}", "WARNING")
|
|
continue
|
|
|
|
self.log_step("✗ No working endpoints found", "ERROR")
|
|
return False
|
|
|
|
def verify_ocr_pdf_exists(self):
|
|
"""Verify the OCR PDF file exists and is valid"""
|
|
self.log_step("Verifying OCR PDF file...")
|
|
if not os.path.exists(OCR_PDF_PATH):
|
|
self.log_step(f"✗ OCR PDF file not found: {OCR_PDF_PATH}", "ERROR")
|
|
return False
|
|
|
|
file_size = os.path.getsize(OCR_PDF_PATH)
|
|
if file_size == 0:
|
|
self.log_step("✗ OCR PDF file is empty", "ERROR")
|
|
return False
|
|
|
|
self.log_step(f"✓ OCR PDF file verified ({file_size} bytes)")
|
|
return True
|
|
|
|
def upload_ocr_pdf(self):
|
|
"""Upload OCR PDF with authentication"""
|
|
self.log_step("Uploading OCR PDF file...")
|
|
|
|
if not self.verify_ocr_pdf_exists():
|
|
return False
|
|
|
|
try:
|
|
with open(OCR_PDF_PATH, 'rb') as file:
|
|
files = {'file': (os.path.basename(OCR_PDF_PATH), file, 'application/pdf')}
|
|
|
|
# Try multiple upload endpoints
|
|
upload_endpoints = [
|
|
f"{BASE_URL}/documents/upload",
|
|
f"{API_URL}/documents/upload",
|
|
f"{BASE_URL}/upload",
|
|
f"{API_URL}/upload"
|
|
]
|
|
|
|
for endpoint in upload_endpoints:
|
|
try:
|
|
self.log_step(f"Trying upload endpoint: {endpoint}")
|
|
response = self.session.post(
|
|
endpoint,
|
|
files=files,
|
|
timeout=30
|
|
)
|
|
|
|
if response.status_code in [200, 201]:
|
|
result = response.json()
|
|
self.log_step("✓ OCR PDF upload successful")
|
|
|
|
# Extract document ID from response
|
|
if 'document_id' in result:
|
|
self.doc_id = result['document_id']
|
|
elif 'id' in result:
|
|
self.doc_id = result['id']
|
|
elif 'doc_id' in result:
|
|
self.doc_id = result['doc_id']
|
|
|
|
if self.doc_id:
|
|
self.log_step(f"Document ID: {self.doc_id}")
|
|
return True
|
|
else:
|
|
self.log_step(f"Upload endpoint {endpoint} returned {response.status_code}: {response.text}", "WARNING")
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
self.log_step(f"Upload endpoint {endpoint} failed: {e}", "WARNING")
|
|
continue
|
|
|
|
self.log_step("✗ All upload endpoints failed", "ERROR")
|
|
return False
|
|
|
|
except Exception as e:
|
|
self.log_step(f"✗ Upload failed: {e}", "ERROR")
|
|
return False
|
|
|
|
def wait_for_indexing(self):
|
|
"""Wait for document to be fully indexed"""
|
|
self.log_step("Waiting for document indexing to complete...")
|
|
|
|
if not self.doc_id:
|
|
self.log_step("✗ No document ID available to check indexing status", "ERROR")
|
|
return False
|
|
|
|
start_time = time.time()
|
|
|
|
while time.time() - start_time < MAX_WAIT_TIME:
|
|
try:
|
|
# Try to get document status
|
|
status_endpoints = [
|
|
f"{BASE_URL}/documents/{self.doc_id}/status",
|
|
f"{API_URL}/documents/{self.doc_id}/status",
|
|
f"{BASE_URL}/status/{self.doc_id}",
|
|
f"{API_URL}/status/{self.doc_id}"
|
|
]
|
|
|
|
for endpoint in status_endpoints:
|
|
try:
|
|
response = self.session.get(endpoint, timeout=10)
|
|
if response.status_code == 200:
|
|
status_data = response.json()
|
|
|
|
# Check various status indicators
|
|
if 'status' in status_data:
|
|
status = status_data['status'].lower()
|
|
if status in ['completed', 'done', 'indexed']:
|
|
self.log_step("✓ Document indexing completed")
|
|
return True
|
|
elif status in ['processing', 'indexing']:
|
|
self.log_step(f"Indexing in progress... ({status})")
|
|
elif status in ['failed', 'error']:
|
|
self.log_step(f"✗ Indexing failed: {status_data.get('message', 'Unknown error')}", "ERROR")
|
|
return False
|
|
|
|
# Alternative status check
|
|
if 'indexed' in status_data and status_data['indexed']:
|
|
self.log_step("✓ Document indexing completed")
|
|
return True
|
|
|
|
except requests.exceptions.RequestException:
|
|
continue
|
|
|
|
# Also check documents list for status
|
|
if self.check_document_status_in_list():
|
|
self.log_step("✓ Document indexed (from list)")
|
|
return True
|
|
|
|
except requests.exceptions.RequestException:
|
|
pass
|
|
|
|
elapsed = int(time.time() - start_time)
|
|
self.log_step(f"Waiting... ({elapsed}s elapsed)")
|
|
time.sleep(POLL_INTERVAL)
|
|
|
|
self.log_step("✗ Indexing timeout reached", "ERROR")
|
|
return False
|
|
|
|
def check_document_status_in_list(self):
|
|
"""Check document status from documents list"""
|
|
try:
|
|
list_endpoints = [
|
|
f"{BASE_URL}/documents",
|
|
f"{API_URL}/documents"
|
|
]
|
|
|
|
for endpoint in list_endpoints:
|
|
try:
|
|
response = self.session.get(endpoint, timeout=10)
|
|
if response.status_code == 200:
|
|
documents = response.json()
|
|
for doc in documents:
|
|
if doc.get('id') == self.doc_id:
|
|
# Check if document has indexing status
|
|
if doc.get('status') in ['completed', 'indexed']:
|
|
return True
|
|
elif doc.get('indexed') is True:
|
|
return True
|
|
return False
|
|
except:
|
|
continue
|
|
return False
|
|
except:
|
|
return False
|
|
|
|
def test_search_functionality(self):
|
|
"""Test search functionality with OCR content"""
|
|
self.log_step("Testing search functionality...")
|
|
|
|
search_payload = {
|
|
"query": TEST_QUERY,
|
|
"top_k": 5
|
|
}
|
|
|
|
search_endpoints = [
|
|
f"{BASE_URL}/search",
|
|
f"{API_URL}/search",
|
|
f"{BASE_URL}/query",
|
|
f"{API_URL}/query",
|
|
f"{BASE_URL}/documents/search",
|
|
f"{API_URL}/documents/search"
|
|
]
|
|
|
|
for endpoint in search_endpoints:
|
|
try:
|
|
self.log_step(f"Testing search endpoint: {endpoint}")
|
|
response = self.session.post(
|
|
endpoint,
|
|
json=search_payload,
|
|
timeout=15
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
results = response.json()
|
|
self.log_step("✓ Search request successful")
|
|
|
|
# Validate search results structure
|
|
if isinstance(results, list) and len(results) > 0:
|
|
self.log_step(f"✓ Search returned {len(results)} results")
|
|
|
|
# Check if results contain relevant content
|
|
for i, result in enumerate(results[:3]):
|
|
if isinstance(result, dict):
|
|
content = result.get('content', result.get('text', str(result)))
|
|
else:
|
|
content = str(result)
|
|
content_preview = content[:100] + "..." if len(content) > 100 else content
|
|
self.log_step(f"Result {i+1}: {content_preview}")
|
|
|
|
return True
|
|
else:
|
|
self.log_step("✗ Search returned no results", "WARNING")
|
|
# Continue to next endpoint
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
self.log_step(f"Search endpoint {endpoint} failed: {e}", "WARNING")
|
|
continue
|
|
|
|
self.log_step("✗ All search endpoints failed", "ERROR")
|
|
return False
|
|
|
|
def run_complete_validation(self):
|
|
"""Run the complete validation workflow"""
|
|
self.log_step("Starting OCR PDF Upload and Search Validation")
|
|
self.log_step("=" * 50)
|
|
|
|
steps = [
|
|
("Server Status Check", self.check_server_status),
|
|
("Authentication", self.authenticate),
|
|
("OCR PDF Verification", self.verify_ocr_pdf_exists),
|
|
("PDF Upload", self.upload_ocr_pdf),
|
|
("Indexing Wait", self.wait_for_indexing),
|
|
("Search Test", self.test_search_functionality)
|
|
]
|
|
|
|
results = []
|
|
for step_name, step_func in steps:
|
|
self.log_step(f"Executing: {step_name}")
|
|
success = step_func()
|
|
results.append((step_name, success))
|
|
|
|
if not success:
|
|
self.log_step(f"✗ Workflow failed at: {step_name}", "ERROR")
|
|
break
|
|
|
|
# Generate final report
|
|
self.log_step("=" * 50)
|
|
self.log_step("VALIDATION RESULTS SUMMARY")
|
|
self.log_step("=" * 50)
|
|
|
|
passed = 0
|
|
total = len(results)
|
|
|
|
for step_name, success in results:
|
|
status = "✓ PASS" if success else "✗ FAIL"
|
|
self.log_step(f"{step_name}: {status}")
|
|
if success:
|
|
passed += 1
|
|
|
|
success_rate = (passed / total) * 100
|
|
self.log_step(f"Success Rate: {passed}/{total} ({success_rate:.1f}%)")
|
|
|
|
if passed == total:
|
|
self.log_step("🎉 COMPLETE WORKFLOW VALIDATION SUCCESSFUL!", "SUCCESS")
|
|
return True
|
|
else:
|
|
self.log_step("❌ WORKFLOW VALIDATION FAILED", "ERROR")
|
|
return False
|
|
|
|
def main():
|
|
"""Main execution function"""
|
|
validator = OCRWorkflowValidator()
|
|
|
|
try:
|
|
success = validator.run_complete_validation()
|
|
sys.exit(0 if success else 1)
|
|
|
|
except KeyboardInterrupt:
|
|
validator.log_step("Validation interrupted by user", "WARNING")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
validator.log_step(f"Unexpected error: {e}", "ERROR")
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
main() |