scripts/test_environment.py

import sys
from pathlib import Path
import os

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from config.settings import Settings
import requests


def test_openai_connection():
    """Test OpenAI API connection"""
    print("\n=== Testing OpenAI API Connection ===")
    
    # Check API key
    if not Settings.OPENAI_API_KEY:
        print("❌ OpenAI API key not found!")
        return False
    print(f"✓ OpenAI API key found (length: {len(Settings.OPENAI_API_KEY)})")
    
    # Test API reachability with more detail
    try:
        print("Testing connection to api.openai.com...")
        response = requests.get(
            "https://api.openai.com",
            timeout=10,
            verify=True  # Force SSL verification
        )
        print(f"✓ OpenAI API reachable (Status: {response.status_code})")
        print(f"Response headers: {dict(response.headers)}")
    except requests.exceptions.SSLError as e:
        print(f"❌ SSL Error connecting to OpenAI: {e}")
        return False
    except requests.exceptions.ConnectionError as e:
        print(f"❌ Connection Error: {e}")
        print("Checking if proxy is needed...")
        # Try to get environment proxy settings
        print(f"HTTP_PROXY: {os.environ.get('HTTP_PROXY')}")
        print(f"HTTPS_PROXY: {os.environ.get('HTTPS_PROXY')}")
        return False
    except Exception as e:
        print(f"❌ Cannot reach OpenAI API: {type(e).__name__}: {e}")
        return False
    
    # Test embeddings with more detail
    try:
        print("\nTesting embeddings creation...")
        embeddings = OpenAIEmbeddings(
            openai_api_key=Settings.OPENAI_API_KEY,
            timeout=30,
            max_retries=2
        )
        print("Embeddings object created, attempting query...")
        result = embeddings.embed_query("test")
        print(f"✓ Embeddings working (vector size: {len(result)})")
        return True
    except Exception as e:
        print(f"❌ Embeddings error: {type(e).__name__}: {e}")
        print("Stack trace:")
        import traceback
        traceback.print_exc()
        return False


def test_chroma_setup():
    """Test ChromaDB setup and functionality"""
    print("\n=== Testing ChromaDB Setup ===")
    
    # Check paths
    chroma_dir = Settings.get_chroma_path()
    print(f"Chroma directory: {chroma_dir}")
    
    if not chroma_dir.exists():
        print("❌ Chroma directory not found!")
        return False
    print("✓ Chroma directory exists")
    
    # Check SQLite file
    sqlite_file = chroma_dir / "chroma.sqlite3"
    if not sqlite_file.exists():
        print("❌ SQLite database not found!")
        return False
    print(f"✓ SQLite database found ({sqlite_file.stat().st_size / (1024*1024):.2f} MB)")
    
    # Test ChromaDB functionality
    try:
        embeddings = OpenAIEmbeddings(
            openai_api_key=Settings.OPENAI_API_KEY
        )
        vector_store = Chroma(
            persist_directory=str(chroma_dir),
            embedding_function=embeddings,
            collection_name=Settings.CHROMA_COLLECTION_NAME
        )
        
        # Check collection
        collection = vector_store._collection
        count = collection.count()
        print(f"✓ Collection loaded ({count} documents)")
        
        # Test search functionality
        results = vector_store.similarity_search_with_score("test", k=1)
        print("✓ Search functionality working")
        
        return True
    except Exception as e:
        print(f"❌ ChromaDB error: {e}")
        return False


def test_huggingface_setup():
    """Test HuggingFace-specific setup"""
    print("\n=== Testing HuggingFace Setup ===")
    
    if not Settings.is_huggingface():
        print("Skipping HuggingFace tests (not in HF environment)")
        return True
    
    # Check HF token
    if not Settings.HF_TOKEN:
        print("❌ HuggingFace token not found!")
        return False
    print("✓ HuggingFace token found")
    
    # Check dataset access
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=Settings.HF_TOKEN)
        files = api.list_repo_files(Settings.HF_DATASET, repo_type="dataset")
        print(f"✓ Dataset accessible ({len(files)} files)")
        return True
    except Exception as e:
        print(f"❌ Dataset access error: {e}")
        return False


def test_environment():
    """Run all environment tests"""
    print(f"\nTesting environment: {Settings.DEPLOYMENT_MODE}")
    
    # Test HuggingFace setup first if applicable
    if Settings.is_huggingface():
        hf_ok = test_huggingface_setup()
        if not hf_ok:
            print("\n❌ HuggingFace setup failed!")
            return False
    
    # Test OpenAI
    openai_ok = test_openai_connection()
    
    # Test ChromaDB
    chroma_ok = test_chroma_setup()
    
    # Summary
    print("\n=== Test Summary ===")
    print(f"OpenAI Connection: {'✓' if openai_ok else '❌'}")
    print(f"ChromaDB Setup: {'✓' if chroma_ok else '❌'}")
    if Settings.is_huggingface():
        print(f"HuggingFace Setup: {'✓' if hf_ok else '❌'}")
    
    return all([openai_ok, chroma_ok])


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