import os import random import json import time import logging from openai import OpenAI import sys sys.path.append('/Users/jway/Desktop/Dev/eda_service') import requests from requests.auth import HTTPDigestAuth from compare_sql import get_database_uri logging.basicConfig(level=logging.INFO) local_headers = { "x-auth-content": json.dumps({ "public_key": "S0T0RH50", # spider: 40T0RJ40, ask_billing: S0T0RH50 }), "X-ORG-ID": "10297887798990446300", "X-App-ID": "372820", } config_mapper = { # org_id, project_id, cluster_id, public_key, private_key, base_url, headers "local-v1": ("1372813089208161328", "1372813089278691291", "1379661944643734204", "public_key", "private_key", "http://127.0.0.1:11000/api/v1beta/apps/chat2query-jebzpXXe", local_headers), "local": ("1372813089208161328", "1372813089278691291", "1379661944643734204", "public_key", "private_key", "http://127.0.0.1:11000/inner/eda/api/v1beta/app/chat2query-jebzpXXe/endpoint", local_headers), "dev": ("30061", "3198952", "10297887798990446300", "G143NFT0", "0a35d8a4-cc21-4b95-b668-a0b087225415", "https://data.dev.tidbcloud.com/api/v1beta/app/chat2query-SijYVzXw/endpoint", {}), "staging": ("1369847559691367421", "10340145369517990849", "10340145369517990849", "C0QN38O0", "a48083d9-53ee-443a-9669-d380103554d4", "https://data.staging.tidbcloud.com/api/v1beta/app/chat2query-eognaXvH/endpoint", {}), } env = os.getenv("ENV", "local") c = config_mapper[env] org_id, project_id, cluster_id, public_key, private_key, base_url, headers = c[0], c[1], c[2], c[3], c[4], c[5], c[6] dataapi_credentials = [(public_key, private_key)] SQLEDITOR_BASE_URL = os.getenv("SQLEDITOR_BASE_URL", "https://us-west-2.nightly.shared.aws.tidbcloud.com/api/v1/eda/orgs/30061/projects/3198952") SQLEDITOR_AUTH_TOKEN = os.getenv("SQLEDITOR_AUTH_TOKEN", "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6Im1oSjk3OVJEQlFhOWU1U1I4ZUxwRCJ9.eyJodHRwczovL3RpZGJjbG91ZC5jb20vbG9naW5fcmVjZWl2ZWRfYXQiOjE3MTE5MzkxOTEyOTUsImxvZ2luX2hhc19vcmciOiIxIiwibG9naW5fb3JnX3R5cGUiOiJwZXJzb25hbCIsImxvZ2luX29yZ19pZCI6IjMwMDYxIiwibG9naW5fY29tcGFueV9uYW1lIjoiIiwiZW1haWwiOiJqaWFud2VpLndhbkBwaW5nY2FwLmNvbSIsImVtYWlsX3ZlcmlmaWVkIjp0cnVlLCJpc3MiOiJodHRwczovL2F1dGguZGV2LnRpZGJjbG91ZC5jb20vIiwiYXVkIjoiS2l3V2FqTEozcHRvOFNzWjdGMzVOYXFOaHdub1JoMWMiLCJpYXQiOjE3MTE5MzkxOTEsImV4cCI6MTcxMTk3NTE5MSwic3ViIjoiZ29vZ2xlLW9hdXRoMnwxMTMyODcxNzg0OTY0OTI3OTQ2MTkiLCJhdF9oYXNoIjoialVQeEZINWY3V3ZQMmhXNVFPLW5hQSIsInNpZCI6InlmOTMzSGdsTnRwTUM3cDNFY0ttMGg0ZHkyVHowTTgxIiwibm9uY2UiOiJSaXVxTmFCdXVvNXRVQ2FldEsyNnZGMmNVMGRJRFYxcCJ9.EN9d-f2RdKeOnVmhgzPwqrsD-CJEy5KqiD1iSJ9_hQmlmbylkNFNauzQMkniUFywCAYm0hYNGEVJdmjpWtws53DUX70efOIYyvQ53PMYVU21K0xJdhtlBMLXtZqprGAGOcNroa2XILuC4CDzOPC55dejydEdg1rQeBbJUqVSWJiN_5KzyguZItEEXTrrus2ucup2yFYd0sysYUEr2koianm-okpI8DnNRJmVczLyJixsoAsmUR8zQ1eXlejOKANhq4LH5ldWsoVnAdpxwRFmDMZGzAgEgVM6O7U_efrWRsXLS_gHLnN8gEqlN48tXVDnKsc0cDeLD_iXIFyffd8WKA") def get_auth_key(): choice = random.choice(dataapi_credentials) return HTTPDigestAuth(choice[0], choice[1]) def get_org_id_by_pubkey(): url = f"{base_url}/org_id" resp = requests.get(url, headers=headers, auth=get_auth_key()) logging.info("get_org_id_by_pubkey got resp %s, %s", resp.text, resp.status_code) return resp.json() def check_connection(): url = f"{base_url}/datacontext/connection_check" database_uri = get_database_uri("new_car_1") payload = { "database_uri": database_uri, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("check_connection got resp %s, %s", resp.text, resp.status_code) return resp.json() def create_data_context(dbname, sqlite=False): url = f"{base_url}/datacontext/context" if sqlite: dbname = dbname.removeprefix("new_") sql_dialect = "sqlite:///spider/database/{}/{}.sqlite".format(dbname, dbname) else: sql_dialect = "mysql+pymysql://2rrudfkjr8MQq3Q.root:9pS7AkerMB1Py3GE@gateway01.us-west-2.prod.aws.tidbcloud.com:4000/{}?ssl_ca%2Fetc%2Fssl%2Fcerts%2Fca-certificates.crt&ssl_verify_identity=true".format(dbname) payload = { "org_id": org_id, "creator": "jiajun", "data": {}, "database_uri": sql_dialect, "force": True, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("create data context %s got resp %s, %s", str(payload), resp.text, resp.status_code) return resp.json()["result"]["data_context_id"] def create_data_context_bigquery(): url = f"{base_url}/datacontext/context" sql_dialect = "bigquery://gcp-cp-bigquery-dev/billing?credentials_base64=ewogICJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsCiAgInByb2plY3RfaWQiOiAiZ2NwLWNwLWJpZ3F1ZXJ5LWRldiIsCiAgInByaXZhdGVfa2V5X2lkIjogImUzZGY3YmFlNmVlZTZiYmUwZTRkMzk2ZTAyNTg1Y2VmY2M5ZDcxNDgiLAogICJwcml2YXRlX2tleSI6ICItLS0tLUJFR0lOIFBSSVZBVEUgS0VZLS0tLS1cbk1JSUV2Z0lCQURBTkJna3Foa2lHOXcwQkFRRUZBQVNDQktnd2dnU2tBZ0VBQW9JQkFRQzBrdVdVd0JyVUxoNEtcbjdhRXE3TEZZdERRZ0J3cXRtQ1hHZEJsS2s4ZzIwVTN4QWN6MzV6RCsvNGxqV0YyeGhFb3hGTmxpalAxWHJHRW1cblV4dTNXa3FUZ3RONzdjcjdQcUQyc3l4cE5yWHUwQUE2TkhJZ2tIbWJQeE5PellMMU5pbndaRXg0S21uek9UU2tcbmRYMGFBOVI4MDBuZ083SEk4M0tra0UvR1ZVQ2pLbXdnSWxINS9acnQ4ZVVPVDlxZHdlcSs1UktZN0hCWUd2UllcbnBBRXJpWWEyUHM2U1dhMlFvRGx2ZjNZNlVyUXFIcGFzOFcyYlM3Y3lFVTM3VHVVTytUNlU0c1RyMGI0R1R3bmtcbmhxZTZWSFpiOUZHS09OTGxBR2g0bmJpakNJUCtrd2xkTVo1Y1dUQlF1RmRsbzJGK0VqZXlrNUpuU2Z4dkJMSGJcbkxBZnNnTmM5QWdNQkFBRUNnZ0VBQjl5T1lMM0ZYVC85d0ZtTXB2OERrSnBWMjZBYTI2aHdXcFhDNVIrT2U2NThcbm4zWEYyL1FsaXJDSzVYZWJLcWVaQ0U2QlVWeFVpa3pmdE9pdzMvaHZIMEYybmlKaDhETFBsTFlEeExma1QxZWdcbmJ1aEtJTGRSaUF3RXRlMVdHb0ZuYlljY1dpMzBpMm1ON0d3T0RJUW1xK0dTQ1pCbnd0a3dnRkFsSFB3b3NvSExcblZ5WmZEbFZ2cEp2d0pMdzN5UEtEdVlZeTBVL05VekhDdUxJMjk2bjlkU0U2SXg3SWlJVVA5c0FVaGpSM0luL3NcbjZKY21aZFgwNElmeENVR2N5eXlvUkRlWUhaaVd4c1JxbWdNUnR3UVNzR0RiQkN3NmV5b0FWRmJEcXRTcC83QTRcbm51QklxeW91T0g1TnRadkwrUWZoeFZUZXR4TDgxaDdoV1g1a3pwZ0ErUUtCZ1FEWFBySFhoVHRLZzFnS0Z0Rmlcbm4vTTNiYUthZUJTbWhBY2tsclJ4VFFSMjJWTHdqRGdwTjFuQVlRL09QcERFTnlOYmxwM09mNkdVZjlqL2xpd3FcbkZUYkpaVmFDVUJ2RSs0bnBxMWRUcGR3S09uYlhzbGpvdE9aVzRLWVNDV3pZNWRKR3JSUW9VbFFYbXpVOUIyU29cblRWcjNSRTZvUHEwV3dDaFF4YW5qY2wxam1RS0JnUURXdzZIeEVxdkFSZ2dyS2JTL3RhVFVZR2c3N0hxd2RiUlBcblRUWit3L3o4dmNzcGVMeGRJcVgyRkNicFkwYUo2bWhSb09XOE9XWHRjbmRVZ0Rmc090WEdtUFlyZk5tM1ZWdFZcbktuby9vMW1FWWZPYTBFc3A5N2FGd045UUpPVUNWNmFsWW9sTm43WnJub1RFWUpQRjFZeHpGUXpiUlE4ZkkvT05cbmZFTnNEZHhYUlFLQmdHamtPNFE4ZisxSzNsOWtTVHRYbVEyYzFzMlIwUFpvaHBhVS9tL3NTVTJsa0Y1NjVtOW5cbmNzeG1RRzZRWExuUjZPQWNJT0RkbjU4ZlFsOGpEbmlmaCtoVEF2cFVFWVJIMUNoVUFDVmFRVVdXd2cxUTh6VExcbjJ5ZmhKK3pJUFF2aHl6djM0a2ZYL3ZUV0JQYXRvS2JIczVtbDV2TGtRRE5EQnlKOWdORy9ETFdaQW9HQkFMZllcbmhMckFCdkx6Z0c2aXB5NHRxc3A3TlArSE1vZjRXc0JBRXRSMXBnNGQwdHJQbGtBTU54SHVzbGk4T1ZnNnA4Nm9cbjg4RXJTWXY0bnhwaElINUM1MW00KzVMQ3NBVjNObFl6cEs3eGdPSUlqK2NKUkM3SzlSNk91eHY2V1dpbDYreExcbjY4bmRYT1hIcjFRRWNMdThVUno2K2lTQ2U2MW1zMGw5QktKdFF5QVpBb0dCQUtFRmFESUUvVnJSZkYrUzc5Wm9cbmFQTElkaSsyM3N3Q29ub0ZLUWh4U3p4Yk1Zbzg5S0U2YWl0STE1YXhDa29SS2JYLzRlL1N3V3luanNWTXlxWVJcbmI1alNKcGlIU05rTC93SlNCa1RLZ2tFaFFMbjVzejRnejNxWE9BS0hNMVlxK0cwOHg3U3hJakNvT0ZTTDA2Vm5cbjBaY0lyNGJST3pGbzdXOTkzUmFUQlJHdFxuLS0tLS1FTkQgUFJJVkFURSBLRVktLS0tLVxuIiwKICAiY2xpZW50X2VtYWlsIjogImVkYXNlcnZpY2VAZ2NwLWNwLWJpZ3F1ZXJ5LWRldi5pYW0uZ3NlcnZpY2VhY2NvdW50LmNvbSIsCiAgImNsaWVudF9pZCI6ICIxMDg4NTQxNzMxMTEzNDk0ODQ4NjQiLAogICJhdXRoX3VyaSI6ICJodHRwczovL2FjY291bnRzLmdvb2dsZS5jb20vby9vYXV0aDIvYXV0aCIsCiAgInRva2VuX3VyaSI6ICJodHRwczovL29hdXRoMi5nb29nbGVhcGlzLmNvbS90b2tlbiIsCiAgImF1dGhfcHJvdmlkZXJfeDUwOV9jZXJ0X3VybCI6ICJodHRwczovL3d3dy5nb29nbGVhcGlzLmNvbS9vYXV0aDIvdjEvY2VydHMiLAogICJjbGllbnRfeDUwOV9jZXJ0X3VybCI6ICJodHRwczovL3d3dy5nb29nbGVhcGlzLmNvbS9yb2JvdC92MS9tZXRhZGF0YS94NTA5L2VkYXNlcnZpY2UlNDBnY3AtY3AtYmlncXVlcnktZGV2LmlhbS5nc2VydmljZWFjY291bnQuY29tIiwKICAidW5pdmVyc2VfZG9tYWluIjogImdvb2dsZWFwaXMuY29tIgp9Cg==" payload = { "org_id": org_id, "creator": "jiajun", "data": {}, "database_uri": sql_dialect, "force": True, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("create data context %s got resp %s, %s", str(payload), resp.text, resp.status_code) return resp.json()["result"]["data_context_id"] def query_job_detail(job_id): url = f"{base_url}/jobs/{job_id}" resp = requests.get(url, headers=headers, auth=get_auth_key()) logging.info("get job %s detail got resp %s", job_id, resp.text) return resp.json()["result"] def create_session_context(data_context_id): url = f"{base_url}/session/context" payload = { "org_id": org_id, "creator": "jiajun", "data": {}, "data_context_id": data_context_id, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("create session_context got %s", resp.text) return resp.json()["result"]["session_context_id"] def question_breakdown(session_context_id, raw_question): url = f"{base_url}/session/{session_context_id}/actions/question_breakdown" payload = { "raw_question": raw_question, "auto_run": True } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("create question_breakdown got %s", resp.text) return resp.json()["result"]["job_id"] def text2sql(session_context_id, task_id): url = f"{base_url}/orgs/{org_id}/text2sql" payload = { "session_context_id": session_context_id, "task_id": task_id, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("create text2sql got %s", resp.text) return resp.json()["result"]["job_id"] def generate_db_schema(data_context_id: int): url = f"{base_url}/datacontext/{data_context_id}/actions/db_exploring" resp = requests.post(url, headers=headers, auth=get_auth_key()) logging.info("generate db schema got resp %s", resp.text) return resp.json()["result"]["job_id"] def query_ai_for_sql(data_summary_id, raw_question, evidence=""): job_id = chat2data(data_summary_id, raw_question, evidence) logging.info("job_id is %s, data_summary_id %s, raw_question %s", job_id, data_summary_id, raw_question) count = 0 while True: try: result = query_job_detail_v2(job_id) count += 1 if result["status"] == "done": inner_result = result["result"]["task_tree"] for k, v in inner_result.items(): if "sql" in v: return v["sql"], v["description"], v["clarified_task"], v.get("raw_generated_sql", ""), v.get("refine_note", "") # noqa logging.info("sql not found, detail: %s", result) return "sql not found", "", "", "", "" if result["status"] == "failed": logging.info("job failed, detail: %s", result) return "job failed", "", "", "", "" if count > 300: logging.info("job failed because of timeout, detail: %s", result) return "job failed", "", "", "", "" except Exception as e: logging.exception("failed to query job detail: %s", str(e)) finally: time.sleep(5) def suggest_questions(session_context_id): url = f"{base_url}/session/{session_context_id}/actions/question_suggestions" resp = requests.post(url, headers=headers, auth=get_auth_key()) logging.info("create suggest_questions got %s", resp.text) return resp.json()["result"]["job_id"] def loop_job(job_id): for i in range(10000): try: result = query_job_detail(job_id) if result["status"] in ("done", "failed"): break except Exception: pass finally: time.sleep(1) def update_like_or_not(session_context_id, task_id, question_id, like_or_not): url = f"{base_url}/session/{session_context_id}/like_or_not" payload = { "task_id": task_id, "question_id": question_id, "like_or_not": like_or_not, } resp = requests.put(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("update_like_or_not got %s", resp.text) return resp.json() def update_task_sql(session_context_id, task_id, question_id, sql): url = f"{base_url}/session/{session_context_id}/sql" payload = { "task_id": task_id, "question_id": question_id, "sql": sql, } resp = requests.put(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("update_task_sql got %s", resp.text) return resp.json() def execute_task_sql(session_context_id, task_id, question_id): url = f"{base_url}/session/{session_context_id}/sql_execution" payload = { "task_id": task_id, "question_id": question_id, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("execute_task_sql got %s", resp.text) return resp.json() def update_table_description(data_context_id, table, description): url = f"{base_url}/datacontext/{data_context_id}/table/description" payload = { "table": table, "description": description, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("update_table_description got %s", resp.text) return resp.json() def update_column_description(data_context_id, table, column, description): url = f"{base_url}/datacontext/{data_context_id}/column/description" payload = { "table": table, "column": column, "description": description, } resp = requests.post(url, headers=headers, json=payload, auth=get_auth_key()) logging.info("update_table_description got %s", resp.text) return resp.json() #################### V2 def create_data_summary_by_cluster(cluster_id: int, dbname: str): url = f"{base_url}/v2/dataSummaries" payload = { "cluster_id": cluster_id, "database": dbname, } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("create_data_summary %s got %s", payload, resp.text) return resp.json() def create_data_summary(dbname, sqlite=False): url = f"{base_url}/v2/dataSummaries" payload = { "database_uri": get_database_uri(dbname, sqlite), "creator": "jiajun" } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("create_data_summary %s got %s", payload, resp.text) return resp.json() def get_data_summary(data_summary_id: int): url = f"{base_url}/v2/dataSummaries/{data_summary_id}" resp = requests.get(url, auth=get_auth_key(), headers=headers) logging.info("get data summary of %s got %s", data_summary_id, resp.text) return resp.json() def update_data_summary(data_summary_id: int): url = f"{base_url}/v2/dataSummaries/{data_summary_id}" payload = { "tables": { "car_makers": { "description": "I_CAN_UPDATE_TABLE_DESCRIPTION", "columns": { "country": { "description": "I_CAN_UPDATE_COLUMN_DESCRIPTION_TOO", }, } } } } resp = requests.put(url, auth=get_auth_key(), headers=headers, json=payload) logging.info("update data summary of %s got %s", data_summary_id, resp.text) return resp.json() def generate_suggest_questions(data_summary_id: int): url = f"{base_url}/v3/suggestQuestions" payload = { "cluster_id": "10297887798990446300", "database": "billing" } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("generate suggest questions for data summary %s got %s", data_summary_id, resp.text) return resp.json() def create_assistant_session(): url = f"{base_url}/v3/sessions" payload = { "cluster_id": "10297887798990446300", "database": "toxicology" } # payload = { # "cluster_id": "10297887798990446300", # "database": "billing" # } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("chat2data with %s got %s", payload, resp.text) return resp.json()["result"]["session_context_id"] def get_assistant_session(session_id: int): url = f"{base_url}/v2/sessions/{session_id}" resp = requests.get(url, auth=get_auth_key(), headers=headers) logging.info("get session messages successed, session_id = %s, result = %s", str(session_id), resp.text) return resp.json() def reset_assistant_session(session_id: int): url = f"{base_url}/v2/sessions/{session_id}/reset" resp = requests.put(url, auth=get_auth_key(), headers=headers) logging.info("reset session id successed, session_id = %s, result = %s", str(session_id), resp.text) return resp.json() def assistant_chat(session_context_id, raw_question, evidence="", query_id="", task_id=""): url = f"{base_url}/v3/sessions/{session_context_id}/chat2data" payload = { "question": raw_question, "feedback_answer_id": query_id, "feedback_task_id": task_id, "evidence": evidence } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("assistant_chat with %s got %s", payload, resp.text) return resp.json()["result"]["job_id"] def chat2data(data_summary_id: int, question: str, evidence=""): url = f"{base_url}/v2/chat2data" payload = { "data_summary_id": data_summary_id, "raw_question": question, "evidence": evidence, } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("chat2data with %s got %s", payload, resp.text) return resp.json()["result"]["job_id"] def chat2sql(data_summary_id: int, question: str): url = f"{base_url}/v2/chat2sql" payload = { "data_summary_id": data_summary_id, "raw_question": question } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("chat2sql with %s got %s", payload, resp.text) return resp.json()["result"]["job_id"] def execute_sql(session_context_id, task_id, question_id): url = f"{base_url}/v2/executeSQL" payload = { session_context_id: session_context_id, task_id: task_id, question_id: question_id, } resp = requests.post(url, json=payload, auth=get_auth_key(), headers=headers) logging.info("executeSQL with %s got %s", payload, resp.text) return resp.json() def query_job_detail_v2(job_id): url = f"{base_url}/v2/jobs/{job_id}" resp = requests.get(url, headers=headers, auth=get_auth_key()) logging.info("get job %s detail got resp %s", job_id, resp.text) return resp.json()["result"] def loopv2(job_id: str): for i in range(10000): try: result = query_job_detail_v2(job_id) if result["status"] in ("done", "failed"): break except Exception as e: logging.exception("failed to query job detail: %s", e) finally: time.sleep(5) def list_assistant_messages(): client = OpenAI(api_key="[REDACTED_OPENAI_API_KEY]") try: thread_messages = client.beta.threads.messages.list("thread_0rAv2tGRZbd9u9v87eRXRlG3") print(thread_messages) except Exception as e: print(e) def sqleditor_create_session(cluster_id, database, session_context_id=0) -> dict: url = f"{SQLEDITOR_BASE_URL}/sessions" headers = { "authorization": f"Bearer {SQLEDITOR_AUTH_TOKEN}", } payload = { "cluster_id": f"{cluster_id}", "database": f"{database}", "session_context_id": session_context_id or 0, } resp = requests.post(url, json=payload, headers=headers) logging.info("sqleditor_create_session got %s", resp.text) return resp.json()["result"] def sqleditor_create_query(session_id, raw_question) -> dict: url = f"{SQLEDITOR_BASE_URL}/multiRoundChat2sql" headers = { "authorization": f"Bearer {SQLEDITOR_AUTH_TOKEN}", } payload = { "session_context_id": f"{session_id}", "question": f"{raw_question}", } resp = requests.post(url, json=payload, headers=headers) logging.info("sqleditor_create_query got %s", resp.text) return resp.json()["result"] def sqleditor_create_feedback(session_id, query_id, feedback) -> dict: url = f"{SQLEDITOR_BASE_URL}/multiRoundFeedback" headers = { "authorization": f"Bearer {SQLEDITOR_AUTH_TOKEN}", } payload = { "session_context_id": f"{session_id}", "query_id": f"{query_id}", "feedback": f"{feedback}", } resp = requests.post(url, json=payload, headers=headers) logging.info("sqleditor_create_feedback got %s", resp.text) return resp.json()["result"] def sqleditor_query_job_detail(job_id) -> dict: url = f"{SQLEDITOR_BASE_URL}/jobs/{job_id}" headers = { "authorization": f"Bearer {SQLEDITOR_AUTH_TOKEN}", } resp = requests.get(url, headers=headers) logging.info("sqleditor_query_job_detail got %s", resp.text) return resp.json()["result"] def loop_sqleditor_job(job_id): for i in range(10000): try: resp_json = sqleditor_query_job_detail(job_id) if resp_json["status"] in ("done", "failed"): return resp_json except Exception as e: logging.exception("failed to query job detail: %s", e) finally: time.sleep(5) if __name__ == "__main__": # data_summary_resp = create_data_summary("game", False)["result"] # job_id = data_summary_resp["job_id"] # data_summary_id = data_summary_resp["data_summary_id"] # loopv2(job_id) # print(data_summary_id) # data_summary_id = 206651 # get_data_summary(data_summary_id) job_id = generate_suggest_questions(206651)["result"]["job_id"] loopv2(job_id) # get_suggest_questions(data_summary_id) # question = "每天费用统计" # job_id = chat2data(206651, question) # loopv2(job_id) # data_context_id = create_data_context("lengjing") # print("data_context_id = {}".format(data_context_id)) # job_id = generate_db_schema(data_context_id) # for i in range(10000): # result = query_job_detail(job_id) # if result["status"] in ("done", "failed"): # break # time.sleep(5) # 90301: lengjing # 90306: baseball # billing-local: 206651 # billing-dev: 485669 # session_id = create_assistant_session() # print(session_id) # case-1: thread_xpTD7um3Yh3ts2OJZRYGaljt prompt 分为: message + instruction # list_assistant_messages() # raw_question = "tell me how mant atoms" # query_id = "" # task_id = "" # job_id = assistant_chat(328903, raw_question, "", query_id, task_id) # for i in range(200): # result = query_job_detail_v2(job_id) # if result["status"] in ("done", "failed"): # break # time.sleep(1) # get_assistant_session(486515) # reset_assistant_session(486515) # bird: toxicology， app_id": 372821 # session_context_id = create_session_context(328648) # print(session_context_id) # raw_question = "Among all chemical compounds identified in the database what percent of compounds form a triple-bond." # job_id = question_breakdown(session_context_id, raw_question) # print(job_id) # loop_job(job_id) # loop_job(suggest_questions(328723)) # session_id = 486856 # resp_json = sqleditor_create_session(10371474764223291851, "bolife_db", session_context_id=486870) # print(resp_json) # session_context_id = resp_json["session_context_id"] # question = "how many valid cards" # resp_json = loop_sqleditor_job(sqleditor_create_query(486856, question)["job_id"]) # print(resp_json) # resp_json = sqleditor_create_feedback(486860, "869a6ec8-243a-4476-8c18-bbd978afeabf", "valid's value only can be 'YES' or 'NO'") # print(loop_sqleditor_job(resp_json["job_id"])) # list_assistant_messages()