test.py

"""
Test suite for Vayumi AI Service
Tests: /ai/embed, /ai/synthesize, /ai/email/filter
"""

import requests
import json
from typing import Dict, Any
import time
import sys


# ============================================================================
# CONFIGURATION
# ============================================================================

BASE_URL = "https://dharambirAgrawal-Veda-Assistant.hf.space"
# For local: BASE_URL = "http://localhost:7860"


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

def print_section(title: str):
    """Print a formatted section header"""
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80 + "\n")


def print_result(response: requests.Response, show_full_response: bool = False):
    """Print formatted API response"""
    print(f"Status Code: {response.status_code}")
    print(f"Response Time: {response.elapsed.total_seconds():.2f}s")
    
    if response.status_code == 200:
        data = response.json()

        if show_full_response:
            print(f"\nFull Response:\n{json.dumps(data, indent=2)}")
        else:
            print(f"\nResponse Preview:")
            if 'response' in data:
                preview = data['response'][:200] + "..." if len(data.get('response', '')) > 200 else data.get('response', '')
                print(f"  Generated Text: {preview}")
            if 'embedding' in data:
                print(f"  Embedding (first 5): {data['embedding'][:5]}")
                print(f"  Dimension: {data['dimension']}")
            if 'embeddings' in data:
                print(f"  Batch Size: {len(data['embeddings'])}")
                print(f"  Dimension: {data['dimension']}")
            if 'is_important' in data:
                print(f"  Is Important: {data['is_important']}")
                print(f"  Summary: {data['summary']}")
            if 'processing_time_ms' in data:
                print(f"  Processing Time: {data['processing_time_ms']}ms")
    else:
        print(f"\n❌ Error: {response.text}")


# ============================================================================
# TEST 1: HEALTH CHECK
# ============================================================================

def test_health_check():
    """Test health endpoint"""
    print_section("TEST 1: Health Check")
    
    response = requests.get(f"{BASE_URL}/health")
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 2: /ai/embed - Single Text Embedding
# ============================================================================

def test_embed_single():
    """Test embedding a single text"""
    print_section("TEST 2: Single Text Embedding")
    
    payload = {
        "text": "Decision to increase marketing budget by 15%"
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/embed",
        json=payload
    )
    
    print_result(response)


# ============================================================================
# TEST 3: /ai/embed/batch - Batch Text Embedding
# ============================================================================

def test_embed_batch():
    """Test embedding multiple texts"""
    print_section("TEST 3: Batch Text Embedding")
    
    payload = {
        "texts": [
            "Meeting with investors tomorrow at 3pm",
            "Increase marketing budget by 15%",
            "Launch new product next quarter",
            "Hire 2 engineers by end of Q2"
        ]
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/embed/batch",
        json=payload
    )
    
    print_result(response)


# ============================================================================
# TEST 4: /ai/synthesize - Daily Briefing
# ============================================================================

def test_synthesize_daily_briefing():
    """Test generating a daily briefing (heavy synthesis)"""
    print_section("TEST 4: Daily Briefing Synthesis")
    
    prompt = """Create a professional daily briefing for a startup founder based on these memories:

MEMORIES:
1. Decision to increase marketing budget by 15% next quarter
2. Sarah assigned to lead marketing campaign
3. Need to hire 2 engineers by end of Q2
4. Product launch scheduled for April 15th
5. Investor call at 3pm tomorrow
6. Q1 revenue exceeded targets by 12%

Write a concise, professional briefing (2-3 paragraphs) highlighting:
- Key decisions made
- Action items requiring attention
- Important upcoming events

Tone: Professional, executive-level"""
    
    payload = {
        "prompt": prompt,
        "max_tokens": 800,
        "temperature": 0.7
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/synthesize",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 5: /ai/synthesize - Email Composition
# ============================================================================

def test_synthesize_email():
    """Test composing a professional email"""
    print_section("TEST 5: Professional Email Composition")
    
    prompt = """Compose a professional email to investors based on this context:

CONTEXT:
- Q1 revenue exceeded targets by 12%
- Product launch scheduled for April 15th
- Secured 3 new enterprise clients
- Expanding to European markets in Q3
- Team grew from 15 to 22 people

Write a concise update email (3-4 paragraphs):
- Thank investors for their support
- Highlight Q1 achievements
- Preview Q2 milestones
- Express optimism about growth trajectory

Tone: Professional, confident, grateful"""
    
    payload = {
        "prompt": prompt,
        "max_tokens": 600,
        "temperature": 0.7
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/synthesize",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 6: /ai/email/filter - Important Email (Job Seeker)
# ============================================================================

def test_email_filter_important():
    """Test filtering an important email"""
    print_section("TEST 6: Email Filter - Important Email")
    
    payload = {
        "sender_email": "hr@techcompany.com",
        "subject": "Interview Invitation for Software Engineer Position",
        "body": """Dear John,

We are pleased to inform you that after reviewing your application, 
we would like to invite you for an interview for the Software Engineer position.

The interview is scheduled for next Monday at 2:00 PM.
Please confirm your availability by replying to this email.

Best regards,
HR Team""",
        "user_type": "job_seeker",
        "user_name": "John Doe",
        "include_details": False
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/email/filter",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 7: /ai/email/filter - Spam/Marketing Email
# ============================================================================

def test_email_filter_spam():
    """Test filtering a marketing email"""
    print_section("TEST 7: Email Filter - Marketing/Spam Email")
    
    payload = {
        "sender_email": "newsletter@randomstore.com",
        "subject": "🎉 FLASH SALE! 50% OFF Everything - Limited Time Only!",
        "body": """Don't miss out on our biggest sale of the year!

50% OFF on ALL items - Shop Now!
Free shipping on orders over $50
Use code: SAVE50

Unsubscribe from this newsletter | View in browser
This email was sent to you because you subscribed to our newsletter.""",
        "user_type": "professional",
        "user_name": "John",
        "include_details": False
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/email/filter",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 8: /ai/email/filter - Detailed Analysis
# ============================================================================

def test_email_filter_detailed():
    """Test email filter with detailed analysis"""
    print_section("TEST 8: Email Filter - Detailed Analysis")
    
    payload = {
        "sender_email": "professor.smith@university.edu",
        "subject": "Urgent: Assignment Deadline Extended",
        "body": """Dear Student,

Due to the technical issues reported by several students, 
I am extending the deadline for the final assignment by 48 hours.

New deadline: Friday, April 18th at 11:59 PM

Please make sure to submit your work on time.

Best,
Prof. Smith""",
        "user_type": "student",
        "user_name": "Student",
        "include_details": True
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/email/filter",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 9: /ai/email/filter/batch - Batch Email Filtering
# ============================================================================

def test_email_filter_batch():
    """Test batch email filtering"""
    print_section("TEST 9: Email Filter - Batch Processing")
    
    payload = {
        "emails": [
            {
                "sender": "hr@company.com",
                "subject": "You've been selected for final interview",
                "body": "Congratulations! We'd like to invite you for the final round interview."
            },
            {
                "sender": "newsletter@store.com",
                "subject": "Weekly Deals Inside!",
                "body": "Check out this week's top deals. Unsubscribe here."
            },
            {
                "sender": "recruiter@linkedin.com",
                "subject": "A recruiter viewed your profile",
                "body": "Your profile was viewed by 5 recruiters this week."
            },
            {
                "sender": "hiring@startup.io",
                "subject": "Offer Letter - Software Developer",
                "body": "We are pleased to offer you the position of Software Developer with a starting salary..."
            }
        ],
        "user_type": "job_seeker",
        "user_name": "John"
    }
    
    print(f"Request Payload:\n{json.dumps(payload, indent=2)}\n")
    
    response = requests.post(
        f"{BASE_URL}/ai/email/filter/batch",
        json=payload
    )
    
    print_result(response, show_full_response=True)


# ============================================================================
# TEST 10: Error Handling
# ============================================================================

def test_error_handling():
    """Test error handling with invalid requests"""
    print_section("TEST 10: Error Handling")
    
    # Test 1: Empty text for embedding
    print("\n--- Test 10.1: Empty text for embedding ---")
    response = requests.post(
        f"{BASE_URL}/ai/embed",
        json={"text": ""}
    )
    print_result(response)
    
    # Test 2: Missing required field for synthesize
    print("\n--- Test 10.2: Missing 'prompt' field for synthesize ---")
    response = requests.post(
        f"{BASE_URL}/ai/synthesize",
        json={"max_tokens": 100}
    )
    print_result(response)
    
    # Test 3: Invalid user_type for email filter
    print("\n--- Test 10.3: Invalid user_type ---")
    response = requests.post(
        f"{BASE_URL}/ai/email/filter",
        json={
            "sender_email": "test@test.com",
            "subject": "Test",
            "body": "Test body",
            "user_type": "invalid_type"
        }
    )
    print_result(response)


# ============================================================================
# TEST 11: Performance Test - Embeddings
# ============================================================================

def test_embedding_performance():
    """Test embedding performance"""
    print_section("TEST 11: Embedding Performance")
    
    texts = [
        f"Test message number {i} for performance testing" for i in range(10)
    ]
    
    print(f"Testing batch embedding with {len(texts)} texts...\n")
    
    start_time = time.time()
    
    response = requests.post(
        f"{BASE_URL}/ai/embed/batch",
        json={"texts": texts}
    )
    
    total_time = time.time() - start_time
    
    print_result(response)
    print(f"\n--- Performance Summary ---")
    print(f"Total Time (including network): {total_time:.2f}s")
    if response.status_code == 200:
        data = response.json()
        print(f"Server Processing Time: {data.get('processing_time_ms', 'N/A')}ms")
        print(f"Texts per second: {len(texts)/total_time:.2f}")


# ============================================================================
# TEST 12: Email Filter Performance
# ============================================================================

def test_email_filter_performance():
    """Test email filter performance"""
    print_section("TEST 12: Email Filter Performance")
    
    emails = [
        {
            "sender": f"sender{i}@company.com",
            "subject": f"Test email subject {i}",
            "body": f"This is test email body number {i} with some content."
        }
        for i in range(5)
    ]
    
    print(f"Testing batch email filtering with {len(emails)} emails...\n")
    
    start_time = time.time()
    
    response = requests.post(
        f"{BASE_URL}/ai/email/filter/batch",
        json={
            "emails": emails,
            "user_type": "professional"
        }
    )
    
    total_time = time.time() - start_time
    
    print_result(response)
    print(f"\n--- Performance Summary ---")
    print(f"Total Time (including network): {total_time:.2f}s")
    if response.status_code == 200:
        data = response.json()
        print(f"Server Processing Time: {data.get('processing_time_ms', 'N/A')}ms")
        print(f"Emails per second: {len(emails)/total_time:.2f}")


# ============================================================================
# MAIN TEST RUNNER
# ============================================================================

def run_all_tests():
    """Run all test cases"""
    print("\n" + "🧪" * 40)
    print("  VAYUMI AI SERVICE - COMPREHENSIVE TEST SUITE")
    print("🧪" * 40)
    
    try:
        # Basic tests
        test_health_check()
        
        # Embedding tests
        test_embed_single()
        test_embed_batch()
        
        # Synthesis tests (requires HF_TOKEN)
        test_synthesize_daily_briefing()
        test_synthesize_email()
        
        # Email filter tests
        test_email_filter_important()
        test_email_filter_spam()
        test_email_filter_detailed()
        test_email_filter_batch()
        
        # Error handling
        test_error_handling()
        
        # Performance
        test_embedding_performance()
        test_email_filter_performance()
        
        print("\n" + "✅" * 40)
        print("  ALL TESTS COMPLETED")
        print("✅" * 40 + "\n")
        
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {str(e)}")


def run_quick_tests():
    """Run only quick tests (skip synthesis)"""
    print("\n" + "🧪" * 40)
    print("  VAYUMI AI SERVICE - QUICK TEST SUITE")
    print("🧪" * 40)
    
    try:
        test_health_check()
        test_embed_single()
        test_email_filter_important()
        test_email_filter_spam()
        
        print("\n" + "✅" * 40)
        print("  QUICK TESTS COMPLETED")
        print("✅" * 40 + "\n")
        
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {str(e)}")


def run_email_tests():
    """Run only email filter tests"""
    print("\n" + "📧" * 40)
    print("  EMAIL FILTER TEST SUITE")
    print("📧" * 40)
    
    try:
        test_email_filter_important()
        test_email_filter_spam()
        test_email_filter_detailed()
        test_email_filter_batch()
        test_email_filter_performance()
        
        print("\n" + "✅" * 40)
        print("  EMAIL FILTER TESTS COMPLETED")
        print("✅" * 40 + "\n")
        
    except Exception as e:
        print(f"\n❌ Test suite failed with error: {str(e)}")


# ============================================================================
# INTERACTIVE MENU
# ============================================================================

def interactive_menu():
    """Interactive test menu"""
    while True:
        print("\n" + "=" * 80)
        print("  VAYUMI AI SERVICE - TEST MENU")
        print("=" * 80)
        print("\n--- Basic Tests ---")
        print("1.  Health Check")
        print("\n--- Embedding Tests ---")
        print("2.  Single Text Embedding")
        print("3.  Batch Text Embedding")
        print("\n--- Synthesis Tests (requires HF_TOKEN) ---")
        print("4.  Daily Briefing Synthesis")
        print("5.  Email Composition Synthesis")
        print("\n--- Email Filter Tests ---")
        print("6.  Filter Important Email")
        print("7.  Filter Spam/Marketing Email")
        print("8.  Filter with Detailed Analysis")
        print("9.  Batch Email Filtering")
        print("\n--- Other Tests ---")
        print("10. Error Handling Tests")
        print("11. Embedding Performance Test")
        print("12. Email Filter Performance Test")
        print("\n--- Test Suites ---")
        print("13. Run Quick Tests (Health, Embed, Email Filter)")
        print("14. Run Email Filter Tests Only")
        print("15. Run All Tests")
        print("\n0.  Exit")
        print(f"\nCurrent URL: {BASE_URL}")
        
        choice = input("\nEnter your choice: ").strip()
        
        if choice == "1":
            test_health_check()
        elif choice == "2":
            test_embed_single()
        elif choice == "3":
            test_embed_batch()
        elif choice == "4":
            test_synthesize_daily_briefing()
        elif choice == "5":
            test_synthesize_email()
        elif choice == "6":
            test_email_filter_important()
        elif choice == "7":
            test_email_filter_spam()
        elif choice == "8":
            test_email_filter_detailed()
        elif choice == "9":
            test_email_filter_batch()
        elif choice == "10":
            test_error_handling()
        elif choice == "11":
            test_embedding_performance()
        elif choice == "12":
            test_email_filter_performance()
        elif choice == "13":
            run_quick_tests()
        elif choice == "14":
            run_email_tests()
        elif choice == "15":
            run_all_tests()
        elif choice == "0":
            print("\n👋 Goodbye!")
            break
        else:
            print("\n❌ Invalid choice. Please try again.")


# ============================================================================
# RUN TESTS
# ============================================================================

if __name__ == "__main__":
    # Check if URL argument provided
    if len(sys.argv) > 1:
        BASE_URL = sys.argv[1]
        print(f"\n🔗 Using custom URL: {BASE_URL}")
    
    # Run interactive menu
    interactive_menu()