import time import os import sqlite3 import glob import csv import logging import json from collections import Counter from concurrent.futures import ProcessPoolExecutor from typing import Tuple import requests from tenacity import before_sleep_log, retry, stop_after_attempt, wait_random_exponential from test_context import create_data_summary, query_ai_for_sql, loopv2 from hardlevel import eval_hardness from compare_sql import compare_row logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) cookies = "" def bird_execute_sql(predicted_sql,ground_truth, db_path): predicted_sql = predicted_sql.replace("\n", " ").replace("\t", " ").replace(";", "") ground_truth = ground_truth.replace("\n", " ").replace("\t", " ").replace(";", "") try: conn = sqlite3.connect(db_path) # Connect to the database cursor = conn.cursor() cursor.execute(predicted_sql) predicted_res = cursor.fetchall() cursor.execute(ground_truth) ground_truth_res = cursor.fetchall() res = 0 if set(predicted_res) == set(ground_truth_res): res = 1 return res except Exception as e: return str(e) def create_test_plan(): payload = { "type": "SQL Generator", "creator": "Jiajun", "name": "SpiderCases", "description": "", "ai_model_id": 1 } headers = { "Cookie": cookies, } url = "https://aitestsuit.oauth.prod.aws.tidbcloud.com/api/v1/plan" resp = requests.post(url, headers=headers, json=payload) logging.info("create test plan got %s", resp.text) def delete_test_case(): case_id = 150001 headers = { "Cookie": cookies, } url = f"https://aitestsuit.oauth.prod.aws.tidbcloud.com/api/v1/case/{case_id}" resp = requests.delete(url, headers=headers) logging.info("delete test case got %s", resp.text) def create_test_case(plan_id, question, database, table, sql, column_names): payload = { "plan_id": plan_id, "question": question, "database": database, "table": table, "sql": sql, "column_names": column_names, "compare_args": {}, } headers = { "Cookie": cookies, } url = "https://aitestsuit.oauth.prod.aws.tidbcloud.com/api/v1/case" resp = requests.post(url, headers=headers, json=payload) logging.info("create test case got %s", resp.text) def parse_cases(csvwriter, data_summary_id, filename, dbname): """ {'db_id': 'concert_singer', 'query': 'SELECT count(*) FROM singer', 'query_toks': ['SELECT', 'count', '(', '*', ')', 'FROM', 'singer'], 'query_toks_no_value': ['select', 'count', '(', '*', ')', 'from', 'singer'], 'question': 'How many singers do we have?', 'question_toks': ['How', 'many', 'singers', 'do', 'we', 'have', '?'], 'sql': {'from': {'table_units': [['table_unit', 1]], 'conds': []}, 'select': [False, [[3, [0, [0, 0, False], None]]]], 'where': [], 'groupBy': [], 'having': [], 'orderBy': [], 'limit': None, 'intersect': None, 'union': None, 'except': None}} """ count = 0 # with open("./train_spider.json") as f: with open(filename) as f: result = json.loads(f.read()) for i in result: database = i["db_id"] is_bird_eval = i.get("SQL") is not None sql = i.get("query") or i.get("SQL") question = i["question"] evidence = i.get("evidence", "") hardness = i.get("difficulty") if is_bird_eval else eval_hardness(i["sql"]) note = "" if database not in [dbname, dbname.removeprefix("new_")]: continue while True: try: time_before = int(time.time()) generated_sql, description, clarified_task, raw_generated_sql, refine_note = query_ai_for_sql(data_summary_id, question, evidence) # noqa time_after = int(time.time()) if not generated_sql: logging.warning( "no sql found, data_summary_id %s, std spider args %s, results %s, %s, %s, %s, %s", data_summary_id, i, generated_sql, description, clarified_task, raw_generated_sql, refine_note, # noqa ) generated_sql = "sql not found" equal = False sql = sql.replace("\t", " ").replace("\n", " ").rstrip(";") + ";" generated_sql = generated_sql.replace("\t", " ").replace("\n", " ").rstrip(";") + ";" count += 1 question = question.replace("\t", " ").replace("\n", " ").replace(",", " ") bird_eval = bird_execute_sql(generated_sql, sql, "./spider/bird/dev_databases/{}/{}.sqlite".format(dbname, dbname)) if is_bird_eval else "N/A" row = [database, question, sql, generated_sql, hardness, equal, note, bird_eval] if os.getenv("DEBUG"): row += [time_after-time_before, clarified_task, raw_generated_sql, refine_note] csvwriter.writerow(row) logging.info("writerow %s", row) # create_test_case(plan_id, question, database, "", sql, []) logging.info("=================== count =%s====================", count) break except Exception as e: logging.exception("failed to query_ai_for_sql: %s, retry after 5 seconds", e) time.sleep(5) @retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(10), before_sleep=before_sleep_log(logger, logging.INFO)) def setup_context(dbname, sqlite=False): resp_json = create_data_summary(dbname, sqlite) result = resp_json["result"] return result["data_summary_id"], result["job_id"] def analysis(csv_filename): c = Counter() human_check_count = 0 with open(csv_filename) as f: for i in csv.DictReader(f): difficulty = i["difficulty"] equal = i["equal"] if equal == "human check": human_check_count += 1 c[difficulty + "human"] += 1 c["total" + "human"] += 1 elif equal == "True" or equal == "TRUE": c[difficulty + "succeed"] += 1 c["total" + "succeed"] += 1 elif equal == "raw query failed": c[difficulty + "raw_query_failed"] += 1 c["total" + "raw_query_failed"] += 1 elif i["generated_sql"] == "sql not found" or equal == "sql not found" or i["generated_sql"] == "job failed": c[difficulty + "sql_not_found"] += 1 c["total" + "sql_not_found"] += 1 elif equal == "run sql failed": c[difficulty + "result_failed"] += 1 c["total" + "result_failed"] += 1 else: c[difficulty + "failed"] += 1 c["total" + "failed"] += 1 c[difficulty] += 1 c["total"] += 1 difficulties = ["easy", "medium", "hard", "extra", "total"] print("difficulty\ttotal\tfailed\traw_query_failed\trun_sql_failed\tsql_not_found\thuman_check\tsucceed\tsuccess_rate\twith_human_check_success_rate") for i in difficulties: total = c[i] failed = c[i + "failed"] raw_query_failed = c[i + "raw_query_failed"] result_failed = c[i + "result_failed"] sql_not_found = c[i + "sql_not_found"] human_check = c[i + "human"] succeed = c[i + "succeed"] success_ratio, human_ratio = 0, 0 if total: success_ratio = (succeed / total) * 100 human_ratio = ((human_check + succeed)/total) * 100 result = i, total, failed, raw_query_failed, result_failed, sql_not_found, human_check, succeed, success_ratio, human_ratio print("{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{:.2f}%\t{:.2f}%".format(*result)) yield result def sort_csv(csv_filename, data_context_id): results = [] with open(csv_filename) as f: for i in csv.DictReader(f): results.append(i) results.sort(key=lambda i: (i["difficulty"], i["question"])) with open(csv_filename, 'w', newline='') as csvfile: csv_header = ["database", "question", "sql", "generated_sql", "difficulty", "equal", "note"] if os.getenv("DEBUG"): csv_header += ["time_cost", "clarified_task", "raw_generated_sql", "refine_note", "data_context_id"] spamwriter = csv.writer(csvfile) spamwriter.writerow(csv_header) for i in results: database = i.get("\ufeffdatabase") or i.get("database") # BOM... row = [database, i["question"], i["sql"], i["generated_sql"], i["difficulty"], i["equal"], i["note"]] if os.getenv("DEBUG"): row += [i["time_cost"], i["clarified_task"], i["raw_generated_sql"], i["refine_note"], data_context_id] spamwriter.writerow(row) def filter_all_sqlites_has_data(): sql = "SELECT name FROM sqlite_master WHERE type='table'" for i in glob.glob("./spider/database/**/*.sqlite"): con = sqlite3.connect(i) cur = con.cursor() res = cur.execute(sql) has_data = False for j in res.fetchall(): count_sql = "SELECT COUNT(*) FROM {}".format(j[0]) result = cur.execute(count_sql).fetchone() if result[0]: has_data = True if has_data: yield "new_" + i.split("/")[-2] def run_db_cases(dbname): cases_filename = "./spider/database/combine.json" csv_filename = "cases_{}_st_sqlite.csv".format(dbname) data_context_id, job_id = setup_context(dbname, True) loopv2(job_id) with open(csv_filename, 'w', newline='') as csvfile: csv_header = ["database", "question", "sql", "generated_sql", "difficulty", "equal", "note", "bird_eval"] if os.getenv("DEBUG"): csv_header += ["time_cost", "clarified_task", "raw_generated_sql", "refine_note"] csvwriter = csv.writer(csvfile) csvwriter.writerow(csv_header) parse_cases(csvwriter, data_context_id, cases_filename, dbname) # sort_csv(csv_filename, data_context_id) # analysis(csv_filename) def run_all_cases(): with ProcessPoolExecutor(max_workers=4) as pool: for i in filter_all_sqlites_has_data(): logging.info("run_db_cases %s", i) pool.submit(run_db_cases, i) def run_cases_in_dbs(dbnames): max_workers = int(os.getenv("MAX_WORKERS", 4)) with ProcessPoolExecutor(max_workers=max_workers) as pool: for i in dbnames: pool.submit(run_db_cases, i) def count_questions(): dbnames = set() for i in filter_all_sqlites_has_data(): dbnames.add(i) total = 0 with open("./combine.json") as f: dev_result = json.loads(f.read()) for i in dev_result: if i["db_id"] in dbnames: total += 1 print(total) if __name__ == "__main__": logging.basicConfig(level=logging.DEBUG) # run_all_cases() os.environ["DEBUG"] = "true" run_all_cases() # # # session_context_id = 60023 # dbnames = ["new_car_1", "new_baseball_1", "new_allergy_1", "new_activity_1", "new_singer"] # dbnames = [ # "new_car_1", "new_baseball_1", "new_allergy_1", "new_activity_1", "new_music_1", # "new_flight_1", "new_movie_1", "new_apartment_rentals", # ] # run_db_cases("new_car_1") # dbnames = ["new_baseball_1"] # run_cases_in_dbs(dbnames) # dbnames = ["new_flight_1", "new_movie_1"] # run_cases_in_dbs(dbnames) # run_db_cases("new_baseball_1") # analysis("./cases_new_baseball_1_st_sqlite.csv")