dvwn/nl2sql-api

1# Test the Hugging Face inference2from backend.src.nl2sql.hf_engine import generate_sql3from backend.src.database.db_manager import get_db_connection, get_schema_context4import pandas as pd5 6def test_single_query():7    print("Initializing Featherless AI SQL generation test...")8    # Fetch the database schema context (ddl) from Chinook9    ddl = get_schema_context10    question = "Identify the artist who has earned the most revenue from customers in Canada."11 12    try:13        generated_sql = generate_sql(question, ddl)14        print(f"\nGenerated SQL:\n{generated_sql}\n")15 16        # Connect to the database and execute the generated SQL17        connection = get_db_connection()18        df = pd.read_sql_query(generated_sql, connection)19        connection.close()20 21        print("\nDatabase Query Result:")22        print(df)23        print("\nTest completed successfully: API connected and SQL is valid.")24 25    except Exception as e:26        print(f"\nTest failed: {e}")27 28if __name__ == "__main__":29    test_single_query()