test_hybrid.py

#!/usr/bin/env python3
"""
Test script for hybrid model loading
"""

import requests
import json
import time

def test_hybrid_api(base_url="https://aurasystems-spanish-embeddings-api.hf.space"):
    """Test the hybrid API"""
    
    print(f"Testing hybrid API at {base_url}")
    
    # Test health endpoint first
    try:
        response = requests.get(f"{base_url}/health")
        print(f"✓ Health endpoint: {response.status_code}")
        if response.status_code == 200:
            health_data = response.json()
            print(f"  Startup models loaded: {health_data.get('startup_models_loaded', False)}")
            print(f"  Available models: {health_data.get('available_models', [])}")
            print(f"  Note: {health_data.get('note', 'N/A')}")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Health endpoint failed: {e}")
        return False
    
    # Test startup model (jina-v3)
    try:
        payload = {
            "texts": ["Hola mundo", "Bonjour le monde"],
            "model": "jina-v3",
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed", json=payload)
        print(f"✓ Startup model (jina-v3): {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Generated {data.get('num_texts', 0)} embeddings")
            print(f"  Dimensions: {data.get('dimensions', 0)}")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Startup model test failed: {e}")
    
    # Test startup model (roberta-ca)
    try:
        payload = {
            "texts": ["Bon dia", "Com estàs?"],
            "model": "roberta-ca",
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed", json=payload)
        print(f"✓ Startup model (roberta-ca): {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Generated {data.get('num_texts', 0)} embeddings")
            print(f"  Dimensions: {data.get('dimensions', 0)}")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ Startup model test failed: {e}")
    
    # Test on-demand model (jina)
    try:
        payload = {
            "texts": ["Texto en español"],
            "model": "jina",
            "normalize": True
        }
        response = requests.post(f"{base_url}/embed", json=payload)
        print(f"✓ On-demand model (jina): {response.status_code}")
        if response.status_code == 200:
            data = response.json()
            print(f"  Generated {data.get('num_texts', 0)} embeddings")
            print(f"  Dimensions: {data.get('dimensions', 0)}")
        else:
            print(f"  Error: {response.text}")
    except Exception as e:
        print(f"✗ On-demand model test failed: {e}")
    
    # Check health again to see all models
    try:
        response = requests.get(f"{base_url}/health")
        if response.status_code == 200:
            health_data = response.json()
            print(f"✓ Final health check:")
            print(f"  All models loaded: {health_data.get('all_models_loaded', False)}")
            print(f"  Available models: {health_data.get('available_models', [])}")
    except Exception as e:
        print(f"✗ Final health check failed: {e}")
    
    return True

if __name__ == "__main__":
    test_hybrid_api()