test_inference_api.py

#!/usr/bin/env python3
"""
Test script for RESDSQL Inference API

This script tests the inference API against expected SQL and NatSQL queries
from the Spider training dataset or a custom test file.

Usage:
    python3 test_inference_api.py --api-url http://localhost:8000/infer --data-file NatSQL/NatSQLv1_6/train_spider.json --target-type natsql --limit 10
"""

import json
import argparse
import re
import time
from pathlib import Path
from typing import List, Dict, Any, Optional
from collections import defaultdict
import sys

try:
    import requests
except ImportError:
    print("Error: 'requests' library is required. Install it with: pip install requests")
    sys.exit(1)


def normalize_sql(sql: str) -> str:
    """
    Normalize SQL query for comparison by:
    - Converting to lowercase
    - Removing extra whitespace
    - Standardizing spacing around operators and keywords
    """
    if not sql:
        return ""
    
    # Convert to lowercase
    sql = sql.lower().strip()
    
    # Remove extra whitespace
    sql = re.sub(r'\s+', ' ', sql)
    
    # Standardize spacing around operators
    sql = re.sub(r'\s*=\s*', ' = ', sql)
    sql = re.sub(r'\s*>\s*', ' > ', sql)
    sql = re.sub(r'\s*<\s*', ' < ', sql)
    sql = re.sub(r'\s*>=\s*', ' >= ', sql)
    sql = re.sub(r'\s*<=\s*', ' <= ', sql)
    sql = re.sub(r'\s*!=\s*', ' != ', sql)
    sql = re.sub(r'\s*<>\s*', ' <> ', sql)
    
    # Standardize spacing around commas
    sql = re.sub(r'\s*,\s*', ', ', sql)
    
    # Standardize spacing around parentheses
    sql = re.sub(r'\s*\(\s*', ' (', sql)
    sql = re.sub(r'\s*\)\s*', ') ', sql)
    
    # Remove trailing/leading spaces
    sql = sql.strip()
    
    return sql


def compare_sql(generated: str, expected: str, normalize: bool = True) -> bool:
    """
    Compare two SQL queries.
    
    Args:
        generated: The SQL generated by the API
        expected: The expected SQL from the dataset
        normalize: Whether to normalize both queries before comparison
    
    Returns:
        True if queries match, False otherwise
    """
    if normalize:
        generated = normalize_sql(generated)
        expected = normalize_sql(expected)
    
    return generated == expected


def call_inference_api(
    api_url: str,
    question: str,
    db_id: str,
    target_type: str = "natsql",
    timeout: int = 30,
    **kwargs
) -> Dict[str, Any]:
    """
    Call the inference API and return the response.
    
    Args:
        api_url: The API endpoint URL
        question: The natural language question
        db_id: The database identifier
        target_type: "sql" or "natsql"
        timeout: Request timeout in seconds
        **kwargs: Additional API parameters
    
    Returns:
        Dictionary with API response or error information
    """
    payload = {
        "question": question,
        "db_id": db_id,
        "target_type": target_type,
        **kwargs
    }
    
    try:
        response = requests.post(api_url, json=payload, timeout=timeout)
        response.raise_for_status()
        return {
            "success": True,
            "data": response.json(),
            "error": None
        }
    except requests.exceptions.Timeout:
        return {
            "success": False,
            "data": None,
            "error": f"Request timeout after {timeout} seconds"
        }
    except requests.exceptions.RequestException as e:
        return {
            "success": False,
            "data": None,
            "error": f"Request failed: {str(e)}"
        }
    except Exception as e:
        return {
            "success": False,
            "data": None,
            "error": f"Unexpected error: {str(e)}"
        }


def load_test_data(data_file: str, limit: Optional[int] = None) -> List[Dict[str, Any]]:
    """
    Load test data from a JSON file.
    
    Args:
        data_file: Path to the JSON file containing test cases
        limit: Maximum number of test cases to load (None for all)
    
    Returns:
        List of test case dictionaries
    """
    with open(data_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if limit:
        data = data[:limit]
    
    return data


def test_single_case(
    test_case: Dict[str, Any],
    api_url: str,
    target_type: str,
    verbose: bool = False
) -> Dict[str, Any]:
    """
    Test a single test case against the API.
    
    Args:
        test_case: Dictionary containing question, db_id, query, and NatSQL
        api_url: The API endpoint URL
        target_type: "sql" or "natsql"
        verbose: Whether to print detailed information
    
    Returns:
        Dictionary with test results
    """
    question = test_case.get("question", "")
    db_id = test_case.get("db_id", "")
    
    # Get expected query based on target type
    if target_type == "natsql":
        expected_query = test_case.get("NatSQL", "")
    else:
        expected_query = test_case.get("query", "")
    
    if not question or not db_id:
        return {
            "success": False,
            "match": False,
            "error": "Missing question or db_id in test case",
            "question": question,
            "db_id": db_id
        }
    
    # Call API
    api_response = call_inference_api(api_url, question, db_id, target_type)
    
    if not api_response["success"]:
        return {
            "success": False,
            "match": False,
            "error": api_response["error"],
            "question": question,
            "db_id": db_id,
            "expected": expected_query
        }
    
    generated_sql = api_response["data"].get("sql", "")
    execution_success = api_response["data"].get("execution_success", None)
    
    # Compare queries
    match = compare_sql(generated_sql, expected_query)
    
    result = {
        "success": True,
        "match": match,
        "question": question,
        "db_id": db_id,
        "generated": generated_sql,
        "expected": expected_query,
        "execution_success": execution_success,
        "error": None
    }
    
    if verbose and not match:
        result["normalized_generated"] = normalize_sql(generated_sql)
        result["normalized_expected"] = normalize_sql(expected_query)
    
    return result


def run_tests(
    api_url: str,
    data_file: str,
    target_type: str = "natsql",
    limit: Optional[int] = None,
    verbose: bool = False,
    delay: float = 0.1,
    output_file: Optional[str] = None
) -> Dict[str, Any]:
    """
    Run tests against the inference API.
    
    Args:
        api_url: The API endpoint URL
        data_file: Path to test data JSON file
        target_type: "sql" or "natsql"
        limit: Maximum number of test cases to run
        verbose: Whether to print detailed results
        delay: Delay between API calls in seconds
        output_file: Optional file to save detailed results
    
    Returns:
        Dictionary with test summary statistics
    """
    print(f"Loading test data from: {data_file}")
    test_cases = load_test_data(data_file, limit)
    print(f"Loaded {len(test_cases)} test cases")
    print(f"Target type: {target_type}")
    print(f"API URL: {api_url}")
    print("=" * 80)
    
    results = []
    stats = {
        "total": len(test_cases),
        "successful": 0,
        "failed": 0,
        "matched": 0,
        "mismatched": 0,
        "errors": 0,
        "by_database": defaultdict(lambda: {"total": 0, "matched": 0, "failed": 0})
    }
    
    for i, test_case in enumerate(test_cases, 1):
        if verbose:
            print(f"\n[{i}/{len(test_cases)}] Testing: {test_case.get('question', 'N/A')[:60]}...")
        
        result = test_single_case(test_case, api_url, target_type, verbose)
        results.append(result)
        
        # Update statistics
        db_id = result.get("db_id", "unknown")
        stats["by_database"][db_id]["total"] += 1
        
        if result["success"]:
            stats["successful"] += 1
            if result["match"]:
                stats["matched"] += 1
                stats["by_database"][db_id]["matched"] += 1
                if verbose:
                    print(f"  ✓ MATCH")
            else:
                stats["mismatched"] += 1
                if verbose:
                    print(f"  ✗ MISMATCH")
                    print(f"    Expected: {result['expected'][:100]}...")
                    print(f"    Generated: {result['generated'][:100]}...")
        else:
            stats["failed"] += 1
            stats["errors"] += 1
            stats["by_database"][db_id]["failed"] += 1
            if verbose:
                print(f"  ✗ ERROR: {result.get('error', 'Unknown error')}")
        
        # Add delay to avoid overwhelming the API
        if delay > 0 and i < len(test_cases):
            time.sleep(delay)
    
    # Calculate accuracy
    if stats["successful"] > 0:
        accuracy = (stats["matched"] / stats["successful"]) * 100
    else:
        accuracy = 0.0
    
    # Print summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)
    print(f"Total test cases: {stats['total']}")
    print(f"Successful API calls: {stats['successful']}")
    print(f"Failed API calls: {stats['failed']}")
    print(f"Matched queries: {stats['matched']}")
    print(f"Mismatched queries: {stats['mismatched']}")
    print(f"Accuracy: {accuracy:.2f}%")
    print("=" * 80)
    
    # Print top databases by accuracy
    if stats["by_database"]:
        print("\nAccuracy by Database (top 10):")
        db_accuracies = []
        for db_id, db_stats in stats["by_database"].items():
            if db_stats["total"] > 0:
                db_accuracy = (db_stats["matched"] / db_stats["total"]) * 100
                db_accuracies.append((db_id, db_accuracy, db_stats["matched"], db_stats["total"]))
        
        db_accuracies.sort(key=lambda x: x[1], reverse=True)
        for db_id, acc, matched, total in db_accuracies[:10]:
            print(f"  {db_id:30s} {acc:6.2f}% ({matched}/{total})")
    
    # Save detailed results if requested
    if output_file:
        output_data = {
            "summary": {
                "total": stats["total"],
                "successful": stats["successful"],
                "failed": stats["failed"],
                "matched": stats["matched"],
                "mismatched": stats["mismatched"],
                "accuracy": accuracy
            },
            "results": results
        }
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, indent=2, ensure_ascii=False)
        print(f"\nDetailed results saved to: {output_file}")
    
    return {
        "stats": stats,
        "accuracy": accuracy,
        "results": results
    }


def main():
    parser = argparse.ArgumentParser(
        description="Test RESDSQL Inference API against expected SQL/NatSQL queries"
    )
    parser.add_argument(
        "--api-url",
        default="http://localhost:8000/infer",
        help="API endpoint URL (default: http://localhost:8000/infer)"
    )
    parser.add_argument(
        "--data-file",
        default="NatSQL/NatSQLv1_6/train_spider.json",
        help="Path to test data JSON file (default: NatSQL/NatSQLv1_6/train_spider.json)"
    )
    parser.add_argument(
        "--target-type",
        choices=["sql", "natsql"],
        default="natsql",
        help="Target query type: sql or natsql (default: natsql)"
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Limit number of test cases to run (default: all)"
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print detailed results for each test case"
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.1,
        help="Delay between API calls in seconds (default: 0.1)"
    )
    parser.add_argument(
        "--output",
        type=str,
        default=None,
        help="Save detailed results to JSON file"
    )
    
    args = parser.parse_args()
    
    # Validate data file exists
    if not Path(args.data_file).exists():
        print(f"Error: Test data file not found: {args.data_file}")
        sys.exit(1)
    
    # Run tests
    try:
        run_tests(
            api_url=args.api_url,
            data_file=args.data_file,
            target_type=args.target_type,
            limit=args.limit,
            verbose=args.verbose,
            delay=args.delay,
            output_file=args.output
        )
    except KeyboardInterrupt:
        print("\n\nTest interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\nError running tests: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()