test_hf_integration.py

"""
Test script for Hugging Face Inference API integration
Run this to verify the setup before deploying to HF Spaces
"""
import os
import sys
import logging

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


def test_hf_token():
    """Test if HF_TOKEN is set"""
    logger.info("Testing HF_TOKEN...")
    token = os.getenv("HF_TOKEN", "")
    
    if not token:
        logger.error("❌ HF_TOKEN not found!")
        logger.info("Please set it: export HF_TOKEN=hf_your_token_here")
        return False
    
    if not token.startswith("hf_"):
        logger.error("❌ Invalid HF_TOKEN format (should start with 'hf_')")
        return False
    
    logger.info(f"✅ HF_TOKEN found: {token[:10]}...")
    return True


def test_imports():
    """Test if all required packages are installed"""
    logger.info("\nTesting imports...")
    
    required_packages = [
        ("huggingface_hub", "Hugging Face Hub"),
        ("gradio", "Gradio"),
        ("chromadb", "ChromaDB"),
        ("sentence_transformers", "Sentence Transformers"),
        ("langchain_text_splitters", "LangChain Text Splitters"),
    ]
    
    all_ok = True
    for package, name in required_packages:
        try:
            __import__(package)
            logger.info(f"✅ {name} installed")
        except ImportError:
            logger.error(f"❌ {name} not installed")
            all_ok = False
    
    return all_ok


def test_hf_api():
    """Test Hugging Face Inference API connection"""
    logger.info("\nTesting HF Inference API...")
    
    try:
        from huggingface_hub import InferenceClient
        from config import HF_TOKEN, HF_MODEL
        
        if not HF_TOKEN:
            logger.error("❌ Cannot test API without HF_TOKEN")
            return False
        
        client = InferenceClient(token=HF_TOKEN)
        
        # Test with a simple prompt
        logger.info(f"Testing model: {HF_MODEL}")
        logger.info("Sending test request...")
        
        response = client.text_generation(
            prompt="Say 'Hello, World!' and nothing else.",
            model=HF_MODEL,
            max_new_tokens=20,
            temperature=0.1,
        )
        
        logger.info(f"✅ API Response: {response}")
        return True
        
    except Exception as e:
        logger.error(f"❌ API Test failed: {e}")
        return False


def test_llm_handler():
    """Test the LLM handler module"""
    logger.info("\nTesting LLM Handler...")
    
    try:
        from llm_handler import LLMHandler
        
        llm = LLMHandler()
        logger.info("✅ LLM Handler initialized")
        
        # Test answer generation
        logger.info("Testing answer generation...")
        test_question = "What is 2+2?"
        test_context = "Basic arithmetic: 2+2 equals 4."
        
        answer = llm.generate_answer(test_question, test_context, stream=False)
        logger.info(f"✅ Generated answer: {answer[:100]}...")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ LLM Handler test failed: {e}")
        return False


def test_vector_store():
    """Test vector store initialization"""
    logger.info("\nTesting Vector Store...")
    
    try:
        from vector_store import VectorStore
        
        vs = VectorStore()
        logger.info("✅ Vector Store initialized")
        
        stats = vs.get_collection_stats()
        logger.info(f"✅ Collection stats: {stats}")
        
        return True
        
    except Exception as e:
        logger.error(f"❌ Vector Store test failed: {e}")
        return False


def main():
    """Run all tests"""
    logger.info("="*60)
    logger.info("RAG System - Hugging Face Integration Tests")
    logger.info("="*60)
    
    tests = [
        ("HF Token", test_hf_token),
        ("Package Imports", test_imports),
        ("HF API Connection", test_hf_api),
        ("LLM Handler", test_llm_handler),
        ("Vector Store", test_vector_store),
    ]
    
    results = {}
    for test_name, test_func in tests:
        try:
            results[test_name] = test_func()
        except Exception as e:
            logger.error(f"❌ {test_name} crashed: {e}")
            results[test_name] = False
    
    # Summary
    logger.info("\n" + "="*60)
    logger.info("Test Summary")
    logger.info("="*60)
    
    for test_name, passed in results.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        logger.info(f"{status} - {test_name}")
    
    all_passed = all(results.values())
    
    logger.info("="*60)
    if all_passed:
        logger.info("🎉 All tests passed! Ready to deploy to HF Spaces.")
        logger.info("\nNext steps:")
        logger.info("1. Create a new Space at https://huggingface.co/new-space")
        logger.info("2. Choose 'Gradio' as SDK")
        logger.info("3. Add HF_TOKEN to Space secrets")
        logger.info("4. Push code to the Space repository")
        return 0
    else:
        logger.error("❌ Some tests failed. Please fix the issues before deploying.")
        return 1


if __name__ == "__main__":
    sys.exit(main())