test_chatbot.py

#!/usr/bin/env python3
"""
Test script for the WUEB Document Assistant
This script demonstrates the core functionality and can be used for testing.
"""

import os
import sys
import logging
from dotenv import load_dotenv

# Add the current directory to Python path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from chatbot import WUEBChatbot
from data_loader import DataLoader
from config import Config

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def test_system_initialization():
    """Test if the system can be initialized properly."""
    print("🔧 Testing system initialization...")
    
    try:
        # Test configuration
        config = Config()
        print(f"✅ Configuration loaded successfully")
        print(f"   - OpenAI Model: {config.OPENAI_MODEL}")
        print(f"   - Embedding Model: {config.EMBEDDING_MODEL}")
        print(f"   - Vector DB Path: {config.VECTOR_DB_PATH}")
        
        # Test chatbot initialization
        chatbot = WUEBChatbot()
        print("✅ Chatbot initialized successfully")
        
        # Test data loader initialization
        data_loader = DataLoader()
        print("✅ Data loader initialized successfully")
        
        return True
        
    except Exception as e:
        print(f"❌ System initialization failed: {str(e)}")
        return False

def test_pdf_directory():
    """Test PDF directory validation."""
    print("\n📁 Testing PDF directory...")
    
    try:
        data_loader = DataLoader()
        validation = data_loader.validate_pdf_directory()
        
        if validation['valid']:
            print(f"✅ PDF directory found: {validation['pdf_directory']}")
            print(f"   - PDF files found: {validation['total_files']}")
            if validation['pdf_files']:
                print(f"   - Files: {', '.join(validation['pdf_files'][:3])}")
                if len(validation['pdf_files']) > 3:
                    print(f"   - ... and {len(validation['pdf_files']) - 3} more")
        else:
            print(f"❌ PDF directory error: {validation['error']}")
            print("   Please create a 'pdfs' directory and add your PDF documents")
        
        return validation['valid']
        
    except Exception as e:
        print(f"❌ PDF directory test failed: {str(e)}")
        return False

def test_document_loading():
    """Test document loading functionality."""
    print("\n📚 Testing document loading...")
    
    try:
        data_loader = DataLoader()
        
        # Check if documents are already loaded
        stats = data_loader.get_processing_stats()
        if stats.get('total_chunks', 0) > 0:
            print(f"✅ Documents already loaded ({stats['total_chunks']} chunks)")
            return True
        
        # Try to load documents
        print("   Attempting to load documents...")
        result = data_loader.load_documents()
        
        if result['success']:
            print(f"✅ Documents loaded successfully")
            print(f"   - Documents processed: {result['documents_processed']}")
            print(f"   - Chunks created: {result['chunks_created']}")
            return True
        else:
            print(f"❌ Document loading failed: {result.get('error', 'Unknown error')}")
            return False
            
    except Exception as e:
        print(f"❌ Document loading test failed: {str(e)}")
        return False

def test_chatbot_queries():
    """Test chatbot with sample queries."""
    print("\n🤖 Testing chatbot queries...")
    
    try:
        chatbot = WUEBChatbot()
        
        # Sample queries to test
        test_queries = [
            "What are the admission requirements?",
            "How do I apply for a program?",
            "What are the tuition fees?",
            "Tell me about the university structure",
            "What are the academic calendar dates?"
        ]
        
        for i, query in enumerate(test_queries, 1):
            print(f"\n   Query {i}: {query}")
            
            # Validate query
            validation = chatbot.validate_query(query)
            if not validation['valid']:
                print(f"   ❌ Query validation failed: {validation['reason']}")
                continue
            
            # Process query
            result = chatbot.process_query(query)
            
            print(f"   Response: {result['response'][:100]}...")
            print(f"   Confidence: {result['confidence']}")
            print(f"   Documents found: {result['total_documents_found']}")
            print(f"   Relevant documents: {result['filtered_documents']}")
            
            if result['confidence'] == 'error':
                print(f"   ❌ Error: {result.get('error', 'Unknown error')}")
        
        return True
        
    except Exception as e:
        print(f"❌ Chatbot query test failed: {str(e)}")
        return False

def test_system_info():
    """Test system information retrieval."""
    print("\nℹ️ Testing system information...")
    
    try:
        chatbot = WUEBChatbot()
        system_info = chatbot.get_system_info()
        
        print("✅ System information retrieved:")
        print(f"   - System Name: {system_info['system_name']}")
        print(f"   - Model: {system_info['model']}")
        print(f"   - Embedding Model: {system_info['embedding_model']}")
        print(f"   - Similarity Threshold: {system_info['similarity_threshold']}")
        
        if 'vector_store_info' in system_info:
            vs_info = system_info['vector_store_info']
            print(f"   - Collection: {vs_info.get('collection_name', 'N/A')}")
            print(f"   - Document Count: {vs_info.get('document_count', 0)}")
        
        return True
        
    except Exception as e:
        print(f"❌ System info test failed: {str(e)}")
        return False

def main():
    """Run all tests."""
    print("🧪 WUEB Document Assistant - System Test")
    print("=" * 50)
    
    # Load environment variables
    load_dotenv()
    
    # Check for OpenAI API key
    if not os.getenv("OPENAI_API_KEY"):
        print("❌ OPENAI_API_KEY not found in environment variables")
        print("   Please create a .env file with your OpenAI API key")
        return False
    
    print("✅ OpenAI API key found")
    
    # Run tests
    tests = [
        ("System Initialization", test_system_initialization),
        ("PDF Directory", test_pdf_directory),
        ("Document Loading", test_document_loading),
        ("Chatbot Queries", test_chatbot_queries),
        ("System Information", test_system_info)
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
        except Exception as e:
            print(f"❌ {test_name} test crashed: {str(e)}")
            results.append((test_name, False))
    
    # Print summary
    print("\n" + "=" * 50)
    print("📊 Test Results Summary")
    print("=" * 50)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} - {test_name}")
        if result:
            passed += 1
    
    print(f"\nOverall: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All tests passed! The system is ready to use.")
        print("\nTo start the application, run:")
        print("   streamlit run app.py")
    else:
        print("⚠️ Some tests failed. Please check the errors above.")
    
    return passed == total

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)