import time import json import copy from collections import deque import re import random import math import os from tenacity import before_sleep_log, retry, retry_if_not_exception_type, stop_after_attempt, wait_random_exponential from consts.consts import LANGUAGE_EN, LANGUAGE_CN, GPT_DEFAULT, GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI, Claude3, Qwen, MODE_TYPE_SPIDER, MODE_TYPE_ASK_BILLING, MODE_TYPE_BIRD, GenSQLErrorMsg, MODE_TYPE_LENGJING, KnowledgeType, ModuleType, KnowledgeSourceType, VectorMode, AssistantName, TaskType, AgentType, EdaTaskJobStatus, SUPPORT_TIDB_SYNTAX_WHITELIST_ORGS, AssignedAppID, SuggestQuestionLevel from models.user_preference import UserPreference from svc.dataapi import get_sql_executor from utils.local import get_local_metrics, get_tl_app_id_pubkey, get_tl_trace_id, get_tl_user_token from utils.credential import get_cluster_info from svc.openai import ai_completion, ai_assistant_chat, ai_assistant_delete_message from utils.sql import has_ddl_sql_permission, is_ddl_sql from utils.database import is_sql_error from errors.errors import AlgorithmException, BadRequestCausedAlgorithmException, PermissionDeniedException, BadRequestException from .context import context import uuid from domain_info.lengjing_v2 import get_domain_term_description, get_domain_db_summary, get_domain_entity from domain_info.index import get_kg_management, get_local_suggest_questions, get_local_question_result from utils.logging import logging logger = logging.getLogger(__name__) class EDATask: def __init__(self, session_context_id, binded_thread_id=""): self.session_context_id = session_context_id self.binded_thread_id = binded_thread_id session_context = context.load_session_context(session_context_id) self.tasks_to_be_resolved = deque(session_context.get("tasks_to_be_resolved") or []) # used for new version self.tasks_tree = {} def _determine_model_by_cloud_provider(self, ai_model): """ Automatically select AI model based on database cluster cloud platform or specific cluster ID Decision priority: 1. Feature flag: QWEN_ENABLED=false/unset (default) → Disable Qwen, keep original model 2. Alibaba Cloud clusters → Force use Qwen.DEFAULT (current: qwen-plus, 128K context) 3. Specific cluster_id (10072363520111170484) → Force use Qwen.DEFAULT 4. Other cases → Keep original ai_model Note: Must explicitly set QWEN_ENABLED=true to enable Qwen routing (disabled by default) Args: ai_model: Original AI model (from user preference or default) Returns: Determined AI model """ try: # ========== 1. Check feature flag ========== # Default false: Must explicitly set QWEN_ENABLED=true to enable Qwen routing qwen_enabled = os.getenv("QWEN_ENABLED", "false").lower() in ("true", "1", "yes") logging.info( "session_id = %s, qwen_enabled from env = %s", self.session_context_id, qwen_enabled ) if not qwen_enabled: logging.info( "session_id = %s, Qwen is disabled by QWEN_ENABLED flag, keeping original model = %s", self.session_context_id, ai_model ) return ai_model # ========== 2. Get cluster_id ========== db_context = context.load_dataset_context() cluster_id = db_context.get("cluster_id") if not cluster_id: logging.info( "session_id = %s, cluster_id not found in context, use original model = %s", self.session_context_id, ai_model ) return ai_model # ========== 3. Get cloud platform information ========== cluster_info = get_cluster_info(cluster_id) cloud_provider = cluster_info.get("cluster", {}).get("cloud_provider_name", "").lower() logging.info( "session_id = %s, cluster_id = %s, cloud_provider = %s, original_model = %s, qwen_enabled = %s", self.session_context_id, cluster_id, cloud_provider, ai_model, qwen_enabled ) # ========== 4. Check if Alibaba Cloud (highest priority) ========== alibaba_cloud_names = ["alibaba", "aliyun", "alicloud"] is_alibaba_cloud = cloud_provider in alibaba_cloud_names if is_alibaba_cloud: logging.info( "session_id = %s, detected Alibaba Cloud, switching to Qwen model, " "cluster_id = %s, cloud_provider = %s, original_model = %s, new_model = %s", self.session_context_id, cluster_id, cloud_provider, ai_model, Qwen.DEFAULT ) return Qwen.DEFAULT # ========== 5. Check specific cluster_id ========== qwen_cluster_ids = ["10072363520111170484"] cluster_id_str = str(cluster_id) is_specific_cluster = cluster_id_str in qwen_cluster_ids if is_specific_cluster: logging.info( "session_id = %s, detected specific cluster_id, switching to Qwen model, " "cluster_id = %s, cloud_provider = %s, original_model = %s, new_model = %s", self.session_context_id, cluster_id, cloud_provider, ai_model, Qwen.DEFAULT ) return Qwen.DEFAULT # ========== 6. Other cases: keep original model ========== logging.info( "session_id = %s, cluster is not eligible for Qwen routing, " "cluster_id = %s, cloud_provider = %s, keeping original model = %s", self.session_context_id, cluster_id, cloud_provider, ai_model ) return ai_model except Exception as e: # If cloud platform detection fails, log error but keep original model logging.warning( "session_id = %s, failed to determine model routing, error = %s, " "keeping original model = %s", self.session_context_id, str(e), ai_model ) return ai_model def task_suggest_questions(self, question_num=5): if question_num <= 0: return [] if question_num > 15: raise Exception("The number of recommended questions generated cannot exceed 15.") db_context = context.load_dataset_context() # Special For Ask Billing preference = UserPreference.load(db_context.get("app_id", 0)) if preference.mode_type == MODE_TYPE_ASK_BILLING: return get_local_suggest_questions(MODE_TYPE_ASK_BILLING, question_num, session_id=self.session_context_id) # Retrieval local suggest questions first (cluster_id + db_name) local_questions = get_local_suggest_questions(str(db_context.get("cluster_id"))+"_"+db_context.get("db_name"), question_num, session_id=self.session_context_id) if len(local_questions) != 0: return local_questions return self.agent_suggest_questions(db_context, question_num, preference.llm_model, preference.language_code, preference.suggest_question_level) def agent_suggest_questions(self, db_context, question_num, ai_model=GPT_DEFAULT, language=LANGUAGE_EN, question_level=SuggestQuestionLevel.SIMPLE): time_before = int(time.time()) db_context = context.load_dataset_context() entity_schema = [] # if data context is not ready, use db schema directly if db_context.get("status") != "done": # format the template table_template = "Table: `%(table_name)s`, All columns are as follows:\n%(columns_str)s" column_template = "- column: `%(name)s`, type: %(type)s" for table in db_context.get("db_schema", {}).keys(): columns_list = [column_template%column for column in db_context["db_schema"].get(table, {}).get("columns", {}).values()] column_str = "\n".join(columns_list) entity_schema.append(table_template%{"columns_str": column_str, **db_context["db_schema"].get(table, {})}) else: for entity in db_context.get("entity", {}).values(): schema = "**entity_name**: %(name)s, **description**: %(summary)s, **key_attributes**: %(attributes)s, **involved_tables**: %(involved_tables)s"%{ "name": entity["name"], "summary": entity["summary"], "attributes": json.dumps(entity["attributes"]), "involved_tables": ",".join(["{'table': %(table)s, 'description': %(description)s}"%{"table": item, "description": db_context["db_schema"].get(item, {}).get("description", "")} for item in entity["involved_tables"]]) } entity_schema.append(schema) step_instruction = "2. Create %(question_num)s diverse questions specifically for Descriptive Data Analysis. Ensure that the questions are varied and do not include any analytical or predictive elements. Avoid starting questions with 'how'; instead, use 'what', 'list', or similar beginnings. Each question should be written in %(language)s."%{"language": language, "question_num": question_num} if question_level == SuggestQuestionLevel.COMPLEX: step_instruction = "2. Formulate %(num)s question(s) for each of the following types of data analysis: Descriptive, Inferential, Diagnostic, Predictive, and Prescriptive. Each question should be written in %(language)s."%{"language": language, "num": math.ceil(question_num/5)} task_prompt = """Act as a data analyst, your objective is to formulate a list of analytical questions based on the provided information. These questions should guide the user in understanding and exploring the data more effectively. ### Description of the Data Source ### %(summary)s ### Entities in Data Source ### %(entity_schema)s Let's think step by step as follows: 1. Use your expertise and skills to draw meaningful insights from the information given about the Data Source. %(step_instruction)s 3. For each question created in step 2, identify a single striking keyword that summarizes the question effectively. Remember, it must be just one word! The output should be a JSON array, as described below: [ { "type": , "question": , "keyword": string } ] Only output json, do not output anything else. Json: """%{"summary": db_context.get("summary", ""), "entity_schema": ("\n".join(entity_schema)), "step_instruction": step_instruction} resp_json = ai_completion("text2sql", task_prompt, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_SUGGEST_QUESTIONS, "session_id": self.session_context_id}) time_after = int(time.time()) logging.info("session_id = %s, suggest questions finished, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(time_after-time_before)) if not isinstance(resp_json, list) or len(resp_json) == 0: logging.info("session_id = %s, suggest questions failed, response_json is not array or it it empty", str(self.session_context_id)) raise Exception("Failed in generating recommended questions") return random.sample(resp_json, question_num if question_num <= len(resp_json) else len(resp_json)) def task_proxy(self, app_id, query, explaination = "", query_id = "", task_id = "", direct_resolve = True, auto_run = True, sql_to_data = False, data_to_chart = False): # query can be question or feedback if not query: raise BadRequestException(GenSQLErrorMsg.INVALID_PARAMS_ERROR) max_level = 1 vector_mode = VectorMode.TABLE_SCHEMA_VECTOR sql_revise_chains = [] language = LANGUAGE_EN ai_model = Claude3.DEFAULT # default mode mode_type = "" if app_id: preference = UserPreference.load(app_id) language = preference.language_code ai_model = preference.llm_model mode_type = preference.mode_type # 根据云平台自动选择模型 (阿里云 → Qwen) ai_model = self._determine_model_by_cloud_provider(ai_model) # flag mapping For Spider Test Mode if mode_type == MODE_TYPE_SPIDER: sql_to_data = False data_to_chart = False # Need to ensure the execution order. sql_revise_chains = ["spider_eval_revise", "ignore_case_revise", "sql_condition_revise"] # flag mapping For Bird Test Mode if mode_type == MODE_TYPE_BIRD: sql_to_data = False data_to_chart = False sql_revise_chains = ["spider_eval_revise"] # generate new query_id new_query_id = str(uuid.uuid4()) logging.info("session_id = %s, query_id = %s, gen new query_id, app_id = %s, input = %s, direct_resolve = %s, auto_run = %s, sql_to_data = %s, data_to_chart = %s", str(self.session_context_id), str(new_query_id), str(app_id), query, str(direct_resolve), str(auto_run), str(sql_to_data), str(data_to_chart)) # init root task input_query = "%(query)s **Explaination**: %(explaination)s"%{"query":query.replace('\n', ' ').strip(), "explaination": explaination.replace('\n', ' ').strip()} if explaination else query.replace('\n', ' ').strip() root_task = { "task": input_query, "clarified_task": input_query, "description": "", "level": 0, "task_id": "0", "query_id": new_query_id, "status": EdaTaskJobStatus.INIT, "created_at": int(time.time()) } if self.binded_thread_id: # Assistant API is not supported. if ai_model not in (GPT_4, GPT_4_TURBO, GPT_4O, GPT_4O_MINI): logging.info("session_id = %s, Assistant API is not supported for this model, app_id = %s, ai_model = %s, query = %s", str(self.session_context_id), str(app_id), ai_model, query) raise Exception("Conversation is not supported for model: %s".format(ai_model)) if query_id and task_id: self.__format_root_task(query, new_query_id, query_id, task_id, root_task) # if query_id and task_id: # logging.info("session_id = %s, query_id = %s, task_id = %s, deal with user's feedback, app_id = %s, query = %s", str(self.session_context_id), str(query_id), str(task_id), str(app_id), query) # yield from self.task_user_feedback(query_id, task_id, query, max_level, direct_resolve, auto_run, sql_to_data, data_to_chart, ai_model= ai_model, language=language) # return # else: # resp_json = self.task_user_intent(query) # if resp_json.get("intent") == "Feedback": # yield from self.task_user_feedback(resp_json.get("related_message_id"), resp_json.get("related_sub_message_id"), query, max_level, direct_resolve, auto_run, sql_to_data, data_to_chart, ai_model= ai_model, language=language) # return # if resp_json.get("intent") == "Follow-up Question": # # supplement query # query = resp_json.get("new_query") # task clarify for the root task.【Dedicated】: extract keyword and fuzzy match columns (release-v20240425-2) self.task_clarification(new_query_id, root_task, direct_resolve, ai_model=ai_model, language=language) if not direct_resolve: yield self.tasks_tree, new_query_id # task loop while len(self.tasks_to_be_resolved) != 0: task = self.tasks_to_be_resolved.popleft() try: self.task_analyse(new_query_id, task, max_level, direct_resolve, ai_model=ai_model) yield self.tasks_tree, new_query_id # break down the given task if task["type"] == TaskType.BREAK_DOWN: self.task_breakdown(new_query_id, task, ai_model=ai_model, language=language) yield self.tasks_tree, new_query_id else: # run task if auto_run: task_tree = self.task_router(new_query_id, task["task_id"], sql_to_data=sql_to_data, data_to_chart=data_to_chart, sql_revise_chains=sql_revise_chains, language=language, ai_model=ai_model) yield task_tree, new_query_id except BadRequestCausedAlgorithmException as e: logging.error("session_id = %s, query_id = %s, task analyse error, task = %s, error = %s", str(self.session_context_id), str(new_query_id), json.dumps(task, ensure_ascii=False), str(e)) raise e except Exception as e: logging.info("session_id = %s, query_id = %s, task analyse error, task = %s, error = %s", str(self.session_context_id), str(new_query_id), json.dumps(task, ensure_ascii=False), str(e)) raise AlgorithmException(str(e)) def __format_root_task(self, query, query_id, linked_query_id, linked_task_id, root_task): session_context = context.load_session_context(self.session_context_id) task = session_context.get("tasks_tree", {}).get(linked_query_id, {}).get(linked_task_id, {}) if not task.get("task"): logging.info("session_id = %s, query_id = %s, linked_query_id = %s, linked_task_id = %s, The specified task for feedback does not exist.", str(self.session_context_id), str(query_id), str(linked_query_id), str(linked_task_id)) raise Exception("The specified task for feedback does not exist.") root_task["feedback"] = query root_task["linked_query_id"] = linked_query_id root_task["linked_task_id"] = linked_task_id # merge the linked query root_task["task"] = query + ". It's feedback for the task: " + task.get("clarified_task", task["task"]) root_task["clarified_task"] = query + ". It's feedback for the task: " + task.get("clarified_task", task["task"]) if task.get("parent_task_id"): if not task.get("sub_task_list"): root_task["type"] = root_task.get("type") # clarify the root task def task_clarification(self, query_id, task, no_clarify_need=False, ai_model=GPT_DEFAULT, language=LANGUAGE_EN): if no_clarify_need: logging.info("session_id = %s, query_id = %s, no clarify needed for task", str(self.session_context_id), str(query_id)) if not task.get("clarified_task"): task["clarified_task"] = task.get("task") elif task.get("clarified_task") and task.get("description"): logging.info("session_id = %s, query_id = %s, has clarify for task", str(self.session_context_id), str(query_id)) else: if self.agent_clarification(task, query_id, ai_model, language, full_schema=True, dry_run=True): resp_json = self.agent_clarification(task, query_id, ai_model, language, full_schema=True, dry_run=False) else: resp_json = self.agent_clarification(task, query_id, ai_model, language, full_schema=False, dry_run=False) # TODO: Add Security Guard task["assumptions"] = resp_json.get("possibleExplanations", []) task["clarified_task"] = resp_json.get("refinedTask", task.get("task")) task["description"] = resp_json.get("description", task.get("description")) # add task to queue and task_tree self.tasks_to_be_resolved = deque([task]) # update tasks_tree self.tasks_tree = {task["task_id"]: task} # store to context session_context = context.load_session_context(self.session_context_id) all_task_trees = session_context.get("tasks_tree", {}) all_task_trees[query_id] = self.tasks_tree session_context["tasks_tree"] = all_task_trees context.save_session_context(self.session_context_id, session_context) # Agent for clarify task: may use tools in future def agent_clarification(self, task, query_id, ai_model, language, full_schema=True, dry_run=False): time_before = int(time.time()) db_context = context.load_dataset_context() # if data context is not ready, throw error if db_context.get("status") != "done": logging.info("session_id = %s, query_id = %s, data context exploration is not yet complete", str(self.session_context_id), str(query_id)) raise AlgorithmException(GenSQLErrorMsg.DATA_CONTEXT_NOT_COMPLETE_ERROR) tables_schema = [] if full_schema: # format the template table_template = "Table: `%(table_name)s`, Description: %(description)s, All columns are as follows:\n%(columns_str)s" column_template = "- column: `%(name)s`, description: %(description)s" for table in db_context.get("db_schema", {}).keys(): columns_list = [column_template%column for column in db_context["db_schema"].get(table, {}).get("columns", {}).values()] column_str = "\n".join(columns_list) tables_schema.append(table_template%{"columns_str": column_str, **db_context["db_schema"].get(table, {})}) else: for entity in db_context.get("entity", {}).values(): schema = "Entity: %(name)s, Description: %(summary)s, Involved_Tables: [%(involved_tables)s], Key_Attributes: %(attributes)s"%{ "name": entity["name"], "summary": entity["summary"], "attributes": json.dumps(entity["attributes"]), "involved_tables": ",".join(["{'table': %(table)s, 'description': %(description)s}"%{"table": item, "description": db_context["db_schema"].get(item, {}).get("description", "")} for item in entity["involved_tables"]]) } tables_schema.append(schema) extra_info = { "module": ModuleType.TASK_CLARIFY, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id"), "dry_run": dry_run } # retrieval term sheet term_sheet = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.TERM_SHEET, ModuleType.ALL, task.get("task"), extra_info=extra_info) # RAG instructions = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.INSTRUCTION, ModuleType.TASK_CLARIFY, task.get("clarified_task"), extra_info=extra_info) relevant_chunked_texts = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.CHUNKED_TEXT, ModuleType.TASK_CLARIFY, task.get("clarified_task"), extra_info=extra_info) task_prompt = """ Act as a data analyst, your objective is to clarify the task requirements more thoroughly in the context of information provided. ### Database summary ### %(summary)s ### Key entities in the database ### %(tables_schema)s %(relevant_chunked_texts)s%(instructions)s --- Task: %(task)s%(term_sheet)s --- Let's think step by step as follows: 1. Identify Main Concepts: Enrich the given task in the context and identify the main concepts or variables mentioned in the task. 2. Check for Ambiguity: Examine the concepts or variables identified in step 1 to determine if any of them are unclear or ambiguous based on the provided information. 3. Generate Explanations In %(language)s: For each unclear concept or variable from step 2, generate a most relevant explanation based on the provided information. 4. Refine the Task In %(language)s: Incorporate the generated explanation from step 3 to refine the task. Ensure that the refined task explicitly conveys the desired outcome, eliminating any ambiguous concepts. 5. Provide Detailed Description In %(language)s: provide a more detailed description for the refined task in step 4, outlining the specific objective, expected outcomes, and any relevant considerations or requirements. The output should be a JSON object, as described below: { "mainConceptsOrVariables": [
,...], "unclearConceptsOrVariables":[{"concept": , "reason": }, ...], "possibleExplanations": [{"concept": , "explanation": string}, ...], "refinedTask": , "description": } Only output json, do not output anything else. Json: """%{"summary":db_context.get("summary", ""), "tables_schema":("\n".join(tables_schema)), "language": language, "task": task["task"], "term_sheet": term_sheet, "instructions": instructions, "relevant_chunked_texts": relevant_chunked_texts} if dry_run: is_success = ai_completion("text2sql", task_prompt, "json", dry_run=True, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, task clarify for root task, dry run, model = %s, is_success = %s", str(self.session_context_id), str(query_id), ai_model, str(is_success)) return is_success else: if self.binded_thread_id: message = "clarify the task: %(task)s%(term_sheet)s"%{"term_sheet": term_sheet, **task} if task.get("feedback"): message = "clarify the task: %(task)s%(term_sheet)s\n**Supplementary Feedback**:\n%(feedback)s"%{"term_sheet": term_sheet, **task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, input_message_id, output_message_id = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_CLARIFY, "json", ai_model=ai_model, instructions=instructions, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task clarify for root task in conversation mode, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) if not resp_json.get("refinedTask"): ai_assistant_delete_message(input_message_id, output_message_id, self.binded_thread_id, extra_info=extra_info) else: resp_json = ai_completion("text2sql", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task clarify for root task, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) return resp_json # determine which workflow to use def task_analyse(self, query_id, task, max_level, no_proxy_need=False, ai_model=GPT_DEFAULT): if no_proxy_need: logging.info("session_id = %s, query_id = %s, no needed to breakdown", str(self.session_context_id), str(query_id)) if not task.get("type"): task["type"] = TaskType.DATA_RETRIEVAL task["agent"] = AgentType.CHAT2DATA elif task["level"] >= max_level: logging.info("session_id = %s, query_id = %s, no need to breakdown, max_level = %s", str(self.session_context_id), str(query_id), str(max_level)) if not task.get("type"): task["type"] = TaskType.DATA_RETRIEVAL task["agent"] = AgentType.CHAT2DATA elif task.get("type"): logging.info("session_id = %s, query_id = %s, task has been analysed, agent = %s, type = %s", str(self.session_context_id), str(query_id), task.get("agent"), task.get("type")) else: if self.agent_analyse(query_id, task, max_level, ai_model, full_schema=True, dry_run=True): resp_json = self.agent_analyse(query_id, task, max_level, ai_model, full_schema=True, dry_run=False) else: resp_json = self.agent_analyse(query_id, task, max_level, ai_model, full_schema=False, dry_run=False) task["reason"] = resp_json.get("reason", "") # TODO: set Agent Type: This determines which workflow the task will follow. task["agent"] = AgentType.PLAN if resp_json.get("category") == "Inferential Analysis" else AgentType.CHAT2DATA # set root task type task["type"] = TaskType.BREAK_DOWN if resp_json.get("category") == "Inferential Analysis" else TaskType.DATA_RETRIEVAL # update self.tasks_tree[task["task_id"]] = task # store to context session_context = context.load_session_context(self.session_context_id) tasks_tree = session_context.get("tasks_tree", {}).get(query_id, {}) tasks_tree[task["task_id"]] = task context.save_session_context(self.session_context_id, session_context) def agent_analyse(self, query_id, task, max_level, ai_model, full_schema=True, dry_run=False): time_before = int(time.time()) db_context = context.load_dataset_context() tables_schema = [] if full_schema: # format the template table_template = "Table: `%(table_name)s`, Description: %(description)s, All columns are as follows:\n%(columns_str)s" column_template = "- column: `%(name)s`, description: %(description)s" for table in db_context.get("db_schema", {}).keys(): columns_list = [column_template%column for column in db_context["db_schema"].get(table, {}).get("columns", {}).values()] column_str = "\n".join(columns_list) tables_schema.append(table_template%{"columns_str": column_str, **db_context["db_schema"].get(table, {})}) else: for entity in db_context.get("entity", {}).values(): schema = "Entity: %(name)s, Description: %(summary)s, Involved_Tables: [%(involved_tables)s], Key_Attributes: %(attributes)s"%{ "name": entity["name"], "summary": entity["summary"], "attributes": json.dumps(entity["attributes"]), "involved_tables": ",".join(["{'table': %(table)s, 'description': %(description)s}"%{"table": item, "description": db_context["db_schema"].get(item, {}).get("description", "")} for item in entity["involved_tables"]]) } tables_schema.append(schema) extra_info = { "module": ModuleType.TASK_ANALYSIS, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } # RAG instructions = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.INSTRUCTION, ModuleType.TASK_ANALYSIS, task.get("clarified_task"), extra_info=extra_info) few_shot_examples = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.FEW_SHOT, ModuleType.TASK_ANALYSIS, task.get("clarified_task"), extra_info=extra_info) task_prompt = """ Act as a data scientist, your objective is to accurately categorize the given task based on the provided information: ### Database Summary ### %(summary)s ### Key Entities in the Database ### %(entity_schema)s ### Task Categories ### **Descriptive Analysis**: This task primarily involves data retrieval and performing descriptive statistical analysis to provide a straightforward summary or interpretation of the data. Here are some examples for reference: ``` - "What are the total sales in the last quarter?" - "How many users registered in the past week?" - "What are the names and email addresses of all customers who made a purchase in the last month?" ``` **Inferential Analysis**: This task goes beyond merely retrieving and describing data (as in descriptive statistics) and aims to make predictions or inferences about a population based on sample data. Here are some examples for reference: ``` - "Analyze the sales data to identify seasonal trends and patterns." - "Predict future sales based on historical data." - "Which factors influencing employee productivity." ``` %(instructions)s --- Task: %(clarified_task)s --- %(few_shot_examples)s Let's think step by step as follows: 1. Clearly define the task objective and choose the most suitable category for the given task from the above categories. 2. Provide the main reason for selecting this category. The output should be a JSON object, as described below: { "category": , "reason": } Only output json, do not output anything else. Json: """%{"summary":db_context.get("summary") or "", "entity_schema":("\n".join(tables_schema)), "instructions": instructions, "few_shot_examples": few_shot_examples, **task} if dry_run: is_success = ai_completion("text2sql", task_prompt, "json", dry_run=True, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, task proxy for task, dry run, model = %s, is_success = %s", str(self.session_context_id), str(query_id), ai_model, str(is_success)) return is_success else: if self.binded_thread_id: message = "categorize the task: %(clarified_task)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, _, _ = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_ANALYSIS, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task proxy for task in conversation mode, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) else: resp_json = ai_completion("text2sql", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task proxy for task, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) return resp_json def task_breakdown(self, query_id, task, init_resp=[], ai_model=GPT_DEFAULT, language=LANGUAGE_EN): # default setting resp_json = init_resp if len(resp_json) == 0: if self.agent_breakdown(query_id, task, ai_model, language, full_schema=True, dry_run=True): resp_json = self.agent_breakdown(query_id, task, ai_model, language, full_schema=True, dry_run=False) else: resp_json = self.agent_breakdown(query_id, task, ai_model, language, full_schema=False, dry_run=False) sub_task_list = [] sub_task_id_list = [] # store to queue for index, refinedTask in enumerate(resp_json): # construct new task subtask_id = task["task_id"]+"-"+str(index) sub_task_id_list.append(subtask_id) new_task = { "task": refinedTask["sub_task"], "clarified_task": refinedTask["sub_task"], "description": refinedTask["description"], "type": refinedTask["step"], "parent_task_id": task["task_id"], "level": task["level"] + 1, "task_id": subtask_id, "created_at": int(time.time()), "status": EdaTaskJobStatus.INIT } sub_task_list.append(new_task) # store to tasks_tree self.tasks_tree[subtask_id] = new_task # add new task to loop for item in list(reversed(sub_task_list)): self.tasks_to_be_resolved.appendleft(item) # store to task_tree task["sub_task_list"] = sub_task_id_list task["status"] = EdaTaskJobStatus.RUNNING self.tasks_tree[task["task_id"]] = task # store to session session_context = context.load_session_context(self.session_context_id) all_task_trees = session_context.get("tasks_tree", {}) all_task_trees[query_id] = self.tasks_tree session_context["tasks_tree"] = all_task_trees context.save_session_context(self.session_context_id, session_context) # agent for breakdown: use tools in future def agent_breakdown(self, query_id, task, ai_model, language, full_schema=True, dry_run=False): time_before = int(time.time()) db_context = context.load_dataset_context() tables_schema = [] if full_schema: # format the template table_template = "Table: `%(table_name)s`, Description: %(description)s, All columns are as follows:\n%(columns_str)s" column_template = "- column: `%(name)s`, description: %(description)s" for table in db_context.get("db_schema", {}).keys(): columns_list = [column_template%column for column in db_context["db_schema"].get(table, {}).get("columns", {}).values()] column_str = "\n".join(columns_list) tables_schema.append(table_template%{"columns_str": column_str, **db_context["db_schema"].get(table, {})}) else: for entity in db_context.get("entity", {}).values(): schema = "Entity: %(name)s, Description: %(summary)s, Involved_Tables: [%(involved_tables)s], Key_Attributes: %(attributes)s"%{ "name": entity["name"], "summary": entity["summary"], "attributes": json.dumps(entity["attributes"]), "involved_tables": ",".join(["{'table': %(table)s, 'description': %(description)s}"%{"table": item, "description": db_context["db_schema"].get(item, {}).get("description", "")} for item in entity["involved_tables"]]) } tables_schema.append(schema) # task_prompt = """ # Act as a data scientist, your objective is to breakdown the given task into smaller subtasks. Here is the information provided: # ### Database Summary ### # %(summary)s # ### Key Entities in the Database ### # %(entity_schema)s # ### Instructions ### # Any given task can only be broken down into the following six sequential steps. Each step should not be repeated, but can be skipped if not applicable. Each corresponding to a unique subtask: # - **Data Retrieval**: Retrieve all necessary data from the data source efficiently, ensuring only essential data is fetched to minimize processing time and resource usage. # - **Data Cleaning and Preprocessing**: Clean and preprocess the retrieved data by addressing missing values, outliers, and inconsistencies, and transform the data into a suitable format for analysis. (Skip if not needed) # - **Data Analysis**: Perform further analysis on the preprocessed data using statistical methods, exploratory data analysis (EDA), and visualization techniques to uncover patterns and insights. (Skip if not needed) # - **Model Building**: If the task involves prediction or classification, select appropriate algorithms, split the data into training and testing sets, and tune hyperparameters. (Skip if not needed) # - **Model Evaluation**: If model building is required, use cross-validation and other robust evaluation techniques to ensure model reliability and generalizability. (Skip if not needed) # - **Summary and Recommendations**: Summarize key findings, draw actionable conclusions, and suggest next steps or business implications. As the final step, it is mandatory. # --- # Task: %(clarified_task)s # --- # Let's think step by step as follows: # 1. Identify all potential steps needed to solve the given task and each potential step should be unique. Ensure the results of all identified steps collectively cover the entire scope of the given task. # 2. Incorporate the given task and provided information, transform each identified step into a subtask in %(language)s. Ensure they are well-defined and clearly related to the given task, so that they can be understood independently. # 3. For each subtask, provide a more detailed description in %(language)s to outline the specific objective, expected outcomes, and any relevant considerations or requirements. # The output should be a JSON array, as described below: # ``` # [ # { # "step": # "sub_task": , # "description": # } # ] # ``` # Only output json, do not output anything else. Json: # """%{"summary": db_context.get("summary") or "", "entity_schema": ("\n".join(tables_schema)), "language": language, **task} extra_info = { "module": ModuleType.TASK_BREAKDOWN, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } # retrieve instructions instructions = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.INSTRUCTION, ModuleType.TASK_BREAKDOWN, task.get("clarified_task"), extra_info=extra_info) # retrieve few-show examples few_shot_examples = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.FEW_SHOT, ModuleType.TASK_BREAKDOWN, task.get("clarified_task"), extra_info=extra_info) task_prompt_v2 = """ Act as a data scientist with access to a cleaned database of relevant data. Your objective is to break down the given task into smaller, actionable subtasks for systematic execution. Here is the information provided: ### Database Summary ### %(summary)s ### Key Entities in the Database ### %(entity_schema)s %(instructions)s --- Task: %(clarified_task)s --- %(few_shot_examples)s Let's think step by step as follows: 1. Identify potential subtasks that contribute to solving the task. Keep the following considerations in mind: ``` - Each subtask should have a specific goal or objective that directly addresses a distinct aspect of the overall task. - The results of all subtasks should collectively cover the entire scope of the original task. - Avoid making subtasks too broad to ensure clarity and prevent overlap. - Resist the tendency to create subtasks solely for the purpose of data visualization. ``` 2. Refine the potential subtasks in %(language)s to ensure they are well-defined with clear conditions for completion. 3. For each refined subtask, provide a more detailed description in %(language)s to outlines the specific objective, expected outcomes, and any relevant considerations or requirements. The output should be a JSON array, as described below: [ { "sub_task": , "description": } ] Only output json, do not output anything else. Json: """%{"summary": db_context.get("summary", ""), "entity_schema": ("\n".join(tables_schema)), "language": language, "few_shot_examples": few_shot_examples, "instructions": instructions, **task} if dry_run: is_success = ai_completion("text2sql", task_prompt_v2, "json", dry_run=True, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, breakdown for task, dry run, model = %s, is_success = %s", str(self.session_context_id), str(query_id), ai_model, str(is_success)) return is_success else: if self.binded_thread_id: # mock result for sample db resp_json = get_local_question_result(str(db_context.get("cluster_id"))+"_"+db_context.get("db_name"), task["task"], ModuleType.TASK_BREAKDOWN) logging.info("session_id = %s, query_id = %s, get mocked result for task breakdown in conversation mode, resp = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False)) if len(resp_json) == 0: message = "breakdown the given task into smaller subtasks: %(clarified_task)s. **Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt_v2, flags=re.DOTALL) resp_json, input_message_id, output_message_id = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_BREAKDOWN, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task breakdown for task in conversation mode, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) if not isinstance(resp_json, list) or len(resp_json) == 0: ai_assistant_delete_message(input_message_id, output_message_id, self.binded_thread_id, extra_info=extra_info) else: resp_json = ai_completion("text2sql", task_prompt_v2, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, task breakdown for task, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) for _, sub_task in enumerate(resp_json): # format sub_tasks if sub_task.get("step") == "Data Retrieval": sub_task["step"] = TaskType.DATA_RETRIEVAL elif sub_task.get("step") == "Data Cleaning and Preprocessing": sub_task["step"] = TaskType.DATA_PROCESS elif sub_task.get("step") == "Data Analysis": sub_task["step"] = TaskType.DATA_ANALYSIS elif sub_task.get("step") == "Model Building": sub_task["step"] = TaskType.MODEL_BUILD elif sub_task.get("step") == "Model Evaluation": sub_task["step"] = TaskType.MODEL_EVALUATION elif sub_task.get("step") == "Summary and Recommendations": sub_task["step"] = TaskType.SUMMARY else: sub_task["step"] = TaskType.DATA_RETRIEVAL logging.info("session_id = %s, query_id = %s, task breakdown for task, define step name finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) return resp_json def task_router(self, query_id, task_id, sql_to_data=True, data_to_chart=True, sql_revise_chains=[], language=LANGUAGE_EN, ai_model=GPT_DEFAULT): # validity check of the task session_context = context.load_session_context(self.session_context_id) task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) if not task.get("task"): logging.info("session_id = %s, query_id = %s, The task ID does not exist, task_id = %s", str(self.session_context_id), str(query_id), str(task_id)) raise AlgorithmException("The task does not exist.") # If it is a breakdown task if task.get("type") == TaskType.BREAK_DOWN: raise AlgorithmException("The current task cannot directly resolved") # if status of task is "Done", return directly if task.get("status") == EdaTaskJobStatus.DONE: self.tasks_tree[task["task_id"]] = task return self.tasks_tree # If it is children task, check run sequence if task.get("parent_task_id"): sub_task_list = session_context.get("tasks_tree", {}).get(query_id, {}).get(task["parent_task_id"], {}).get("sub_task_list", []) try: current_task_index = sub_task_list.index(task_id) if current_task_index > 0: last_task = session_context.get("tasks_tree", {}).get(query_id, {}).get(sub_task_list[current_task_index-1], {}) if last_task.get("status") != EdaTaskJobStatus.DONE: raise AlgorithmException("The prior task remains unexecuted.") except ValueError as e: logging.info("session_id = %s, query_id = %s, check run sequence error, error = %s, task_id = %s", str(self.session_context_id), str(query_id), str(e), str(task_id)) raise AlgorithmException("Internal Error") # Router if task.get("type") == TaskType.DATA_RETRIEVAL: logging.info("session_id = %s, query_id = %s, route to text2sql, task_id = %s", str(self.session_context_id), str(query_id), str(task_id)) resp = self.task_text2sql(query_id, task_id, VectorMode.TABLE_SCHEMA_VECTOR, sql_to_data, data_to_chart, sql_revise_chains, language, ai_model) if sql_to_data: # insert column_data and row_data # copy_tasks_tree = copy.deepcopy(self.tasks_tree) # copy_tasks_tree[resp.get("task_id")]["columns"] = resp.get("columns", []) # copy_tasks_tree[resp.get("task_id")]["rows"] = resp.get("rows", []) self.tasks_tree[task_id]["columns"] = resp.get("columns", []) self.tasks_tree[task_id]["rows"] = resp.get("rows", []) if task.get("type") == TaskType.DATA_PROCESS: logging.info("session_id = %s, query_id = %s, route to data process, task_id = %s", str(self.session_context_id), str(query_id), str(task_id)) self.task_data_process(query_id, task_id, language=language, ai_model=ai_model) if task["type"] == TaskType.DATA_ANALYSIS: logging.info("session_id = %s, query_id = %s, route to data analysis, task_id = %s", str(self.session_context_id), str(query_id), str(task_id)) self.task_data_analysis(query_id, task_id, language=language, ai_model=ai_model) if task["type"] == TaskType.MODEL_BUILD: self.task_model_building(query_id, task_id, language=language, ai_model=ai_model) if task["type"] == TaskType.MODEL_EVALUATION: self.task_model_evaluation(query_id, task_id, language=language, ai_model=ai_model) if task["type"] == TaskType.SUMMARY: logging.info("session_id = %s, query_id = %s, route to data summary, task_id = %s", str(self.session_context_id), str(query_id), str(task_id)) self.task_summary(query_id, task_id, language=language, ai_model=ai_model) # if the task is the last task if task.get("parent_task_id"): # retrieve session context again, avoid to overwrite new_session_context = context.load_session_context(self.session_context_id) parent_task = new_session_context.get("tasks_tree", {}).get(query_id, {}).get(task["parent_task_id"], {}) if task_id == parent_task.get("sub_task_list", [])[-1]: parent_task["status"] = EdaTaskJobStatus.DONE context.save_session_context(self.session_context_id, new_session_context) # update self.task_tree self.tasks_tree[task["parent_task_id"]] = parent_task return self.tasks_tree def task_data_process(self, query_id, task_id, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): time_before = int(time.time()) session_context = context.load_session_context(self.session_context_id) # query the current task task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # query parent task parent_task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task["parent_task_id"], {}) # query log of data retrieval task【To Delete】 if task.get("linked_query_id"): # it is feedback resolved_log = session_context.get("task_resolved_log", {}).get(task["linked_query_id"], {}) else: resolved_log = session_context.get("task_resolved_log", {}).get(query_id, {}) task_prompt = """ Act as a data scientist, you've completed the first step of the given analysis task, which is Data Retrieval. Now your objective is to clean and process the results from the Data Retrieval to align with the analytical requirements of the analysis task. Here is the provided information: ### Results of Data Retrieval Subtask ### - Columns: %(columns)s - Rows: %(rows)s --- Task: %(clarified_task)s\nDescription: %(description)s --- Let's think step by step as follows and generate the result in %(language)s: 1. Review the results of the Data Retrieval Subtask and the Given Analysis Task. Determine if data cleaning and processing are necessary. 2. Briefly describe your processing plan. If no further processing is needed, describe your reason. 3. Outline detailed data processing steps (If no further processing is needed, leave this section blank). Ensure each step must include the following: - The objective or goal of this step. - Specific actions to be taken. The output should be a JSON object, as described below: { "description": "steps": [ { "objective": "actions": } ] } Only output json, do not output anything else. Json: """%{"language": language, **parent_task, **resolved_log} extra_info = { "module": ModuleType.TASK_DATA_PROCESS, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } if self.binded_thread_id: message = "Clean and process the data results for the task : %(clarified_task)s\n**Task Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, _, _ = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_DATA_PROCESS, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data process for task in conversation mode, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) else: resp_json = ai_completion("data_analysis", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data process for task, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) task["description"] = resp_json.get("description", "") task["steps"] = resp_json.get("steps", []) # update status task["status"] = EdaTaskJobStatus.DONE # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) def task_data_analysis(self, query_id, task_id, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): time_before = int(time.time()) session_context = context.load_session_context(self.session_context_id) # query the current task task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # query parent task parent_task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task["parent_task_id"], {}) # query log of data retrieval task【To Delete】 if task.get("linked_query_id"): # it is feedback resolved_log = session_context.get("task_resolved_log", {}).get(task["linked_query_id"], {}) else: resolved_log = session_context.get("task_resolved_log", {}).get(query_id, {}) task_prompt = """ Act as a data scientist, you've retrieved the necessary data for a specified task and have completed data cleaning and processing. Now, your objective is to explore the data to identify patterns, relationships, or insights that are relevant to the given analysis task. Here is the provided information: ### Data for Analysis ### - Columns: %(columns)s - Rows: %(rows)s --- Task: %(clarified_task)s\nDescription: %(description)s --- Let's think step by step as follows and generate the result in %(language)s: 1. **General Analysis Plan**: Provide a high-level overview of your data analysis plan based on the given analysis task and data. 2. **Detailed Analysis Plan**: Break down your analysis plan into specific steps. Ensure each step includes the following: - The goal of this step. - The actions to be taken to achieve the goal, including the data analysis method you will use and how to implement it. 3. **Insights Extraction**: Based on your analysis plan, list all the insights you have derived from the provided data. The output should be a JSON object, as described below: { "description": "steps": [ { "objective": "actions": } ], "insights": [, ...] } Only output json, do not output anything else. Json: """%{"language": language, **parent_task, **resolved_log} extra_info = { "module": ModuleType.TASK_DATA_ANALYSE, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } if self.binded_thread_id: message = "Identify insights that are relevant to the task : %(clarified_task)s\n**Task Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, _, _ = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_DATA_ANALYSE, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data analyisis for task in conversation mode, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) else: resp_json = ai_completion("data_analysis", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data analyisis for task, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) task["description"] = resp_json.get("description", "") task["steps"] = resp_json.get("steps", []) task["insights"] = resp_json.get("insights", []) # update status task["status"] = EdaTaskJobStatus.DONE # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) def task_model_building(self, query_id, task_id, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): time_before = int(time.time()) session_context = context.load_session_context(self.session_context_id) # query the current task task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # query log of data retrieval task【To Delete】 if task.get("linked_query_id"): # it is feedback resolved_log = session_context.get("task_resolved_log", {}).get(task["linked_query_id"], {}) else: resolved_log = session_context.get("task_resolved_log", {}).get(query_id, {}) task_prompt = """ As a data scientist, your objective is to recommend the most suitable algorithms to construct a model for the prediction or classification analysis task. Here is all the information provided: ### Data for Analysis ### - Columns: %(columns)s - Rows: %(rows)s --- Task: %(clarified_task)s\nDescription: %(description)s --- Let's think step by step as follows: 1. Recommend the appropriate algorithm based on the problem and the given data. 2. Provide your reasoning for this choice. The output should be a JSON object, as described below: { "recommendations": string "reason": string, } Only output json, do not output anything else. Json: """%{**task, **resolved_log} extra_info = { "module": ModuleType.TASK_MODEL_BUILDING, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } if self.binded_thread_id: message = "construct a algorithms model to address the task effectively: %(clarified_task)s\n**Task Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, _, _ = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_MODEL_BUILDING, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, model build recommendation for task in conversation mode, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) else: resp_json = ai_completion("data_analysis", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data summary for task, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) # update status task["status"] = EdaTaskJobStatus.DONE task["recommendations"] = resp_json.get("recommendations") task["reason"] = resp_json.get("reason") # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) def task_model_evaluation(self, query_id, task_id, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): session_context = context.load_session_context(self.session_context_id) # query the current task task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # update status task["status"] = EdaTaskJobStatus.DONE # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) def task_summary(self, query_id, task_id, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): time_before = int(time.time()) session_context = context.load_session_context(self.session_context_id) # query the current task task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # query parent task parent_task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task["parent_task_id"], {}) # query insight data_insights = [] for subtask_id in parent_task.get("sub_task_list", []): if session_context.get("tasks_tree", {}).get(query_id, {}).get(subtask_id, {}).get("type") == TaskType.DATA_ANALYSIS: insights = session_context.get("tasks_tree", {}).get(query_id, {}).get(subtask_id, {}).get("insights", []) data_insights = [str(index) + ". " + insight for index, insight in enumerate(insights)] break # query log of data retrieval task【To Delete】 if task.get("linked_query_id"): # it is feedback resolved_log = session_context.get("task_resolved_log", {}).get(task["linked_query_id"], {}) else: resolved_log = session_context.get("task_resolved_log", {}).get(query_id, {}) task_prompt = """ Act as a data scientist, you've retrieved the necessary data and completed the data analysis for the given task. Now, your objective is to make a summary and provide recommendations based on the provided information: ### Data for Analysis ### - Columns: %(columns)s - Rows: %(rows)s ### Insight of Data Analysis ### %(insights)s --- Task: %(clarified_task)s\nDescription: %(description)s --- Let's think step by step to summarize key findings, draw actionable conclusions, and suggest next steps or business implications in %(language)s The output should be a JSON object, as described below: { "conclusions": string, "recommendations": string } Only output json, do not output anything else. Json: """%{"insights": "\n".join(data_insights), "language": language, **parent_task, **resolved_log} extra_info = { "module": ModuleType.TASK_SUMMARY, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } if self.binded_thread_id: message = "make a summary and provide recommendations for the task: %(clarified_task)s\n**Task Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, _, _ = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_SUMMARY, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data summary for task in conversation mode, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) else: resp_json = ai_completion("data_analysis", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, data summary for task, finished, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) task["conclusions"] = resp_json.get("conclusions", "") task["recommendations"] = resp_json.get("recommendations", "") # update status task["status"] = EdaTaskJobStatus.DONE # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) def task_text2sql(self, query_id, task_id, vector_mode = VectorMode.TABLE_SCHEMA_VECTOR, sql_to_data = True, data_to_chart = True, sql_revise_chains=[], language=LANGUAGE_EN, ai_model=GPT_DEFAULT): time_before = int(time.time()) session_context = context.load_session_context(self.session_context_id) task = session_context.get("tasks_tree", {}).get(query_id, {}).get(task_id, {}) # execute sql directly if task.get("sql"): columns, rows, sql_error = self.execute_sql(query_id, task["sql"]) if sql_error: task["sql_error"] = sql_error # whether need to data2chart if data_to_chart and not sql_error: resp_json = self.task_data2chart(query_id, task["sql"], columns, rows[:3], task, ai_model=ai_model, language=language) task["chart_options"] = resp_json.get("chart") or {} task["recommendations"] = resp_json.get("data_analysis_methods") or {} context.save_session_context(self.session_context_id, session_context) # update tasks_tree self.tasks_tree[task["task_id"]] = task logging.info("session_id = %s, query_id = %s, directly execute SQL and data2chart, task = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(task, ensure_ascii=False), str(int(time.time()) - time_before)) return {"columns": columns, "rows": rows, "sql_error": sql_error, **task} table_schemas = context.load_dataset_context().get("db_schema", {}) necessary_tables = [{"table": key, "columns": []} for key, _ in table_schemas.items()] # Determine whether all db schema can import to prompt if vector_mode == VectorMode.TABLE_SCHEMA_VECTOR and self.text_to_sql(query_id, task, necessary_tables, vector_mode, False, False, [], dry_run=True, language=language, ai_model=ai_model): logging.info("session_id = %s, query_id = %s, text2sql with all tables and columns, model = %s", str(self.session_context_id), str(query_id), ai_model) resp_json = self.text_to_sql(query_id, task, necessary_tables, vector_mode, sql_to_data, data_to_chart, sql_revise_chains, language=language, ai_model=ai_model) else: # Embedding search with query. if vector_mode == VectorMode.COLUMN_SCHEMA_VECTOR: necessary_tables_dic = {} # format the embedding input embedding_input = task["clarified_task"] or task["task"] if task["description"]: embedding_input += ", %(description)s"%{**task} vector_storage = context.init_vector_storage(vector_mode, query_id + "_" + str(time_before)) relevant_columns = vector_storage.similarity_search_with_score(embedding_input, 50) # clear collection vector_storage.delete_collection() for item in relevant_columns: column_list = necessary_tables_dic.get(item[0].metadata["table"], []) column_list.append(item[0].metadata["column"]) necessary_tables_dic[item[0].metadata["table"]] = column_list necessary_tables = [{"table": key, "columns": value} for key, value in necessary_tables_dic.items()] logging.info("session_id = %s, query_id = %s, column schema vector to get relevant tables and columns, resp = %s", str(self.session_context_id), str(query_id), json.dumps(necessary_tables)) # text2sql resp_json = self.text_to_sql(query_id, task, necessary_tables, vector_mode, sql_to_data, data_to_chart, sql_revise_chains, language=language, ai_model=ai_model) else: vector_storage = context.init_vector_storage(VectorMode.TABLE_SCHEMA_VECTOR, query_id + "_" + str(time_before)) relevant_tables = vector_storage.similarity_search_with_score(task["clarified_task"], 30) # clear collection vector_storage.delete_collection() try: necessary_tables = self.get_necessary_tables(query_id, task, relevant_tables, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, table schema vector to get relevant tables and columns, resp = %s", str(self.session_context_id), str(query_id), json.dumps(necessary_tables)) resp_json = self.text_to_sql(query_id, task, necessary_tables, vector_mode, sql_to_data, data_to_chart, sql_revise_chains, language=language, ai_model=ai_model) except Exception as e: # single table has two many columns if GenSQLErrorMsg.TOKEN_EXCEED_ERROR in str(e): logging.info("session_id = %s, query_id = %s, get necessary tables or text2sql, token exceed and change to embedding retrieval, model = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, str(int(time.time()) - time_before)) # turn to retrieval relevant columns return self.task_text2sql(query_id, task_id, VectorMode.COLUMN_SCHEMA_VECTOR, sql_to_data, data_to_chart, sql_revise_chains, language, ai_model) logging.info("session_id = %s, query_id = %s, failed to get necessary tables or text2sql, error = %s", str(self.session_context_id), str(query_id), str(e)) raise Exception(GenSQLErrorMsg.REQUEST_OPENAI_ERROR) task["sql"] = resp_json.get("sql_query", "") task["chart_options"] = resp_json.get("chart", {}) task["recommendations"] = resp_json.get("data_analysis_methods", {}) if resp_json.get("sql_error"): task["sql_error"] = resp_json.get("sql_error") # update status task["status"] = EdaTaskJobStatus.DONE # TODO: For task breakdown, store data of sql execute, but now, it only have data_retrieval task # if task.get("parent_task_id"): # if not session_context.get("task_resolved_log"): # session_context["task_resolved_log"] = {} # task_resolved_log = session_context["task_resolved_log"].get(query_id, {}) # task_resolved_log["columns"] = resp_json.get("columns", []) # task_resolved_log["rows"] = resp_json.get("rows", []) if len(resp_json.get("rows", [])) <= 50 else random.sample(resp_json.get("rows", []), 50) # session_context["task_resolved_log"][query_id] = task_resolved_log # update tasks_tree for this query self.tasks_tree[task["task_id"]] = task # update context context.save_session_context(self.session_context_id, session_context) logging.info("session_id = %s, query_id = %s, text2sql for task, finished, task = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(task, ensure_ascii=False), str(int(time.time()) - time_before)) return {"columns": resp_json.get("columns", []), "rows": resp_json.get("rows", []), **task} def get_necessary_tables(self, query_id, task, relevant_tables, ai_model=GPT_DEFAULT) -> list: time_before = int(time.time()) relevant_table_names = [item[0].metadata["table"] for item in relevant_tables] logging.info("session_id = %s, query_id = %s, retrieval top30 relevant tables with clarified task by vector, relevant_tables = %s", str(self.session_context_id), str(query_id), str(relevant_table_names)) # If all relevant tables can be feed into prompt if self.llm_retrieval_tables(query_id, task, relevant_table_names, dry_run=True, ai_model=ai_model): necessary_tables = self.llm_retrieval_tables(query_id, task, relevant_table_names, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, retrieval tables and columns at once, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(necessary_tables), str(int(time.time()) - time_before)) return necessary_tables logging.info("session_id = %s, query_id = %s, Begin to get necessary tables with map reduce, model = %s, relevant_table_names = %s", str(self.session_context_id), str(query_id), ai_model, str(relevant_table_names)) patch_num_num = 7 if ai_model == GPT_4: patch_num_num = 4 # LLM Map Reduce necessary_tables = [] llm_times = 0 index = 0 patch_relevant_tables = [] while index < len(relevant_table_names): if index == llm_times * patch_num_num + patch_num_num: llm_times += 1 resp = self.llm_retrieval_tables(query_id, task, patch_relevant_tables, ai_model=ai_model) for item in resp: # post process if (item["table"] not in patch_relevant_tables) or (item["table"] in [table_item["table"] for table_item in necessary_tables]): continue necessary_tables.append(item) logging.info("session_id = %s, query_id = %s, Getting necessary tables with map reduce, current_necessary_tables = %s", str(self.session_context_id), str(query_id), json.dumps(necessary_tables)) patch_relevant_tables = [] patch_relevant_tables.append(relevant_table_names[index]) if index == (len(relevant_table_names) - 1) and len(patch_relevant_tables) != 0: resp = self.llm_retrieval_tables(query_id, task, patch_relevant_tables, ai_model=ai_model) for item in resp: # post process if (item["table"] not in patch_relevant_tables) or (item["table"] in [table_item["table"] for table_item in necessary_tables]): continue necessary_tables.append(item) logging.info("session_id = %s, query_id = %s, Getting necessary tables with map reduce, current_necessary_tables = %s", str(self.session_context_id), str(query_id), json.dumps(necessary_tables)) index += 1 logging.info("session_id = %s, query_id = %s, End to get necessary tables with map reduce, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), json.dumps(necessary_tables), str(int(time.time() - time_before))) return necessary_tables def llm_retrieval_tables(self, query_id, task, relevant_tables, dry_run=False, ai_model=GPT_DEFAULT): if len(relevant_tables) == 0: return [] db_context = context.load_dataset_context() tables_schema = [] tables_relationship = [] for table in relevant_tables: # Get table detail table_info = "Table: `%(table_name)s`, Description: %(description)s, Columns are as follows:\n"%db_context["db_schema"].get(table, {}) # Get column detail columns_list = ["- column: %(name)s, description: %(description)s"%column for column in list(db_context["db_schema"].get(table, {}).get("columns", {}).values())] tables_schema.append(table_info+("\n".join(columns_list))) # table relationship between necessary_tables for relationship in db_context.get("table_relationship", {}).get(table, []): if relationship["referencing_table"] == table and relationship["referenced_table"] in relevant_tables: tables_relationship.append("%(referencing_table)s.%(foreign_key_column)s = %(referenced_table)s.%(primary_key_column)s"%relationship) extra_info = {"session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id"), "module": ModuleType.TASK_RETRIEVAL_TABLES, "dry_run": dry_run} # RAG instructions = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.INSTRUCTION, ModuleType.TASK_RETRIEVAL_TABLES, task.get("clarified_task"), extra_info=extra_info) term_sheet = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.TERM_SHEET, ModuleType.TASK_RETRIEVAL_TABLES, task.get("clarified_task"), extra_info=extra_info) # Format prompt task_prompt = """ Act as an expert in text-to-SQL, your objective is to map a given task to relevant tables and columns based on the provided database schema and table relationships. Below is the information you have: ### Table Schema ### %(tables_schema)s ### Table relationships ### %(tables_relationship)s %(instructions)s --- Task: %(clarified_task)s%(term_sheet)s --- Let's think step by step as follows: 1. Carefully analyze the task to understand the specific objective, expected outcomes, and any relevant considerations or requirements. 2. Based on your analysis in step 1, identify all tables that are relevant to the task. Do not introduce non-existent tables. 3. For each table identified in step 2, identify all relevant columns needed to accomplish the task. Do not introduce non-existent columns. The output should only contains all relevant tables and columns identitied in a json format as described below: { "relevant_tables": ["", ...], "relevant_columns": { "": ["", ...], ... } } Do not output anything else. Json: """%{"tables_schema": ("\n\n".join(tables_schema)), "tables_relationship": ("\n".join(tables_relationship)), "instructions": instructions, "clarified_task": task.get("clarified_task"), "term_sheet": term_sheet} if dry_run: return ai_completion("text2sql", task_prompt, "json", dry_run=True, ai_model=ai_model) else: resp_json = ai_completion("text2sql", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) return [{"table": item, "columns": []} for item in resp_json.get("relevant_tables", [])] def text_to_sql(self, query_id, task, necessary_tables, vector_mode, sql_to_data, data_to_chart, sql_revise_chains, dry_run=False, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): # get db schema db_context = context.load_dataset_context() table_template = "Table: `%(table_name)s`, Description: %(description)s, Columns are as follows:\n%(columns_str)s" column_template = "- column: `%(name)s`, type: %(type)s, description: %(description)s, sample_data: %(sample_data)s" tables_schema = [] tables_relationship = [] # Vector search will only search for some columns of the table. if vector_mode == VectorMode.COLUMN_SCHEMA_VECTOR: table_template = "Table: `%(table_name)s`, Columns are as follows:\n%(columns_str)s" for item in necessary_tables: table = item.get("table", "") if vector_mode == VectorMode.COLUMN_SCHEMA_VECTOR: columns_list = [column_template%db_context["db_schema"].get(table, {}).get("columns", {}).get(column, {}) for column in item.get("columns", [])] # Add PK for each relevant table primary_key = db_context["db_schema"].get(table, {}).get("primary_key") if primary_key not in item.get("columns", []): primary_key_column = db_context["db_schema"].get(table, {}).get("columns", {}).get(primary_key, {}) if primary_key_column: columns_list.append(column_template%db_context["db_schema"].get(table, {}).get("columns", {}).get(primary_key, {})) # Add FK and relationship for each relevant table for relationship in db_context.get("table_relationship", {}).get(table, []): if relationship["referencing_table"] == table and relationship["referenced_table"] in [item.get("table", "") for item in necessary_tables]: if "- %(referenced_table)s.%(primary_key_column)s = %(referencing_table)s.%(foreign_key_column)s"%relationship not in tables_relationship: tables_relationship.append("- %(referencing_table)s.%(foreign_key_column)s = %(referenced_table)s.%(primary_key_column)s"%relationship) if relationship["foreign_key_column"] not in item.get("columns", []): columns_list.append(column_template%db_context["db_schema"].get(table, {}).get("columns", {}).get(relationship["foreign_key_column"], {})) else: columns_list = [column_template%column for column in db_context["db_schema"].get(table, {}).get("columns", {}).values()] # table relationship between necessary_tables for relationship in db_context.get("table_relationship", {}).get(table, []): if relationship["referencing_table"] == table and relationship["referenced_table"] in [item.get("table", "") for item in necessary_tables]: if "- %(referenced_table)s.%(primary_key_column)s = %(referencing_table)s.%(foreign_key_column)s"%relationship not in tables_relationship: tables_relationship.append("- %(referencing_table)s.%(foreign_key_column)s = %(referenced_table)s.%(primary_key_column)s"%relationship) column_str = "\n".join(columns_list) tables_schema.append(table_template%{"columns_str": column_str, **db_context["db_schema"].get(table, {})}) # get the type of database database_type = "MySQL" if db_context.get("database_uri", "").lower().startswith("sqlite"): database_type = "SQLite" if db_context.get("database_uri", "").lower().startswith("bigquery"): database_type = "BigQuery" if db_context.get("database_uri", "").lower().startswith("postgresql"): database_type = "PostgreSQL" # gen SQL if dry_run: return self.text_to_sql_single_step(query_id, task, tables_schema, tables_relationship, database_type, dry_run=True, ai_model=ai_model, language=language, necessary_tables=necessary_tables, app_id=db_context.get("app_id", 0)) # store task_prompt and resp_json for text2SQL. For Conversation in SQLEditor sql_query = self.text_to_sql_single_step(query_id, task, tables_schema, tables_relationship, database_type, dry_run=False, ai_model=ai_model, language=language, necessary_tables=necessary_tables, app_id=db_context.get("app_id", 0)) # No SQL Generated if not sql_query: # TODO: if it is for subtask of breakdown workflow, don't throw exception, keep running if task.get("parent_task_id"): return {"sql_query": "No SQL Generated", "sql_error": "", "columns": [], "rows": []} if len(list(db_context["db_schema"].keys())) == 0: logging.info("session_id = %s, query_id = %s, Gen SQL Fail, Database Empty, task = %s", str(self.session_context_id), str(query_id), json.dumps(task, ensure_ascii=False)) raise BadRequestCausedAlgorithmException(GenSQLErrorMsg.DATABASE_EMPTY_ERROR%{"db_name": db_context.get("db_name", "")}) logging.info("session_id = %s, query_id = %s, Gen SQL Fail, No SQL Generated, task = %s", str(self.session_context_id), str(query_id), json.dumps(task, ensure_ascii=False)) raise AlgorithmException(GenSQLErrorMsg.NO_SQL_GENERATED) # post process for SQL optimized_sql_query, sql_error = self.post_process_for_gen_sql(query_id, task, tables_schema, tables_relationship, database_type, sql_query, sql_revise_chains, ai_model=ai_model) # Pre-define columns_data = [] rows_data = [] if sql_error: return {"sql_query": optimized_sql_query, "sql_error": sql_error, "columns": columns_data, "rows": rows_data} # need execute sql and return result data: when no sql_self_refine in sql_revise_chains if sql_to_data: columns_data, rows_data, sql_error = self.execute_sql(query_id, optimized_sql_query) # need to chart chart_info = {} if data_to_chart and not sql_error: chart_info = self.task_data2chart(query_id, sql_query, columns_data, rows_data[:3], task, ai_model=ai_model, language=language) return {"sql_query": optimized_sql_query, "sql_error": sql_error, "columns": columns_data, "rows": rows_data, **chart_info} def text_to_sql_single_step(self, query_id, task, tables_schema, tables_relationship, database_type, dry_run=False, ai_model=GPT_DEFAULT, language=LANGUAGE_EN, **kwargs): time_before = int(time.time()) # Load preference preference = UserPreference.load(kwargs.get("app_id", 0)) if preference.mode_type == MODE_TYPE_BIRD or preference.mode_type == MODE_TYPE_SPIDER: return self.text_to_sql_for_leaderboard(query_id, task, tables_schema, tables_relationship, dry_run=dry_run, ai_model=ai_model) extra_info = {"session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id"), "module": ModuleType.TASK_TEXT2SQL, "dry_run": dry_run} meta_data = {"involved_tables": [item.get("table") for item in kwargs.get("necessary_tables", [])]} # RAG instructions = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.INSTRUCTION, ModuleType.TASK_TEXT2SQL, task.get("clarified_task"), meta_data=meta_data, extra_info=extra_info) term_sheet = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.TERM_SHEET, ModuleType.TASK_TEXT2SQL, task.get("clarified_task"), meta_data=meta_data, extra_info=extra_info) few_shot_examples = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.FEW_SHOT, ModuleType.TASK_TEXT2SQL, task.get("clarified_task"), meta_data=meta_data, extra_info=extra_info) relevant_chunked_texts = "" # retrieval tidb extension syntax if kwargs.get("app_id", 0) == AssignedAppID.SQLEDITOR: logging.info("session_id = %s, query_id = %s, text to SQL with single step for SQLEditor", str(self.session_context_id), str(query_id)) relevant_chunked_texts = get_kg_management(KnowledgeSourceType.REMOTE).retrieval_knowledge(KnowledgeType.CHUNKED_TEXT, ModuleType.TASK_TEXT2SQL, task.get("clarified_task"), meta_data=meta_data, extra_info=extra_info) # Supporting TiDB Syntax, use whitelist if get_local_metrics().get("org_id") in SUPPORT_TIDB_SYNTAX_WHITELIST_ORGS: ai_model = GPT_4_TURBO # Format prompt task_prompt = """Act as an expert in text-to-SQL, your objective is to write an SQL statement using %(database_type)s best practices to achieve the desired results for a given task. Here is the provided information: ### Table schemas ### %(tables_schema)s ### Table relationships ### %(tables_relationship)s %(relevant_chunked_texts)s%(instructions)s --- Task: %(clarified_task)s%(term_sheet)s --- %(few_shot_examples)s Let's think step by step as follows: 1. Clarify the Task in %(language)s: Thoroughly understand the task's specific goal, expected outcomes, and any relevant conditions. Ensure that you have a clear understanding of the task's requirements. 2. Identify Keywords: Identify all relevant keywords associated with tables and columns mentioned in the clarified task. 3. Map Keywords to Columns and Tables: Match the identified keywords in step 2 to specific columns and tables within the provided table schemas. For each keyword, determine the corresponding column name and table name where the keyword is relevant. 4. Generate SQL Query: Leverage the mapping results obtained from step 3, along with the indicated table relationships to construct a comprehensive SQL query that fulfills the task's requirements. If no SQL query can be generated, leave it empty. 5. Rewrite SQL Query: After you've generated the SQL query in step 4, refine it in line with these criteria: Use the table.column format when multiple tables are involved in the SQL query to avoid ambiguity. Ensure proper table and column references, making sure they exist in the appropriate tables. Enclose all table names and column names with backticks (`) to ensure correct referencing and avoid any syntax issues. The output should be a JSON object, as described below: { "clarified_task": string, "keywords": string[], "mapping": [{"keyword": , "related_column": , "related_table": }], "sql": string, "rewritten_sql": } Only output json, do not output anything else. Json: """%{"tables_schema": ("\n\n".join(tables_schema)), "tables_relationship": ("\n".join(tables_relationship)), "database_type": database_type, "term_sheet": term_sheet, "language": language, "instructions": instructions, "few_shot_examples": few_shot_examples, "relevant_chunked_texts": relevant_chunked_texts, **task} if dry_run: is_success = ai_completion("text2sql", task_prompt, "json", dry_run=True, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, text to SQL with single step, dry run, model = %s, is_success = %s", str(self.session_context_id), str(query_id), ai_model, str(is_success)) return is_success else: if self.binded_thread_id: message = "Write an SQL statement to achieve the desired results for the task: %(clarified_task)s **Task Description**: %(description)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, input_message_id, output_message_id = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_TEXT2SQL, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, text to SQL with single step in conversation mode, finished, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) # post process if not resp_json.get("rewritten_sql"): # Gen result has problem, need to delete ai_assistant_delete_message(input_message_id, output_message_id, self.binded_thread_id, extra_info=extra_info) else: resp_json = ai_completion("text2sql", task_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, text to SQL with single step, finished, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) # update the clarified task task["clarified_task"] = resp_json.get("clarified_task", task.get("clarified_task")) sql_query = resp_json.get("rewritten_sql", resp_json.get("5. Rewrite SQL Query", "")) # when 'rewritten_sql' is empty, try to get the 'sql' if not sql_query: sql_query = resp_json.get("sql", resp_json.get("4. Generate SQL Query", "")) if len(tables_schema) == 0 and (not sql_query.strip().upper().startswith("CREATE")) and (not sql_query.strip().upper().startswith("SHOW")): if 'INFORMATION_SCHEMA' not in sql_query.strip().upper(): sql_query = "" return sql_query def text_to_sql_for_leaderboard(self, query_id, task, tables_schema, tables_relationship, dry_run=False, ai_model=GPT_DEFAULT): time_before = int(time.time()) # Format task supplementary explanation task_description_prompt ="""\n**Supplementary Explanation**: %(description)s"""%{**task} if task.get("description") else "" # Format prompt context_prompt = """ Act as an expert in text-to-SQL, your objective is to write an SQL statement using SQLite best practices to achieve the desired results for a given task. Here is the provided information: ### Table schemas ### %(tables_schema)s ### Table relationships ### %(tables_relationship)s --- **Task**: %(task)s%(task_description_prompt)s --- """%{"tables_schema": ("\n\n".join(tables_schema)), "tables_relationship": ("\n".join(tables_relationship)), "task_description_prompt": task_description_prompt, **task} instruction_prompt = """ Let's think step by step as follows and generate a single JSON object in the JSON format that includes the results of each step: { 1. Clarify the Task: Thoroughly understand the task's specific goal, expected outcomes, and any relevant conditions. Ensure that you have a clear understanding of the task's requirements. {"clarified_task": string} 2. Identify Keywords: Identify all relevant keywords associated with tables and columns mentioned in the clarified task. {"keywords": string[]} 3. Map Keywords to Columns and Tables: Match the identified keywords in step 2 to specific columns and tables within the provided table schemas. For each keyword, determine the corresponding column name and table name where the keyword is relevant. {"mapping": [{"keyword": , "related_column": , "related_table": }]} 4. Recheck the Mapping Result: Double-check the mapping result to ensure that for each keyword, the related column exists in the relevant table based on the provided table schema. {"rechecked_mapping": [{"keyword": , "related_column": , "related_table": }]} 5. Generate SQL Query: Leverage the rechecked mapping results obtained from step 4, along with the indicated table relationships to construct a comprehensive SQL query that fulfills the task's requirements. If no SQL statement can be generated, leave it empty. {"sql": string} 6. Rewrite SQL Query: After you've generated the SQL query in step 5, refine it in line with these criteria: Use the table.column format when multiple tables are involved in the SQL query to avoid ambiguity. Ensure proper table and column references, making sure they exist in the appropriate tables. Enclose all table names and column names with backticks (`) to ensure correct referencing and avoid any syntax issues. {"rewritten_sql": } } """ if ai_model in (Claude3.DEFAULT, Claude3.SONNET, Claude3.HAIKU): instruction_prompt = """ Let's think step by step as follows: 1. Clarify the Task: Thoroughly understand the task's specific goal, expected outcomes, and any relevant conditions. Ensure that you have a clear understanding of the task's requirements. 2. Identify Keywords: Identify all relevant keywords associated with tables and columns mentioned in the clarified task. 3. Map Keywords to Columns and Tables: Match the identified keywords in step 2 to specific columns and tables within the provided table schemas. For each keyword, determine the corresponding column name and table name where the keyword is relevant. 4. Recheck the Mapping Result: Double-check the mapping result to ensure that for each keyword, the related column exists in the relevant table based on the provided table schema. 5. Generate SQL Query: Leverage the rechecked mapping results obtained from step 4, along with the indicated table relationships to construct a comprehensive SQL query that fulfills the task's requirements. If no SQL statement can be generated, leave it empty. 6. Rewrite SQL Query: After you've generated the SQL query in step 5, refine it in line with these criteria: Use the table.column format when multiple tables are involved in the SQL query to avoid ambiguity. Ensure proper table and column references, making sure they exist in the appropriate tables. Enclose all table names and column names with backticks (`) to ensure correct referencing and avoid any syntax issues. Provide the output in the following JSON format: ``` { "clarified_task": string, "keywords": string[], "mapping": [{"keyword": , "related_column": , "related_table": }], "rechecked_mapping": [{"keyword": , "related_column": , "related_table": }], "sql": string, "rewritten_sql": } ``` Please ensure that only the JSON output is provided, do not output anything else. Json: """ if dry_run: is_success = ai_completion("text2sql", context_prompt + instruction_prompt, "json", dry_run=True, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, text to SQL for leaderboard, dry run, model = %s, is_success = %s", str(self.session_context_id), str(query_id), ai_model, str(is_success)) return is_success else: extra_info = { "module": ModuleType.TASK_TEXT2SQL, "session_id": self.session_context_id, "query_id": query_id, "task_id": task.get("task_id") } try: resp_json = ai_completion("text2sql", context_prompt + instruction_prompt, "json", dry_run=False, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, text to SQL for leaderboard, finished, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) # update the clarified task task["clarified_task"] = resp_json.get("clarified_task", resp_json.get("1. Clarify the Task", "")) sql_query = resp_json.get("rewritten_sql", resp_json.get("6. Rewrite SQL Query", "")) # when 'rewritten_sql' is empty, try to get the 'sql' if not sql_query: sql_query = resp_json.get("sql", resp_json.get("5. Generate SQL Query", "")) return sql_query except Exception as e: # Token Exceed, Use Embedding Retrieval if GenSQLErrorMsg.TOKEN_EXCEED_ERROR in str(e): logging.info("session_id = %s, query_id = %s, text to SQL for leaderboard, token exceed and change to embedding retrieval, model = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, str(int(time.time()) - time_before)) resp_json = self.task_text2sql(query_id, task["task_id"], vector_mode = VectorMode.COLUMN_SCHEMA_VECTOR, sql_to_data = False, data_to_chart = False, sql_revise_chains=[], ai_model=ai_model) logging.info("session_id = %s, query_id = %s, text to SQL for leaderboard, change to embedding retrieval and successed, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) return resp_json.get("sql", "") else: raise BadRequestException(str(e)) def post_process_for_gen_sql(self, query_id, task, tables_schema, tables_relationship, database_type, sql_query, sql_revise_chains, ai_model=GPT_DEFAULT): logging.info("session_id = %s, query_id = %s, post process for gen sql, before, raw_sql_query = %s, sql_revise_chain = %s, ai_model = %s", str(self.session_context_id), str(query_id), sql_query, json.dumps(sql_revise_chains), ai_model) if not sql_query: return sql_query, "The SQL is empty" optimized_sql_query = sql_query.strip() sql_error = "" # in chat, not supported # if self.binded_thread_id: # return optimized_sql_query if optimized_sql_query.endswith(";") else optimized_sql_query + ";", sql_error for chain_name in sql_revise_chains: if chain_name == "spider_eval_revise": optimized_sql_query = self.revise_sql_for_spider_eval(query_id, task, database_type, optimized_sql_query, ai_model=ai_model) if chain_name == "ignore_case_revise": optimized_sql_query = self.revise_sql_for_ignore_case(query_id, task, database_type, optimized_sql_query) if chain_name == "sql_condition_revise": optimized_sql_query = self.revise_sql_for_refine_SQL_condition(query_id, task, database_type, optimized_sql_query, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, post process for gen sql, after revised by chains, raw_sql_query = %s, sql_revise_chain = %s, ai_model = %s, optimized_sql_query = %s", str(self.session_context_id), str(query_id), sql_query, json.dumps(sql_revise_chains), ai_model, optimized_sql_query) # SQL self refine if "SELECT" in optimized_sql_query.strip().upper(): optimized_sql_query, sql_error = self.explain_sql_and_refine(query_id, task, tables_schema, tables_relationship, database_type, optimized_sql_query, sql_error, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, post process for gen sql, after SQL self refine, raw_sql_query = %s, sql_revise_chain = %s, ai_model = %s, optimized_sql_query = %s", str(self.session_context_id), str(query_id), sql_query, json.dumps(sql_revise_chains), ai_model, optimized_sql_query) # when only one table listed in NT and the sql generated used table.column format if 'INFORMATION_SCHEMA' not in optimized_sql_query.strip().upper(): pattern = re.compile(r"`([a-zA-Z0-9_ ]+)`\.`[a-zA-Z0-9_ ]+`") matches = pattern.findall(optimized_sql_query) involved_tables = list(set(matches)) if len(involved_tables) == 1: comparedStr = "`" + involved_tables[0] + "`." optimized_sql_query = optimized_sql_query.replace(comparedStr, "") # For Spider: Revise the `value` to `Value` db_context = context.load_dataset_context() preference = UserPreference.load(db_context.get("app_id", 0)) if preference.mode_type == MODE_TYPE_SPIDER: optimized_sql_query = optimized_sql_query.replace("`value`", "`Value`") # add semicolon optimized_sql_query = optimized_sql_query.strip() if not optimized_sql_query.endswith(";"): optimized_sql_query += ";" logging.info("session_id = %s, query_id = %s, post process for gen sql, after revised by other processor, raw_sql_query = %s, sql_revise_chain = %s, ai_model = %s, optimized_sql_query = %s", str(self.session_context_id), str(query_id), sql_query, json.dumps(sql_revise_chains), ai_model, optimized_sql_query) return optimized_sql_query, sql_error def explain_sql_and_refine(self, query_id, task, tables_schema, tables_relationship, database_type, sql_query, sql_error, ai_model=GPT_DEFAULT): new_sql_error, is_need_refine = self.explain_sql(query_id, sql_query) if is_need_refine and not sql_error: # SQL Refine Once rewrited_sql = self.sql_self_refine(query_id, task, tables_schema, tables_relationship, database_type, sql_query, new_sql_error, ai_model=ai_model) return self.explain_sql_and_refine(query_id, task, tables_schema, tables_relationship, database_type, rewrited_sql, new_sql_error, ai_model=ai_model) logging.info("session_id = %s, query_id = %s, explain sql and refine finished, sql_query = %s, sql_error = %s", str(self.session_context_id), str(query_id), sql_query, new_sql_error) return sql_query, new_sql_error def revise_sql_for_spider_eval(self, query_id, task, database_type, sql_query, ai_model=GPT_4): time_before = int(time.time()) # revision by llm task_prompt = """ Act as a data analyst, your objective is to review and modify only the SELECT list of a given SQL to ensure it only includes the columns that are directly required by the task. Please note: **Do not consider the correctness of other parts of the SQL query.** **Do not add new columns or modify the column names in the select list.** ### Here are some examples for guidance ### Task: Which year had the most matches? Given SQL: SELECT `year`, count(*) as `match_num` FROM `matches` GROUP BY `year` ORDER BY `match_num` DESC LIMIT 1; Answer: The task is to find the specific year with the most matches. In the given SQL, the select list includes two columns: `year` and `match_num`, but `match_num` is not directly required by the task, it is only needed in ORDER BY clause for sorting and can be replaced by using `COUNT` clause. Therefore, we should remove it from the select list, replace it with `COUNT` clause in ORDER BY clause, and keep the rest of the SQL unchanged. Modified SQL: SELECT `year` FROM `matches` GROUP BY `year` ORDER BY count(*) DESC LIMIT 1; Task: When do all the researcher role staff start to work, and when do they stop working? Given SQL: SELECT `project_staff`.`staff_id`, `project_staff`.`date_from`, `project_staff`.`date_to` FROM `project_staff` WHERE `project_staff`.`role_code` COLLATE NOCASE = 'researcher'; Answer: The task is to find the start and end dates of work for all staff with the role of researcher. In the given SQL, the select list includes three columns: `staff_id`, `date_from`, and `date_to`, but the task does not require the `staff_id` directly. Therefore, we should remove it from the select list and keep the rest of the SQL unchanged. Modified SQL: SELECT `project_staff`.`date_from`, `project_staff`.`date_to` FROM `project_staff` WHERE `project_staff`.`role_code` COLLATE NOCASE = 'researcher'; Task: What is the grade of each high schooler? Given SQL: SELECT `highschooler`.`id`, `highschooler`.`grade` FROM `highschooler`; Answer: The task is to find the grade of each high schooler. In the given SQL, the select list includes two columns: `id` and `grade`. Since the task is to find the grade of each high schooler, we need both the `id` and `grade` columns in the select list to identify each high schooler and their corresponding grade. Therefore, the select list does not need any modification and the modified SQL is the same as the given SQL. Modified SQL: SELECT `highschooler`.`id`, `highschooler`.`grade` FROM `highschooler`; Task: Find the names of all the customers and staff members. Given SQL: SELECT `customers`.`customer_details` FROM `customers` UNION SELECT `staff`.`staff_details` FROM `staff`; Answer: The task is to find the names of all the customers and staff members. In the given SQL, the select list includes `customer_details` and `staff_details`. We are assuming that these columns contain the names of customers and staff members respectively, as we are not allowed to add new columns or modify the column names in the select list. Therefore, the select list does not need any modification and the modified SQL is the same as the given SQL. Modified SQL: SELECT `customers`.`customer_details` FROM `customers` UNION SELECT `staff`.`staff_details` FROM `staff`; **Task**: %(task)s **Given SQL**: %(sql_query)s **Answer**: . Modified SQL: """%{"sql_query": sql_query, "task": task.get("task")} optimized_sql = sql_query # determine if more than one column are selected start = optimized_sql.find("SELECT") + len("SELECT") end = optimized_sql.find("FROM") selected_columns = optimized_sql[start:end].strip().split(',') # Modified log resp_str = "" if len(selected_columns) > 1: logging.info("session_id = %s, query_id = %s, revise sql for spider eval, begin, sql_query = %s", str(self.session_context_id), str(query_id), optimized_sql) resp_str = ai_completion("text2sql", task_prompt, ai_model=ai_model) # Format the result try: optimized_sql = resp_str.split("Modified SQL:")[1].strip().strip('\n').strip().strip('.').strip(';') except Exception as e: logging.info("session_id = %s, query_id = %s, revise sql for spider eval, error, error = %s", str(self.session_context_id), str(query_id), str(e)) logging.info("session_id = %s, query_id = %s, revise sql for spider eval, after, sql_query = %s, time_cost = %s", str(self.session_context_id), str(query_id), optimized_sql, str(int(time.time()) - time_before)) # store `raw_generated_sql` and `refine_note` if not task.get("raw_generated_sql"): task["raw_generated_sql"] = sql_query task["refine_note"] = task.get("refine_note", "") + ("\n===\n"+resp_str) return optimized_sql def revise_sql_for_ignore_case(self, query_id, task, database_type, sql_query): logging.info("session_id = %s, query_id = %s, revise sql for ignore case, begin, sql_query = %s", str(self.session_context_id), str(query_id), sql_query) optimized_sql = sql_query if database_type == "SQLite": # add clause `COLLATE NOCASE` for case-insensitive in SQLite sql_snippets = optimized_sql.split(" = ") for index in range(1, len(sql_snippets)): if sql_snippets[index].strip().startswith("'"): sql_snippets[index-1] += " COLLATE NOCASE" optimized_sql = " = ".join(sql_snippets) if " NOT IN " in optimized_sql: optimized_sql = optimized_sql.replace(" NOT IN ", " COLLATE NOCASE NOT IN ") elif " IN " in optimized_sql: optimized_sql = optimized_sql.replace(" IN ", " COLLATE NOCASE IN ") else: pass logging.info("session_id = %s, query_id = %s, revise sql for ignore case, after, sql_query = %s", str(self.session_context_id), str(query_id), optimized_sql) return optimized_sql def revise_sql_for_refine_SQL_condition(self, query_id, task, database_type, sql_query, ai_model=GPT_4): time_before = int(time.time()) # revision by llm task_prompt = """ Act as an expert in %(database_type)s, your objective is to review and determine whether the given SQL needs to be restructured by replacing the `IN` or `UNION` clause with the `INTERSECT` clause to better meet the task's requirements. Please adhere to the following instructions: **Do not modify or delete any column names used in the given SQL query.** **Do not add any new columns in the restructured SQL query.** ### Below are some examples ### Task: What is the number of nations that use English and Dutch? Given SQL: SELECT COUNT(DISTINCT `country`.`code`) FROM `country` JOIN `countrylanguage` ON `country`.`code` = `countrylanguage`.`countrycode` WHERE `countrylanguage`.`language` COLLATE NOCASE IN ('English', 'Dutch') GROUP BY `country`.`code` HAVING COUNT(DISTINCT `countrylanguage`.`language`) = 2 Answer: The given SQL query is trying to find the number of countries that use both English and Dutch. However, the use of the `IN` clause in this case might not be the best approach. The `IN` clause will return true if the language is either English or Dutch, which means it will count countries that speak either of these languages, not necessarily both. To ensure that we only count countries that speak both languages, we can use the `INTERSECT` clause to restructure the given SQL. Restructured SQL: SELECT COUNT(*) FROM (SELECT `country`.`code` FROM `country` JOIN `countrylanguage` ON `country`.`code` = `countrylanguage`.`countrycode` WHERE `countrylanguage`.`language` COLLATE NOCASE = 'English' INTERSECT SELECT `country`.`code` FROM `country` JOIN `countrylanguage` ON `country`.`code` = `countrylanguage`.`countrycode` WHERE `countrylanguage`.`language` COLLATE NOCASE = 'Dutch') Task: Which airlines have departing flights from both APG and CVO airports? Given SQL: SELECT `a`.`airline` FROM `airlines` AS `a` INNER JOIN `flights` AS `f` ON `a`.`uid` = `f`.`airline` WHERE `f`.`sourceairport` COLLATE NOCASE IN ('APG', 'CVO') GROUP BY `a`.`airline` HAVING COUNT(DISTINCT `f`.`sourceairport`) = 2 Answer: The given SQL query is designed to find airlines that have departing flights from both APG and CVO airports. It does this by joining the `airlines` and `flights` tables on the `uid` and `airline` fields respectively. It then filters the results to only include flights where the `sourceairport` is either 'APG' or 'CVO'. The `GROUP BY` clause groups the results by airline, and the `HAVING` clause ensures that only airlines with flights departing from both airports are included in the final result. The given SQL query is correctly structured to meet the task's requirements, therefore, there is no need to replace the `IN` clause with the `INTERSECT` clause in this case. The restructured SQL is the same as the given SQL. Restructured SQL: SELECT `a`.`airline` FROM `airlines` AS `a` INNER JOIN `flights` AS `f` ON `a`.`uid` = `f`.`airline` WHERE `f`.`sourceairport` COLLATE NOCASE IN ('APG', 'CVO') GROUP BY `a`.`airline` HAVING COUNT(DISTINCT `f`.`sourceairport`) = 2 Task: How many concerts occurred in 2014 or 2015? Given SQL: SELECT COUNT(`concert`.`concert_id`) FROM `concert` WHERE `concert`.`year` COLLATE NOCASE IN ('2014', '2015') Answer: The given SQL query is trying to count the number of concerts that occurred in either 2014 or 2015. The `IN` clause is appropriate here because we are checking if the `year` column value is in a list of specific values (2014, 2015). The `INTERSECT` clause is used to return the common records between two select statements, which is not the requirement here. So, there is no need to replace the `IN` clause with the `INTERSECT` clause in this case. The restructured SQL is the same as the given SQL. Restructured SQL: SELECT COUNT(`concert`.`concert_id`) FROM `concert` WHERE `concert`.`year` COLLATE NOCASE IN ('2014', '2015') **Task**: %(task)s **Given SQL**: %(sql_query)s **Answer**: . Restructured SQL: """%{"sql_query": sql_query, "task": task.get("task"), "database_type": database_type} optimized_sql = sql_query # determine weather need to revise for INTERSECT if "UNION" in optimized_sql or "COLLATE NOCASE IN" in optimized_sql: logging.info("session_id = %s, query_id = %s, revise sql for refine SQL condition, begin, sql_query = %s", str(self.session_context_id), str(query_id), optimized_sql) resp_str = ai_completion("text2sql", task_prompt, ai_model=ai_model) try: optimized_sql = resp_str.split("Restructured SQL:")[1].strip().strip('\n').strip().strip('```').strip('.').strip(';') except Exception as e: logging.info("session_id = %s, query_id = %s, revise sql for refine SQL condition, error, error = %s", str(self.session_context_id), str(query_id), str(e)) logging.info("session_id = %s, query_id = %s, revise sql for refine SQL condition, after, sql_query = %s, time_cost = %s", str(self.session_context_id), str(query_id), optimized_sql, str(int(time.time()) - time_before)) if not task.get("raw_generated_sql"): task["raw_generated_sql"] = sql_query task["refine_note"] = task.get("refine_note", "") + ("\n===\n"+resp_str) return optimized_sql def sql_self_refine(self, query_id, task, tables_schema, tables_relationship, database_type, sql_query, sql_error, ai_model=GPT_DEFAULT): logging.info("session_id = %s, query_id = %s, SQL self refine, before, model = %s, sql_query = %s, sql_error = %s", str(self.session_context_id), str(query_id), ai_model, sql_query, sql_error) time_before = int(time.time()) task_prompt = """ Act as a SQL expert in %(database_type)s, your objective is to identify and fix an SQL query error for a specific task. Here is the information provided: ### Table schemas ### %(tables_schema)s ### Table relationships ### %(tables_relationship)s ### Original SQL Query ### %(sql_query)s ### Error Message ### %(sql_error)s --- Task: %(clarified_task)s --- Let's think step by step as follows: 1. Identify the cause: Thoroughly review the original SQL query and its associated error message to identify the root cause of the SQL error. 2. Propose a solution: Based on your analysis in step 1, propose a suitable solution to address the identified error and align the query with the task's requirements. 3. Generate the corrected SQL: Implement the necessary modifications to the SQL query according to the solution proposed in step 2. Ensure that the corrected SQL query obtains the desired outcomes of the task, while aligning with the subsequent specific criteria: - Use the table.column format when multiple tables are involved in the SQL query to avoid ambiguity - Enclose all table names and column names with backticks (`) to ensure correct referencing and avoid any syntax issues. The Output should in JSON format as described below: { "cause": , "solution": , "rewritten_sql": } Only output json, do not output anything else. Json: """%{"tables_schema": ("\n\n".join(tables_schema)), "tables_relationship": ("\n".join(tables_relationship)), "sql_query": sql_query, "sql_error": sql_error, "database_type": database_type, **task} extra_info = { "module": ModuleType.TASK_SQL_SELF_REFINE, "session_id": self.session_context_id, "query_id":query_id, "task_id": task.get("task_id") } if self.binded_thread_id: message = "Identify and fix an SQL query error for the task: %(clarified_task)s"%{**task} pattern = r'---\nTask:.*?\n---\n\n' instructions = re.sub(pattern, '', task_prompt, flags=re.DOTALL) resp_json, input_message_id, output_message_id = ai_assistant_chat(AssistantName.EDA_DEFAULT, self.binded_thread_id, message, ModuleType.TASK_SQL_SELF_REFINE, "json", instructions=instructions, ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, sql self refine in conversation mode, finished, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) if not resp_json.get("rewritten_sql"): # Gen result has problem, need to delete ai_assistant_delete_message(input_message_id, output_message_id, self.binded_thread_id, extra_info=extra_info) else: resp_json = ai_completion("text2sql", task_prompt, "json", ai_model=ai_model, extra_info=extra_info) logging.info("session_id = %s, query_id = %s, SQL self refine, after, model = %s, resp = %s, time_cost = %s", str(self.session_context_id), str(query_id), ai_model, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()) - time_before)) # post processor optimized_sql_query = resp_json.get("rewritten_sql", sql_query) return optimized_sql_query # give recommendation and chart type def task_data2chart(self, query_id, sql_query, columns, rows, task, ai_model=GPT_DEFAULT, language=LANGUAGE_EN): # add pre-process for single column if len(columns) == 1: logging.info("session_id = %s, query_id = %s, data2chart for single column", str(self.session_context_id), str(query_id)) return {} time_before = int(time.time()) task_prompt = """ Act as a data scientist, your objective is to recommend suitable visualization and analysis methods to gain insights for a given task. Below are the details: --- Task: %(clarified_task)s Description: %(description)s --- ### SQL Query To Collect Result Data Needed To Resolve The Task ### %(sql_query)s ### Sample Data Result Of The SQL Query ### - **Columns**: %(columns)s - **Rows**: %(rows)s ### Chart options provided ### - **PieChart**: Suitable for displaying the proportion or distribution of categorical data. Specify 'label' for categories and 'value' for their corresponding values. Option: {label: Column, value: Column}. - **LineChart**: Appropriate for visualizing the relationship between two continuous variables over a continuous range. You need to specify the 'x' and 'y' columns, and optionally a 'pivot_column'. Option: {x: Column, y: Column | Column[], pivot_column: Column}. - **BarChart**: Used to compare discrete categories or groups. Specify 'x' for categories and 'y' for values. Optionally, a 'pivot_column' can be used to display multiple series per category. Option: {x: Column, y: Column | Column[], pivot_column: Column}. - **Table**: Suitable when none of the above chart types are appropriate. Specify the columns to display. Option: {columns: Column[]}. Let's think step by step as follows: 1. Analyze the sample data result and identify the data type of each column. 2. Based on the analysis in Step 1, choose a suitable chart type from the provided options to visualize the result data for the given task. 3. Verify if the chosen chart type exists in the provided options. If not, replace it with a suitable chart type from the provided options. 4. Suggest the most appropriate data analysis method to gain insights from the data, with a brief explanation of its application. The Output should in JSON format as described below: { "data_type": [{"name": , "type": }, ...], "chosen_chart": {"chart_type":