import time import json from collections import deque from typing import Union from consts.consts import MODE_EXPLORE, MODE_REFINE, LANGUAGE_EN, LANGUAGE_CN, GPT_DEFAULT, GPT_4, MODE_TYPE_LENGJING, GenSQLErrorMsg, VectorMode, ModuleType from errors.errors import AlgorithmException, BadRequestException, SystemException from models.user_preference import UserPreference from svc.dataapi import get_sql_executor, DBExecutorABC from utils.local import set_local_metrics from utils.logging import logging from svc.openai import ai_completion from .context import context from concurrent.futures import ThreadPoolExecutor, as_completed from svc.executor import ThreadPoolWithLocalExecutor from domain_info.lengjing_v2 import get_domain_column_description # configure log # logging.basicConfig(format="%(asctime)s, %(funcName)s, %(levelname)s, %(message)s", datefmt="%Y-%m-%d %H:%M:%S") COLUMN_EXCLUDE_TYPE = ["JSON", "TEXT", "MEDIUMTEXT", "LONGTEXT", "LONGBLOB", "BLOB", "MEDIUMBLOB", "TINYBLOB"] def __get_table_schema(tb_name: str, db_executor: DBExecutorABC, db_name): try: time_before_for_table = int(time.time()) logging.info("db_name = %s, begin to process specific table, table = %s", db_name, tb_name) tb_context = {"table_name": tb_name.lower(), "description": ""} columns = db_executor.get_columns(tb_name) for column in columns: column["name"] = column["name"].lower() column["type"] = str(column["type"]).split(" COLLATE ")[0] # 'description' is for user column["description"] = "" column["statistics"] = "" tb_context["columns"] = {column["name"]: column for column in columns} # Get primary key pk_info = db_executor.get_primary_key(tb_name) if len(pk_info["constrained_columns"]) != 0: tb_context["primary_key"] = pk_info["constrained_columns"][0].lower() # Get foreign keys fk_list = db_executor.get_foreign_keys(tb_name) if len(fk_list) != 0: tb_context["foreign_keys"] = [fk["constrained_columns"][0].lower() for fk in fk_list] # Get column statistics logging.info("db_name = %s, begin to get column statistics, table = %s", db_name, tb_name) df = db_executor.sql_to_dataframe("select * from `%s` limit 200" % tb_name) statistics = df.describe(include="all") statistics = statistics.fillna("") # store to dataset_context logging.info("db_name = %s, end to get column statistics, table = %s", db_name, tb_name) for column_name, column_statistics in statistics.to_dict().items(): column_statistics_filter = {key: val for key, val in column_statistics.items() if val} lower_column_name = column_name.lower() if "columns" not in tb_context or lower_column_name not in tb_context["columns"]: logging.warning("db_name = %s, column name = %s, not in tb_context", db_name, lower_column_name) continue tb_context["columns"][lower_column_name]["statistics"] = json.dumps(column_statistics_filter, ensure_ascii=False, default=str) # Store sample data to column logging.info("db_name = %s, begin to store sample data to column, table = %s", db_name, tb_name) sample_data_for_column = df.head(3).fillna("").to_dict(orient='list') # if the sample data is large, omited for column_name, column_sample_data in sample_data_for_column.items(): if tb_context["columns"].get(column_name.lower()): if len(json.dumps(column_sample_data, ensure_ascii=False, default=str)) < 180: tb_context["columns"][column_name.lower()]["sample_data"] = json.dumps(column_sample_data, ensure_ascii=False, default=str) else: tb_context["columns"][column_name.lower()]["sample_data"] = "[]" # Store sample data logging.info("db_name = %s, begin to store sample data to table, table = %s", db_name, tb_name) sample_data_list = [{k.lower(): v for k, v in row.items()} for row in df.head(5).fillna("NaN").to_dict(orient='records')] table_sample_data = { "table_name": tb_name.lower(), "sample_data_str": json.dumps(sample_data_list, ensure_ascii=False, default=str), } time_after_for_table = int(time.time()) logging.info("db_name = %s, end to process specific table, table = %s, resp = %s, time_cost = %s", db_name, tb_name, json.dumps(tb_context), str(time_after_for_table - time_before_for_table)) return fk_list, table_sample_data, tb_context, tb_name except Exception as e: logging.exception("db_name = %s, failed to get table schema, table = %s, error = %s", db_name, tb_name, str(e)) raise e def setup_data_context(db_context: dict, app_id, public_key, user_token: Union[str, None] = None) -> dict: if db_context.get("status") is not None: logging.info("db_name = %s, the data context has been set.", db_context.get("db_name")) return db_context try: time_before = int(time.time()) db_executor = get_sql_executor(db_context, app_id, public_key, user_token=user_token) tables = db_executor.get_table_names() # store the db name firstly db_context["db_name"] = db_executor.get_database_name() logging.info("db_name = %s, start to setup data context", db_context.get("db_name")) db_schema = {} tables_sample_data = {} multi_fk_infos = {} futures = [] with ThreadPoolExecutor(max_workers=5) as executor: for tb_name in tables: futures.append(executor.submit(__get_table_schema, tb_name, db_executor, db_context.get("db_name"))) for future in as_completed(futures): if future.exception(): logging.error("db_name = %s, failed to get table schema, error = %s", db_context.get("db_name"), str(future.exception())) continue fk_list, table_sample_data, tb_context, tb_name = future.result() lower_tb_name = tb_name.lower() tables_sample_data[lower_tb_name] = table_sample_data multi_fk_infos[lower_tb_name] = fk_list db_schema[lower_tb_name] = tb_context db_context["statistics"] = { "tables_count": len(db_schema), "columns_count": sum(len(db_schema[tb_name.lower()]["columns"]) for tb_name in db_schema.keys()) } # Status for data context: None ——> inited ——> Done db_context["status"] = "inited" db_context["db_schema"] = db_schema db_context["tables_sample_data"] = tables_sample_data # Get init table relationship setup_table_relationships(db_context, multi_fk_infos) time_after = int(time.time()) logging.info("db_name = %s, end to setup data context, time_cost = %s", db_context.get("db_name"), str(time_after-time_before)) return db_context except Exception as e: logging.exception("db_name = %s, failed to setup data context, error = %s", db_context.get("db_name"), str(e)) raise SystemException(GenSQLErrorMsg.GET_DATABASE_COLUMN_ERROR) def setup_table_relationships(db_context, multi_fk_infos): logging.info("db_name = %s, begin to setup init table relationships, fk_infos = %s", db_context.get("db_name", ""), json.dumps(multi_fk_infos)) fk_infos = multi_fk_infos if len(fk_infos) != 0: time_before = int(time.time()) db_context["table_relationship"] = {} for table_name, relationship_list in fk_infos.items(): if len(relationship_list) != 0: db_context["table_relationship"][table_name.lower()] = [] for relationship in relationship_list: if len(relationship["constrained_columns"]) == 0: continue referencing_info = { "referencing_table": table_name.lower(), "foreign_key_column": relationship["constrained_columns"][0].lower(), "referenced_table": relationship["referred_table"].lower(), "primary_key_column": relationship["referred_columns"][0].lower(), "orig_table_relship": True } db_context["table_relationship"][table_name.lower()].append(referencing_info) # double storage referenced_table_relationships = db_context["table_relationship"].get(relationship["referred_table"].lower(), []) referenced_table_relationships.append(referencing_info) db_context["table_relationship"][relationship["referred_table"].lower()] = referenced_table_relationships time_after = int(time.time()) logging.info("db_name = %s, end to setup init table relationships, resp = %s, time_cost = %s", db_context.get("db_name", ""), json.dumps(db_context["table_relationship"]), str(time_after-time_before)) # describe column for each table def explore_column_description(db_context: dict, table: str, mode: str, language=LANGUAGE_EN, ai_model=GPT_DEFAULT) -> dict: if not db_context.get("db_schema", {}).get(table): raise AlgorithmException("Table schema is empty: %s"%table) table_schema = db_context["db_schema"][table] if mode == MODE_EXPLORE and table_schema.get("status") is not None: return table_schema patch_column_num = 80 max_workers = 5 if ai_model == GPT_4: patch_column_num = 40 max_workers = 1 time_before = int(time.time()) logging.info("db_name = %s, begin to explore column description. model = %s, table = %s", db_context.get("db_name", ""), ai_model, table) # multi thread executor = ThreadPoolWithLocalExecutor(max_workers=max_workers) work_list = [] llm_times = 0 index = 0 patch_columns_detail = [] all_columns_detail = list(table_schema.get("columns", {}).values()) while index < len(all_columns_detail): if index == llm_times * patch_column_num + patch_column_num: llm_times += 1 work_list.append(executor.submit(explore_patch_column_description, db_context, table, patch_columns_detail, mode, language, ai_model)) patch_columns_detail = [] column_detail = "column: %(name)s, type: %(type)s"%all_columns_detail[index] # comment if mode == MODE_EXPLORE and all_columns_detail[index].get("comment"): column_detail += (", comment: %(comment)s"%all_columns_detail[index]) # refine and column description exist if mode == MODE_REFINE and all_columns_detail[index].get("description"): column_detail += (", description: %(description)s"%all_columns_detail[index]) if all_columns_detail[index].get("type") not in COLUMN_EXCLUDE_TYPE: # TODO: statistic to be added later # if all_columns_detail[index].get("statistics") != "{}": # column_detail += (", statistics: %(statistics)s"%all_columns_detail[index]) if all_columns_detail[index].get("sample_data") != "[\"NaN\", \"NaN\", \"NaN\"]": column_detail += (", sample_data: %(sample_data)s"%all_columns_detail[index]) # TODO use relationship to refine # if mode == MODE_REFINE: # for item in db_context.get("table_relationship", {}).get(table, []): # if item.get("referenced_table", "") == table and item.get("primary_key_column", "") == all_columns_detail[index].get("name"): # # label primary key. TODO: add relationship to refine # column_detail += (", is_primary_key: True.") # if item.get("referencing_table", "") == table and item.get("foreign_key_column", "") == all_columns_detail[index].get("name"): # # label foreign key. TODO: add relationship to refine # column_detail += (", is_foreign_key: True.") patch_columns_detail.append(column_detail) if index == (len(all_columns_detail) - 1) and len(patch_columns_detail) != 0: work_list.append(executor.submit(explore_patch_column_description, db_context, table, patch_columns_detail, mode, language, ai_model)) index += 1 for future in as_completed(work_list): resp = future.result() for item in resp: if table_schema["columns"].get(item["column"]): table_schema["columns"][item["column"]]["description"] = item.get("description", "") # column description of table is ready table_schema["status"] = "des_ready" context.save_dataset_context(db_context) time_after = int(time.time()) logging.info("db_name = %s, end to explore column description. model = %s, table = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, table, json.dumps(table_schema, ensure_ascii=False), str(time_after-time_before)) return table_schema # explore patch columns description def explore_patch_column_description(db_context, table, patch_columns_detail, mode, language, ai_model, metrics=None): set_local_metrics(metrics) if len(patch_columns_detail) == 0: return [] time_before = int(time.time()) logging.info("db_name = %s, begin to explore patch column description. model = %s, table = %s", db_context.get("db_name", ""), ai_model, table) # TODO: Get inner domian knowledge term_description = get_domain_column_description(db_context.get("db_name"), language=language) # Explore process main_task = """ Act as a data scientist, your task is to infer the meaning and significance of each column in a given table. Here's all the information provided: **Table Name**: %(table)s ### Table schema ### %(column_description)s ## Supplementary Explanations ### %(term_description)s Let's think step by step as follows: 1. Table Name Analysis: Begin by reviewing the table name to gain an initial understanding of the type of data it might contain. 2. Column Analysis: Along with the #Supplementary Explanations# list above(if any), proceed with a comprehensive analysis of each column's details to gain insights into what the column represents. 3. Provide Column Descriptions In %(language)s: Based on the insights from Step 2, describe each column's meaning and significance. 4. Refine Column Descriptions In %(language)s: Ensure that the column descriptions created in Step 3 are concise and specific. Eliminate redundant information like column type and column statistics, and focus solely on the most relevant details. The output should be a JSON array, as described below: [ { "column": , "description": }, ... ] Only output json, do not output anything else. Json: """%{"db_name": db_context.get("db_name", ""), "table": table, "column_description":("\n".join(patch_columns_detail)),"language": language, "term_description": term_description} resp_json = ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_EXPLORE_COLUMN}) logging.info("db_name = %s, end to explore patch column description. model = %s, table = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, table, json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) return resp_json # explore table description for each table def explore_table_description(table, mode, language=LANGUAGE_EN, ai_model=GPT_DEFAULT): db_context = context.load_dataset_context() if not db_context.get("db_schema", {}).get(table): raise AlgorithmException("Table schema is empty: %s"%table) table_schema = db_context["db_schema"][table] if mode == MODE_EXPLORE and table_schema.get("status") == "done": return table_schema patch_column_num = 80 max_workers = 5 if ai_model == GPT_4: patch_column_num = 40 max_workers = 1 time_before = int(time.time()) logging.info("db_name = %s, begin to explore table description. model = %s, table = %s", db_context.get("db_name", ""), ai_model, table) # multi thread executor = ThreadPoolWithLocalExecutor(max_workers=max_workers) work_list = [] llm_times = 0 index = 0 patch_columns_detail = [] all_columns_detail = ["column: %(name)s, type: %(type)s, description: %(description)s"%item for item in table_schema.get("columns", {}).values()] # store patch result patch_result = [] while index < len(all_columns_detail): if index == llm_times * patch_column_num + patch_column_num: llm_times += 1 work_list.append(executor.submit(explore_patch_table_description, db_context, table, patch_columns_detail, mode, language, ai_model)) patch_columns_detail = [] patch_columns_detail.append(all_columns_detail[index]) if index == (len(all_columns_detail) - 1) and len(patch_columns_detail) != 0: work_list.append(executor.submit(explore_patch_table_description, db_context, table, patch_columns_detail, mode, language, ai_model)) index += 1 for future in as_completed(work_list): patch_result.append(future.result()) if len(patch_result) == 0: return {} resp_json = patch_result[0] if len(patch_result) > 1: resp_json = merge_patch_table_description(db_context.get("db_name"), table, patch_result, language, ai_model) table_schema["description"] = resp_json.get("description", "") table_schema["type"] = resp_json.get("type", "") table_schema["key_attributes"] = resp_json.get("key_attributes", []) table_schema["entity"] = resp_json.get("entity", "") if not table_schema.get("primary_key"): table_schema["primary_key"] = resp_json["primary_key"] # None ——> des_ready ——> done table_schema["status"] = "done" context.save_dataset_context(db_context) time_after = int(time.time()) logging.info("db_name = %s, end to explore table description. model = %s, table = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, table, json.dumps(resp_json, ensure_ascii=False), str(time_after-time_before)) return resp_json # divided into multiple blocks using vertical partitioning of columns, explore each block def explore_patch_table_description(db_context, table, patch_columns_detail, mode, language, ai_model, metrics=None): set_local_metrics(metrics) time_before = int(time.time()) logging.info("db_name = %s, begin to explore patch table description. model = %s, table = %s", db_context.get("db_name", ""), ai_model, table) main_task = """ Act as a data scientist, your objective is to perform a comprehensive analysis of a given table. Here's the information provided: ** Table name **: %(table)s ### Table schema ### %(column_description)s Let's think step by step as follows: 1. Table Review: Begin by reviewing the table name and each column's details to gain insights and understand the information represented by the table. 2. Table Description: Based on the insights gained in Step 1, describe the purpose and meaning of the given table in %(language)s. Do not emphasize what columns it includes in your description and be user-friendly. 3. Primary Key Inference: Determine which column serves as the primary key. If there are multiple potential primary key columns, only choose the most possible one. 4. Key Attributes Identification: Identify and list up to a maximum of 5 key attributes from the table that play a vital role in understanding its purpose and content. 5. Table Type Analysis: Analyze the table with regards to OLAP (Online Analytical Processing) and categorize it as one of the following: Dimension, Bridge, or Fact. 6. Entity Inference: Determine the primary subject matter or entity that the table focuses on. The output should be a JSON object, as described below: { "insights": , "description": , "primary_key": , "key_attributes": [, ...], "type": , "entity": } Only output json, do not output anything else. Json: """ % {"db_name": db_context.get("db_name", ""), "table": table, "column_description":"\n".join(patch_columns_detail), "language": language} resp_json = ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_EXPLORE_TABLE}) logging.info("db_name = %s, end to explore patch table description. model = %s, table = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, table, json.dumps(resp_json, ensure_ascii=False), str(int(time.time())-time_before)) return resp_json # merge the inference of each block def merge_patch_table_description(db_name, table, patch_result, language, ai_model): time_before = int(time.time()) logging.info("db_name = %s, begin to merge patch table description. model = %s, table = %s", db_name, ai_model, table) patch_result_detail = [] for item in patch_result: item["key_attributes_str"] = json.dumps(item.get("key_attributes", [])) patch_result_detail.append("Description = %(description)s, Key_attributes = %(key_attributes_str)s, Type = %(type)s, Entity_belonged = %(entity)s, Primary_key = %(primary_key)s"%item) main_task = """ As a data scientist, your objective is to create a comprehensive understanding of a given table, which has been divided into multiple blocks using vertical partitioning of columns. The information for each block is provided below: ** Table Name **: %(table)s ### Information For Each Block ### %(patch_result_detail)s Let's think step by step as follows: 1. Table Description: Write a detailed description of the table's purpose and its significance in %(language)s. Ensure that the description includes details of each block. 2. Table Type Analysis: Analyze the type in each block and determine the final type for the given table. Categorize the table as one of the following: Dimension, Bridge, or Fact, based on its usage and role in data analysis. 3. Entity Inference: Analyze the entity associated with each block and determine the primary subject matter or entity that the table focuses on. 4. Primary Key Inference: Analyze the primary key noted by each block and determine which column serves as the primary key. If there are multiple potential primary key columns, choose the most relevant one based on its importance and uniqueness within the table. 5. Key Attributes Merge: Combine all key attributes from the individual blocks to form the table's key attributes. The output should be a JSON object, as described below: { "description":
, "type": , "entity": , "primary_key": , "key_attributes": [, ...] } Only output json, do not output anything else. Json: """%{"table": table, "patch_result_detail":"\n".join(patch_result_detail), "language": language} resp_json = ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_EXPLORE_TABLE}) # post progress if resp_json.get("table_description"): resp_json["description"] = resp_json.get("table_description") if resp_json.get("table_type"): resp_json["type"] = resp_json.get("table_type") logging.info("db_name = %s, end to merge table descriptions. model = %s, table = %s, resp = %s, time_cost = %s", str(db_name), ai_model, table, json.dumps(resp_json, ensure_ascii=False), str(int(time.time()-time_before))) return resp_json # explore table relationship for each table def explore_tables_relationship(mode, similar_count, ai_model=GPT_DEFAULT): db_context = context.load_dataset_context() # if not db_context.get("status"): # raise AlgorithmException("Database schema is empty.") if not db_context.get("db_schema"): logging.info("db_name = %s, fail to explore table relationships, there is no tables", db_context.get("db_name", "")) return {} if len(db_context.get("table_relationship", {})) != 0: logging.info("db_name = %s, end to explore table relationships, table relationships has exist", db_context.get("db_name", "")) return db_context["table_relationship"] time_before = int(time.time()) logging.info("db_name = %s, begin to explore table relationships. model = %s", db_context.get("db_name", ""), ai_model) table_relationships = {} # Get vector storage vector_storage = context.init_vector_storage(VectorMode.TABLE_SCHEMA_VECTOR, "explore_tables_relationship_"+str(time_before)) # multi thread max_workers = 5 if ai_model == GPT_4: max_workers = 1 executor = ThreadPoolWithLocalExecutor(max_workers=max_workers) work_list = [] for table, _ in db_context["db_schema"].items(): table_schema = db_context["db_schema"].get(table, {}) document = "%(table_name)s: %(description)s, key_attributes = %(key_attributes)s"%{ "table_name": table, "description": table_schema.get("description", ""), "key_attributes": json.dumps(table_schema.get("key_attributes", "")) } # query relevant tables by embedding (include self) relevant_tables = vector_storage.similarity_search_with_score(document, similar_count) relevant_tables_schema = [] # default value for specified table specified_table_schema = {} for item in relevant_tables: relevant_table_name = item[0].metadata["table"] columns = [{"name": column["name"], "type": column["type"], "description": column["description"]} for column in db_context["db_schema"][relevant_table_name]["columns"].values()] if relevant_table_name != table: # todo: statistic is important relevant_table_schema = db_context["db_schema"].get(relevant_table_name, {}) primary_key = relevant_table_schema.get("primary_key", "") primary_key_type = relevant_table_schema.get("columns", {}).get(primary_key, {}).get("type", "") primary_key_discription = relevant_table_schema.get("columns", {}).get(primary_key, {}).get("description", "") table_template = "table: %(table_name)s, description: %(description)s, primary_key_column: { name: %(primary_key)s, type: %(primary_key_type)s, description: %(primary_key_discription)s }"%{ "table_name": relevant_table_name, "description": relevant_table_schema.get("description", ""), "primary_key": primary_key, "primary_key_type": primary_key_type, "primary_key_discription": primary_key_discription} relevant_tables_schema.append(table_template) else: specified_table_schema = {"table": table, "primary_key_column": table_schema.get("primary_key", ""), "columns": columns} logging.info("db_name = %s, get most relevant tables for specific table, table = %s, relevant tables = %s", db_context.get("db_name", ""), table, json.dumps([item[0].metadata["table"] for item in list(relevant_tables)])) # add to threadpool work_list.append(executor.submit(explore_relationship, db_context.get("db_name", ""), relevant_tables_schema, specified_table_schema, table, ai_model)) # clear collections vector_storage.delete_collection() # post process for future in as_completed(work_list): resp = future.result() for relationship in resp: # foreign key is null if not relationship.get("foreign_key_column"): continue # Both referenced table and referencing table should exist in db schema all_tables = db_context["db_schema"].keys() if relationship["referencing_table"] not in all_tables or relationship["referenced_table"] not in all_tables: continue # foreign key should exist in the referenced table if relationship.get("foreign_key_column") not in db_context["db_schema"].get(relationship["referencing_table"], {}).get("columns", {}).keys(): continue # primary key is error if db_context["db_schema"].get(relationship["referenced_table"], {}).get("primary_key") != relationship["primary_key_column"]: continue # foreign key is equal to primary key if db_context["db_schema"].get(relationship["referencing_table"], {}).get("primary_key") == relationship["foreign_key_column"]: continue # double storage as_referencing_table = table_relationships.get(relationship["referencing_table"], []) # avoid repeat is_existed = False for item in as_referencing_table: if item["referenced_table"] == relationship["referenced_table"]: is_existed = True break if not is_existed: as_referencing_table.append(relationship) table_relationships[relationship["referencing_table"]] = as_referencing_table as_referenced_table = table_relationships.get(relationship["referenced_table"], []) is_existed = False for item in as_referenced_table: if item["referencing_table"] == relationship["referencing_table"]: is_existed = True break if not is_existed: as_referenced_table.append(relationship) table_relationships[relationship["referenced_table"]] = as_referenced_table db_context["table_relationship"] = table_relationships context.save_dataset_context(db_context) time_after = int(time.time()) logging.info("db_name = %s, end to explore table relationships. model = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, json.dumps(table_relationships), str(time_after-time_before)) return table_relationships def explore_relationship(db_name, relevant_tables_schema, specified_table_schema, table, ai_model, metrics=None): set_local_metrics(metrics) time_before = int(time.time()) logging.info("db_name = %s, begin to explore relationship for specific table. model = %s, table = %s", db_name, ai_model, table) if not specified_table_schema.get("table") or len(relevant_tables_schema) == 0: return [] # For the specified table, filter it's primary key column columns_exclude_primary = [] for column in specified_table_schema.get("columns", []): if column.get("name", "") != specified_table_schema.get("primary_key_column"): columns_exclude_primary.append(column) # The inference effect of all tables together is better than that of single table to single table main_task = """ Act as a data analyst, your task is to identify all the referential integrity relationships between a specified table and each of the other tables. Here is the information provided: ### Schema for the specified table ### ** Table **: %(table)s ** Columns(exclude the primary key column) **: %(columns)s ### Information about the other tables, including their descriptions and primary key columns ### %(relevant_tables_schema)s ### Definition of Referential Integrity Relationship ### A referential integrity relationship is established by linking the primary key column of one table (referenced table) to the foreign key column of another table (referencing table). The primary key uniquely identifies each record in the referenced table, while the foreign key in the referencing table refers to the primary key in the referenced table. Let's think step by step as follows: 1. Iterate through each column in the specified table. 2. Identify the foreign key columns in the specified table by examining their descriptions or other indicators. 3. For each foreign key column in the specified table, search for a related primary key column in the other tables based on name, description and type: - Note that the primary key column has been provided for each of the other tables. - If multiple related primary key column are found, choose the most relevant one. 4. Once a referential integrity relationship is found between a foreign key column and a primary key column: - Specify the nature of each inferred relationship(e.g., 1:1, 1:N). The output should be a JSON array, as described below: [ { "referencing_table": , "foreign_key_column": , "referenced_table": , "primary_key_column": , "relationship": }, ... ] Only output json, do not output anything else. Json: """%{"relevant_tables_schema": "\n".join(relevant_tables_schema), "table": specified_table_schema["table"], "columns": json.dumps(columns_exclude_primary)} resp = ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_EXPLORE_TABLE_RELATIONSHIPS}) logging.info("db_name = %s, end to explore relationship for specific table. model = %s, table = %s, resp = %s, time_cost = %s", db_name, ai_model, table, json.dumps(resp), str(int(time.time()-time_before))) return resp def entities_extract(language=LANGUAGE_EN, ai_model=GPT_DEFAULT) -> Union[dict, None]: db_context = context.load_dataset_context() if not db_context.get("db_schema"): logging.info("db_name = %s, fail to extract entities, there is no tables", db_context.get("db_name", "")) return {} time_before = int(time.time()) logging.info("db_name = %s, begin to extract entities. model = %s", db_context.get("db_name", ""), ai_model) # hot fix: filter for relationship_dic relationship_dic = __filter_relationships(db_context) # Sort sorted_tables = deque([]) for table, relationship in relationship_dic.items(): sorted_tables.append({"table": table, "num": len(relationship)}) # resort new_sorted_tables = deque(sorted(sorted_tables, key=lambda element: element.get("num", 0), reverse=True)) logging.info("db_name = %s, retrieval all tables based on relationships. tables_num = %s", db_context.get("db_name", ""), str(len(list(new_sorted_tables)))) # get top20 table top20_tables = __retrieval_top20_tables([table["table"] for table in list(new_sorted_tables)[:20]], list(db_context["db_schema"].keys())) logging.info("db_name = %s, retrieval top20 tables. tables = %s", db_context.get("db_name", ""), str(top20_tables)) # Get top20 important tables schema top20_tables_schema = [] for table in top20_tables: table_schema = "Table: %(table_name)s, Description: %(description)s, Key_attributes: %(key_attributes)s, Table_type: %(type)s"%{ "table_name": db_context["db_schema"][table]["table_name"], "description": db_context["db_schema"][table]["description"], "key_attributes": json.dumps(db_context["db_schema"][table]["key_attributes"]), "type": db_context["db_schema"][table]["type"] } # Add table relationship related_tables = [] for item in relationship_dic.get(table, []): if item["referencing_table"] == table and item["referenced_table"] in top20_tables: related_tables.append({"table": item["referenced_table"], "relationship": "%(referencing_table)s.%(foreign_key_column)s = %(referenced_table)s.%(primary_key_column)s"%item}) if len(related_tables) != 0: table_schema += ", Related_tables: %s"%json.dumps(related_tables) top20_tables_schema.append(table_schema) main_entities = {} if len(top20_tables_schema) == 0: # use table as entity, < 20 for table in list(db_context.get("db_schema", {}).values())[:20]: main_entities[table["table_name"]] = {"name": table["table_name"], "summary": table["description"], "attributes": table.get("key_attributes", []), "involved_tables": [table["table_name"]]} else: resp_json = entities_extract_by_ai(db_context.get("db_name", ""), top20_tables_schema, language, ai_model) # store for item in resp_json: main_entities[item["name"]] = item db_context["entity"] = main_entities context.save_dataset_context(db_context) # explore_db_description() time_after = int(time.time()) logging.info("db_name = %s, end to extract entities, resp = %s, time_cost = %s", db_context.get("db_name", ""), json.dumps(main_entities, ensure_ascii=False), str(time_after-time_before)) return main_entities def __filter_relationships(db_context): relationship_dic = {} all_tables = db_context.get("db_schema", {}).keys() logging.info("db_name = %s, begin to filter original table relationships. all_tables = %s, table_relationship = %s", db_context.get("db_name", ""), str(all_tables),json.dumps(db_context.get("table_relationship", {}))) for table, relationships in db_context.get("table_relationship", {}).items(): if not table in all_tables: continue filtered_relationships = [] for item in relationships: if item.get("referencing_table") in all_tables and item.get("referenced_table") in all_tables: filtered_relationships.append(item) relationship_dic[table] = filtered_relationships logging.info("db_name = %s, after to filter original table relationships. table_relationship = %s", db_context.get("db_name", ""), json.dumps(relationship_dic)) return relationship_dic def __retrieval_top20_tables(candidate_tables, all_tables): set_candidate_tables = set(candidate_tables) set_all_tables = set(all_tables) diff = list(set_all_tables - set_candidate_tables) if len(set_candidate_tables) < 20: return list(set_candidate_tables) + diff[:20-len(set_candidate_tables)] else: return list(set_candidate_tables) def entities_extract_by_ai(db_name, top20_tables_schema, language, ai_model): logging.info("db_name = %s, begin to extract entities by ai, model = %s", db_name, ai_model) main_task = """ Act as a data scientist, your objective is to cluster the provided tables into larger entities to improve data analysis and comprehension of the data source. ### Here are the tables ### %(tables_schema)s Let's think step by step as follows: 1. Review Details: Carefully review the details of each table to understand its purpose and content. 2. Identify Larger Entities: Identify groups of related tables that can be grouped together based on common themes or strong relationships. 3. Define Larger Entities: Assign meaningful names to the created larger entities that accurately represent the grouped tables and reflect their theme. 4. Provide Descriptions: Create comprehensive descriptions for each larger entity in %(language)s, emphasizing their purpose and the insights they contribute to a holistic understanding of the data source. 5. Identify key Attributes: List the essential attributes for each larger entity, ensuring they are derived from the key attributes of the grouped tables. Use the format: table.column to specify the source of each attribute. The output should be a JSON array, as described below: [ { "name": , "summary": , "attributes": [, ...], "involved_tables": [, ...] }, ... ] Only output json, do not output anything else. Json: """%{"tables_schema": "\n".join(top20_tables_schema), "language": language} return ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_EXTRACT_ENTITIES}) def explore_db_description(language=LANGUAGE_EN, ai_model=GPT_DEFAULT): db_context = context.load_dataset_context() if not db_context.get("entity"): logging.info("db_name = %s, fail to explore db description, there is no entity", db_context.get("db_name", "")) return "" # raise AlgorithmException("Entity extraction is not complete.") time_before = int(time.time()) logging.info("db_name = %s, begin to explore db description. model = %s", db_context.get("db_name", ""), ai_model) entities_schema = "\n".join(["Entity_name: %(name)s, Entity_description: %(summary)s, Involved_tables: %(involved_tables)s, Key_attributes: %(attributes)s"%{ "name": item["name"], "summary": item["summary"], "attributes": json.dumps(item["attributes"]), "involved_tables": json.dumps(item["involved_tables"]) } for item in db_context["entity"].values()]) main_task = """ Act as a data analyst, you are tasked with describing a data source in %(language)s that involves the following entities: --- %(entities_schema)s --- Let's think step by step as follows: 1. **Purpose**: Infer the purpose of the data source. Do not include the detail about entities in your inference. 2. **Domain**: Determine the domain of the data source. 3. **Business Impact**: Describe in simple terms how the data source aids in business operations, analytics, or decision-making processes. 4. **Real-World Example**: Provide an example of real-world scenarios where the data source can plays a pivotal role. 5. **User-Friendly Description**: Combine the results from the previous steps into a user-friendly description. Do not emphasize what it includes in your description. 6. **Summary**: Write a summary for the data source based on the description from step 5, highlighting the key points, and in no more than 10 words. The output should be a JSON object, as described below: { "purpose": string, "domain": string, "business_impact": string, "example": string, "description": string, "summary": string } Only output json, do not output anything else. Json: """%{"entities_schema": entities_schema, "language": language} resp_json = ai_completion("understand_db", main_task, "json", ai_model=ai_model, extra_info={"module": ModuleType.TASK_DESCRIBE_DB}) db_context["summary"] = resp_json["description"] db_context["short_summary"] = resp_json["summary"] # todo: get all entity name(keywords) try: keywords = [] for item in db_context["entity"].keys(): if " and " in item: key_word_list = item.split(" and ") keywords.append(key_word_list[0]) keywords.append(key_word_list[1]) else: keywords.append(item) db_context["summary_keywords"] = keywords except Exception as e: logging.info("get keywords for db summary error, reason = %s", str(e)) db_context["status"] = "done" context.save_dataset_context(db_context) time_after = int(time.time()) logging.info("db_name = %s, end to explore db description. model = %s, resp = %s, time_cost = %s", db_context.get("db_name", ""), ai_model, json.dumps(resp_json, ensure_ascii=False), str(time_after-time_before)) return db_context["summary"] def update_table_description(table, description): db_context = context.load_dataset_context() if not db_context.get("db_schema", {}).get(table): raise AlgorithmException("Table not found: %s" % table) table_schema = db_context["db_schema"][table] if table_schema["description"] != description: table_schema["description"] = description db_context["db_schema"][table] = table_schema context.save_dataset_context(db_context) def update_column_description(table, column, description): db_context = context.load_dataset_context() if not db_context.get("db_schema", {}).get(table, {}).get("columns", {}).get(column): raise AlgorithmException("Column not found: %s" % table) table_schema = db_context["db_schema"][table] if table_schema["columns"][column]["description"] != description: table_schema["columns"][column]["description"] = description db_context["db_schema"][table] = table_schema context.save_dataset_context(db_context)