test_api.py

"""
API Testing Script for Lovebird Reply Generator

This script helps you test the deployed Hugging Face Space API.

Usage:
    python test_api.py --space_url https://your-space-name.hf.space
"""

import argparse
import requests
import json
from typing import Dict, Any


def test_intent_api(space_url: str, text: str):
    """Test Intent classification API"""
    api_url = f"{space_url}/api/intent"
    
    response = requests.post(api_url, json={"data": [text]})
    
    if response.status_code == 200:
        result = response.json()
        print("\n✅ Intent API Response:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        return result
    else:
        print(f"\n❌ Intent API Error: {response.status_code}")
        print(response.text)
        return None


def test_tone_api(space_url: str, text: str):
    """Test Tone classification API"""
    api_url = f"{space_url}/api/tone"
    
    response = requests.post(api_url, json={"data": [text]})
    
    if response.status_code == 200:
        result = response.json()
        print("\n✅ Tone API Response:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        return result
    else:
        print(f"\n❌ Tone API Error: {response.status_code}")
        print(response.text)
        return None


def test_reply_api(space_url: str, conversation: str):
    """Test Reply Suggestion API"""
    api_url = f"{space_url}/api/reply"
    
    response = requests.post(api_url, json={"data": [conversation]})
    
    if response.status_code == 200:
        result = response.json()
        print("\n✅ Reply Suggestion API Response:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        return result
    else:
        print(f"\n❌ Reply API Error: {response.status_code}")
        print(response.text)
        return None


def test_gradio_api(space_url: str, conversation: str, fn_index: int = 2):
    """Test using Gradio API endpoint"""
    api_url = f"{space_url}/api/predict"
    
    payload = {
        "data": [conversation],
        "fn_index": fn_index  # 2 = Reply Suggestion tab
    }
    
    response = requests.post(api_url, json=payload)
    
    if response.status_code == 200:
        result = response.json()
        print("\n✅ Gradio API Response:")
        print(json.dumps(result, indent=2, ensure_ascii=False))
        return result
    else:
        print(f"\n❌ Gradio API Error: {response.status_code}")
        print(response.text)
        return None


def main():
    parser = argparse.ArgumentParser(description="Test Lovebird API endpoints")
    parser.add_argument(
        "--space_url",
        type=str,
        required=True,
        help="URL of your Hugging Face Space (e.g., https://your-space-name.hf.space)"
    )
    parser.add_argument(
        "--test_conversation",
        type=str,
        default="Tối nay anh có lịch đột xuất. ||| Thế mai được không?",
        help="Conversation to test (format: 'male ||| female')"
    )
    parser.add_argument(
        "--test_text",
        type=str,
        default="Tối nay anh có lịch đột xuất. ||| Thế mai được không?",
        help="Text to test for intent/tone"
    )
    
    args = parser.parse_args()
    
    print("=" * 60)
    print("Testing Lovebird API")
    print("=" * 60)
    print(f"Space URL: {args.space_url}")
    print(f"Test conversation: {args.test_conversation}")
    print()
    
    # Test 1: Intent API
    print("=" * 60)
    print("Test 1: Intent Classification")
    print("=" * 60)
    intent_result = test_intent_api(args.space_url, args.test_text)
    
    # Test 2: Tone API
    print("\n" + "=" * 60)
    print("Test 2: Tone Classification")
    print("=" * 60)
    tone_result = test_tone_api(args.space_url, args.test_text)
    
    # Test 3: Reply Suggestion API
    print("\n" + "=" * 60)
    print("Test 3: Reply Suggestion (Direct API)")
    print("=" * 60)
    reply_result = test_reply_api(args.space_url, args.test_conversation)
    
    # Test 4: Gradio API (alternative)
    print("\n" + "=" * 60)
    print("Test 4: Reply Suggestion (Gradio API)")
    print("=" * 60)
    gradio_result = test_gradio_api(args.space_url, args.test_conversation)
    
    # Summary
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    print(f"Intent API: {'✅ PASS' if intent_result else '❌ FAIL'}")
    print(f"Tone API: {'✅ PASS' if tone_result else '❌ FAIL'}")
    print(f"Reply API: {'✅ PASS' if reply_result else '❌ FAIL'}")
    print(f"Gradio API: {'✅ PASS' if gradio_result else '❌ FAIL'}")
    
    # Show generated reply if available
    if reply_result and "data" in reply_result:
        reply_data = reply_result["data"][0] if isinstance(reply_result["data"], list) else reply_result["data"]
        if isinstance(reply_data, list) and len(reply_data) > 1:
            suggestion = reply_data[1]  # Second element is the suggestion text
            print(f"\n💬 Generated Reply: {suggestion}")


if __name__ == "__main__":
    main()