import random import json import os import time import json from abc import ABC, abstractmethod from consts.consts import ModuleType, KnowledgeType, KnowledgeSourceType, EmbeddingModel, GPT_4 from models.domain_knowledge import DomainKnowledge, KnowledgeBase from utils.local import get_session, get_local_metrics, set_local_metrics from errors.errors import BadRequestException, NotFoundException, PermissionDeniedException from langchain_community.vectorstores.chroma import Chroma from langchain.embeddings.openai import OpenAIEmbeddings from langchain.schema import Document from sql_metadata import Parser from svc.openai import ai_completion from svc.context import context from utils.logging import logging script_dir = os.path.dirname(os.path.realpath(__file__)) TERM_SHEET = "term_sheet" FEW_SHOT_EXAMPLES = "few_shot" COMMON_KNOWLEDGE = "all" FOR_TEXT2SQL = "text2sql" FOR_RETRIEVAL_TEABLES = "retrieval_tables" SUGGEST_QUESTIONS = "suggest_questions" class KnowledgeManagement(ABC): @abstractmethod def retrieval_knowledge(self, type, module, query, meta_data={}, extra_info={}) -> str: raise NotImplementedError() class LocalKnowledgeManagement(KnowledgeManagement): def __init__(self) -> None: pass def retrieval_knowledge(self, type, module, query, meta_data={}, extra_info={}) -> str: # TO DELETE app_id = "ask-billing" logging.info("session_id = %s, query_id = %s, type = %s, app_id = %s, module = %s, query = %s, meta_data = %s, extra_data = %s, start to retrieval local knowledge", str(extra_info.get("session_id")), str(extra_info.get("query_id")), type, str(app_id), module, query, str(meta_data), str(extra_info)) kg_str = "" domain_infos = get_file_content(str(app_id)) if type == KnowledgeType.FEW_SHOT: kg_str = self.__get_few_shot_examples(domain_infos, app_id, module, query, meta_data, extra_info) if type == KnowledgeType.TERM_SHEET: kg_str = self.__get_term_sheet_for_task(domain_infos, app_id, module, query, meta_data, extra_info) if type == KnowledgeType.INSTRUCTION: if module == ModuleType.TASK_RETRIEVAL_TABLES: kg_str = self.__get_instructions_for_retrieval_tables(domain_infos, app_id, query, meta_data, extra_info) if module == ModuleType.TASK_TEXT2SQL: kg_str = self.__get_instructions_for_text2sql(domain_infos, app_id, query, meta_data, extra_info) return kg_str def __get_few_shot_examples(self, domain_infos, app_id, module, query, meta_data, extra_info) -> str: if len(domain_infos) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local examples finished, domain infos is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" # Only support for Text2SQL if module != ModuleType.TASK_TEXT2SQL: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local examples failed, module is %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), module) return "" examples = domain_infos.get(FEW_SHOT_EXAMPLES, []) if len(examples) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local examples failed, examples is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" retrieval_num = meta_data.get("num", 3) if retrieval_num == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local examples failed, retrieval_num = 0", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" necessary_tables = meta_data.get("involved_tables", []) retrieval_examples = [] for example in examples: if len(example.get("involved_tables")) == 0: retrieval_examples.append(example) else: if _is_subset(example.get("involved_tables", []), necessary_tables): retrieval_examples.append(example) if len(retrieval_examples) >= retrieval_num: break # TODO: similarity search with query if len(retrieval_examples) >= 8: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local examples success, len = %s, retrieval_examples = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(len(retrieval_examples)), str(retrieval_examples)) logging.info("session_id = %s, query_id = %s, app_id = %s, meta_data = %s, retrieval local examples success, retrieval_examples = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(meta_data), str(retrieval_examples)) # Format the result examples_info = "" if len(retrieval_examples) != 0: examples_info = "\nFor your guidance, here are some examples:\n###\n{}\n###\n".format("\n\n".join(["Task: %(task)s\nSQL: %(sql)s"%item for item in retrieval_examples])) return examples_info def __get_term_sheet_for_task(self, domain_infos, app_id, module, query, meta_data, extra_info) -> str: if len(domain_infos) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local term sheet finished, domain infos is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" desc_list = [] new_task = query.lower() nouns_desc = domain_infos.get(TERM_SHEET, []) if len(nouns_desc) != 0: for item in nouns_desc: candidate_desc = "" for noun in item.get("term", []): if noun.lower() in new_task: candidate_desc = "- {}: {}".format(noun, item.get("description")) if len(candidate_desc) != 0: desc_list.append(candidate_desc) return_desc = "" if len(desc_list) != 0: return_desc = "\n\nTerm Explanation:\n{}".format("\n".join(desc_list)) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local term sheet sucess, task = %s, desc = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), query, return_desc) return return_desc def __get_instructions_for_retrieval_tables(self, domain_infos, app_id, query, meta_data, extra_info) -> str: if len(domain_infos) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local instructions for retrieval_tables finished, domain infos is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" knowledge_list = ["%(index)s. %(item)s"%{"index": index+1, "item": item} for index, item in enumerate(domain_infos.get(COMMON_KNOWLEDGE, []) + domain_infos.get(FOR_RETRIEVAL_TEABLES, []))] knowledge = "" if len(knowledge_list) != 0: knowledge = "### Additional Guidelines:\nPlease prioritize the subsequent details as they are crucial.\n{}\n".format("\n".join(knowledge_list)) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local instructions for retrieval_tables sucess, desc = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), knowledge) return knowledge def __get_instructions_for_text2sql(self, domain_infos, app_id, query, meta_data, extra_info) -> str: if len(domain_infos) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval local instructions for text2sql finished, domain infos is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" knowledge_list = ["%(index)s. %(item)s"%{"index": index+1, "item": item} for index, item in enumerate(domain_infos.get(COMMON_KNOWLEDGE, []) + domain_infos.get(FOR_TEXT2SQL, []))] knowledge = "" if len(knowledge_list) != 0: knowledge = "### Additional Guidelines:\nPlease prioritize the subsequent details as they are crucial.\n{}\n".format("\n".join(knowledge_list)) logging.info("session_id = %s, query_id = %s, app_id= %s, retrieval local instructions for text2sql success, desc = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), knowledge) return knowledge class RemoteKnowledgeManagement(KnowledgeManagement): def __init__(self) -> None: pass def retrieval_knowledge(self, type, module, query, meta_data={}, extra_info={}) -> str: db_context = context.load_dataset_context() app_id = db_context.get("app_id", 0) # insert db_name into meta_data meta_data["db_name"] = db_context.get("db_name", "") logging.info("session_id = %s, query_id = %s, app_id = %s, db_name = %s, cluster_id = %s, type = %s, module = %s, query = %s, meta_data = %s, retrieval remote knowledge started", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), db_context.get("db_name"), db_context.get("cluster_id"), type, module, query, str(meta_data)) # Compatible with old logic knowledge_base_ids = [] # query knowledge_base_id first knowledge_base_list = KnowledgeBase.get_all_by_cluster_id_and_db_name(get_session(), db_context.get("cluster_id", ""), db_context.get("db_name", "")) if len(knowledge_base_list) != 0: knowledge_base_ids = [item.id for item in knowledge_base_list] logging.info("session_id = %s, query_id = %s, app_id = %s, db_name = %s, cluster_id = %s, retrieval remote knowledge base IDs success, result = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), db_context.get("db_name"), db_context.get("cluster_id"), str(knowledge_base_ids)) # query knowledge items records = DomainKnowledge.get_by_app_id_and_module(get_session(), type, app_id, knowledge_base_ids, module) if len(records) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, type = %s, module = %s, retrieval remote knowledge finished, query records is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), type, module) return "" # if retrieval failed, can not influence the process try: if type == KnowledgeType.FEW_SHOT: return self.__get_few_shot_examples(records, app_id, module, query, meta_data, extra_info) elif type == KnowledgeType.TERM_SHEET: return self.__get_term_sheet_for_task(records, app_id, module, query, meta_data, extra_info) elif type == KnowledgeType.INSTRUCTION: return self.__get_instructions_for_task(records, app_id, module, query, meta_data, extra_info) elif type == KnowledgeType.CHUNKED_TEXT: return self.__get_relevant_chunked_text_for_task(records, app_id, module, query, meta_data, extra_info) else: logging.info("session_id = %s, query_id = %s, app_id = %s, type = %s, module = %s, retrieval remote knowledge failed, type is not supported", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), type, module) except Exception as e: logging.info("session_id = %s, query_id = %s, app_id = %s, type = %s, module = %s, retrieval remote knowledge error, error = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), type, module, str(e)) return "" def __get_few_shot_examples(self, records, app_id, module, query, meta_data, extra_info) -> str: ''' Retrieval few-shot examples for text2sql. Example for meta_data(it set by user): { "involved_tables": [], // current relevant tables for text2sql "num": 0, // retrieval number of examples "embedding_model": "", // name of embedding model } ''' # Params check if not query: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples failed, query is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" # Only support text2sql if module != ModuleType.TASK_TEXT2SQL: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples failed, only support text2sql", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" time_before = int(time.time()) candidate_examples = [] # filter by condition for item in records: involved_tables_of_example = item[1].get("involved_tables", []) item[0]["id"] = item[2] if len(involved_tables_of_example) == 0: # if text2sql example has not noted it's involved tables, it can be used for all candidate_examples.append(item[0]) else: # if text2sql example has noted it's involved tables, it need be filtered if _is_subset(involved_tables_of_example, meta_data.get("involved_tables", [])): candidate_examples.append(item[0]) # format the result examples_info = "" if len(candidate_examples) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples failed, filtered example is empty, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(int(time.time()) - time_before)) return examples_info if len(candidate_examples) <= meta_data.get("num", 5): examples_info = "\n### For your guidance, here are some examples ###\n```\n{}\n```\n".format("\n\n".join(["Task: %(question)s\nSQL: %(answer)s"%item for item in candidate_examples])) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples success by condition, examples = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(candidate_examples), str(int(time.time()) - time_before)) # send metric self.__send_metrics(module, records, [item.get("id") for item in candidate_examples], dry_run=extra_info.get("dry_run", False)) return examples_info # embedding and similarity search document_list = [Document(page_content=item.get("question"), metadata={"answer": item.get("answer"), "id": item.get("id", 0)}) for item in candidate_examples if item.get("question") and item.get("answer")] retrieval_documents = self.__retrieval_docs(query, meta_data.get("num", 3), document_list, extra_info, embedding_model=meta_data.get("embedding_model", EmbeddingModel.ADA_002)) logging.info("session_id = %s, query_id = %s, app_id = %s, all retrieval documents of few_shot, documents = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_documents)) retrieval_examples = [{"question": item[0].page_content, "answer": item[0].metadata["answer"]} for item in retrieval_documents] if len(retrieval_examples) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples failed by similarity search, example is empty, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(int(time.time()) - time_before)) return "" logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote examples success by similarity search, example = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_examples), str(int(time.time()) - time_before)) examples_info = "\n### For your guidance, here are some examples ###\n```\n{}\n```\n".format("\n\n".join(["Task: %(question)s\nSQL: %(answer)s"%item for item in retrieval_examples])) # send metric self.__send_metrics(module, records, [item[0].metadata["id"] for item in retrieval_documents], dry_run=extra_info.get("dry_run", False)) return examples_info def __get_term_sheet_for_task(self, records, app_id, module, query, meta_data, extra_info) -> str: ''' Mainly used for query augmented. Retrieval by keyword matching by default ''' # Params check if not query: logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote term sheet failed, query is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id)) return "" time_before = int(time.time()) # format the result return_desc = "" if len(records) <= 100: # keyword matching desc_list = [] retrieved_ids = [] new_query = query.lower() for record in records: item = record[0] candidate_desc = "" for term in sorted(item.get("term", [])): if term.lower() in new_query: candidate_desc = "- {}: {}".format(term, item.get("description")) if len(candidate_desc) != 0: desc_list.append(candidate_desc) retrieved_ids.append(record[2]) if len(desc_list) != 0: # only select the top10 desc_list = desc_list[:10] return_desc = "\n{}".format("\n".join(desc_list)) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote term sheet sucess by keyword matching, task = %s, term_sheet = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), query, str(desc_list), str(int(time.time()) - time_before)) # send metrics self.__send_metrics(module, records, retrieved_ids, dry_run=extra_info.get("dry_run", False)) else: # similarity search: calculate the cosine similarity between ":" and "query" document_list = [] for record in records: item = record[0] if not item.get("description"): continue for term in item.get("term", []): content = term + ": " + item["description"] document_list.append(Document(page_content=content, metadata={"id": record[2]})) retrieval_documents = self.__retrieval_docs(query, meta_data.get("num", 10), document_list, extra_info, embedding_model=meta_data.get("embedding_model", EmbeddingModel.ADA_002)) logging.info("session_id = %s, query_id = %s, app_id = %s, all retrieval documents of term_sheet, documents = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_documents)) # TODO: the score threshold need to be optimized retrieval_examples = [("- " + item[0].page_content) for item in retrieval_documents if item[1] < 0.5] retrieved_ids = [item[0].metadata.get("id") for item in retrieval_documents if item[1] < 0.5] if len(retrieval_examples) != 0: return_desc = "\n{}".format("\n".join(retrieval_examples)) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote term sheet sucess by similarity search, task = %s, term_sheet = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), query, str(retrieval_examples), str(int(time.time()) - time_before)) # send metrics self.__send_metrics(module, records, retrieved_ids, dry_run=extra_info.get("dry_run", False)) return return_desc def __get_instructions_for_task(self, records, app_id, module, query, meta_data, extra_info) -> str: if len(records) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, module = %s, retrieval remote instructions failed, record is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), module) return "" time_before = int(time.time()) candidate_instructions = [] # filter by meta_data for item in records: item[0]["id"] = item[2] # pre-process: insert the db_name into instruction pertaining to querying database metadata dynamically if 'querying database metadata' in item[0].get("instruction") and meta_data.get("db_name"): item[0]["instruction"] += " The name of database is: %(db_name)s"%{**meta_data} if len(item[1].get("involved_tables", [])) == 0: candidate_instructions.append(item[0]) else: if _is_subset(item[1]["involved_tables"], meta_data.get("involved_tables", [])): candidate_instructions.append(item[0]) # format the result if len(candidate_instructions) <= meta_data.get("num", 10): logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote instructions success, less than retrieval num, instructions = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(candidate_instructions), str(int(time.time()) - time_before)) # send metrics self.__send_metrics(module, records, [item.get("id") for item in candidate_instructions], dry_run=extra_info.get("dry_run", False)) return "\n### Additional Instructions ###\nPlease prioritize the subsequent details as they are crucial.\n{}\n".format("\n".join(["- %s"%(item.get("instruction")) for item in candidate_instructions])) # similarity search document_list = [Document(page_content=item.get("instruction"), meta_data={"id": item.get("id")}) for item in candidate_instructions] retrieval_documents = self.__retrieval_docs(query, meta_data.get("num", 10), document_list, extra_info, embedding_model=meta_data.get("embedding_model", EmbeddingModel.ADA_002)) logging.info("session_id = %s, query_id = %s, app_id = %s, all retrieval documents of instruction, documents = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_documents)) # TODO: the score threshold need to be optimized retrieval_instructions = [item[0].page_content for item in retrieval_documents if item[1] < 0.5] retrieved_ids = [item[0].metadata.get("id") for item in retrieval_documents if item[1] < 0.5] # send metric self.__send_metrics(module, records, retrieved_ids, dry_run=extra_info.get("dry_run", False)) logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote instructions sucess by similarity search, instructions = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_instructions), str(int(time.time()) - time_before)) return "\n### Additional Instructions ###\nPlease prioritize the subsequent details as they are crucial.\n{}\n".format("\n".join(["- %s"%(value) for value in retrieval_instructions])) def __get_relevant_chunked_text_for_task(self, records, app_id, module, query, meta_data, extra_info): if len(records) == 0: logging.info("session_id = %s, query_id = %s, app_id = %s, module = %s, retrieval remote chunked text failed, record is empty", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), module) return "" time_before = int(time.time()) # format the result if len(records) <= meta_data.get("num", 5): logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote chunked text success, less than retrieval num, records = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(records), str(int(time.time()) - time_before)) return "\n### Supplementary Knowledge ###\n{}\n".format("\n".join(["%s. %s"%(index+1, item[0].get("chunked_text", "")) for index, item in enumerate(records)])) # similarity search document_list = [Document(page_content=item[0].get("chunked_text")) for item in records] retrieval_documents = self.__retrieval_docs(query, meta_data.get("num", 5), document_list, extra_info, embedding_model=meta_data.get("embedding_model", EmbeddingModel.ADA_002)) logging.info("session_id = %s, query_id = %s, app_id = %s, all retrieval document of chunked text, documents = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), str(retrieval_documents)) retrieval_examples = [item[0].page_content for item in retrieval_documents if item[1] < 0.5] logging.info("session_id = %s, query_id = %s, app_id = %s, retrieval remote chunked text sucess by similarity search, task = %s, chunked_text = %s, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(app_id), query, str(retrieval_examples), str(int(time.time()) - time_before)) if len(retrieval_examples) != 0: return "\n### Supplementary Knowledge ###\n{}\n".format("\n".join(["%s. %s"%(index+1, item) for index, item in enumerate(retrieval_examples)])) return "" def __setup_chroma(self, document_list, extra_info, embedding_model=EmbeddingModel.ADA_002) -> Chroma: time_before = int(time.time()) embeddings = OpenAIEmbeddings(model=embedding_model) vector_storage = Chroma.from_documents(documents=document_list, embedding=embeddings, collection_name=(str(extra_info.get("session_id")) + "_" + str(time_before))) logging.info("session_id = %s, query_id = %s, set up chromadb finished, time_cost = %s", str(extra_info.get("session_id")), str(extra_info.get("query_id")), str(int(time.time()) - time_before)) return vector_storage def __retrieval_docs(self, query, num, document_list, extra_info, embedding_model=EmbeddingModel.ADA_002): vector_storage = self.__setup_chroma(document_list, extra_info, embedding_model=embedding_model) retrieval_documents = vector_storage.similarity_search_with_score(query, num) # clear collection cache vector_storage.delete_collection() return retrieval_documents def __send_metrics(self, module, records, retrieved_ids, dry_run=False): if dry_run: return if len(retrieved_ids) != 0 or len(records) != 0: retrieved_knowledge = [] for item in records: for id in retrieved_ids: if item[2] == id: retrieved_knowledge.append({ "knowledge_base_id": item[4], "id": item[2], "type": item[3], "meta_data": item[1], "raw_data": item[0] }) metrics = get_local_metrics() all_knowledge_metrics = metrics.get("recalled_knowledge_data", {}) module_knowledge_metrics = all_knowledge_metrics.get(module, []) module_knowledge_metrics += retrieved_knowledge # update all_knowledge_metrics[module] = module_knowledge_metrics metrics["recalled_knowledge_data"] = all_knowledge_metrics set_local_metrics(metrics) def query_knowledge(self, session, app_id, knowledge_base_id, _type=""): domain_items = [] items = DomainKnowledge.get_all_by_app_id(session, app_id, knowledge_base_id, _type) for i in items: domain_items.append(i.to_json()) return domain_items def insert_knowledge(self, session, _type: str, app_id: int, module: str, raw_data: dict, meta_data: dict, created_by: str, knowledge_base_id: int): # The type of knowledge can not be modify if not self.__validate_knowledge(_type, raw_data): logging.info("app_id = %s, insert knowledge failed, This knowledge does not meet the content format requirements, type = %s, raw_data = %s", str(app_id), _type, str(raw_data)) raise Exception("This knowledge does not meet the content format requirements.") # The module of knowledge need to be redefine if _type == KnowledgeType.FEW_SHOT: # Currently, only few-shot for text2sql is supported. if self.__is_sql_statement(raw_data.get("answer")): module = ModuleType.TASK_TEXT2SQL # Extract the names of the tables used in the SQL.(if user not specified) meta_data["involved_tables"] = Parser(raw_data.get("answer")).tables else: logging.info("app_id = %s, insert knowledge failed, only text2SQL examples are supported and the SQL is not valid, type = %s, raw_data = %s", str(app_id), _type, str(raw_data)) raise Exception("Only Text2SQL examples are supported, and your SQL is not valid.") if _type == KnowledgeType.INSTRUCTION: # limit the length of instruction if len(raw_data.get("instruction")) > 512: raise Exception("The length of the instruction is too long. Please shorten it to within 500 characters.") module = self.__guess_insturction_for(raw_data.get("instruction"), app_id) return DomainKnowledge.create(session, raw_data, _type, app_id, module, meta_data, created_by, knowledge_base_id) def update_knowledge(self, session, id, knowledge_base_id, raw_data: dict, meta_data: dict): knowledge = DomainKnowledge.get_by_id(session, id) # record not exist if not knowledge: raise NotFoundException("The specified knowledge does not exist.") if knowledge.knowledge_base_id != knowledge_base_id: raise PermissionDeniedException("The associated knowledge_base_id does not match the current id.") # The type of knowledge can not be modify if not self.__validate_knowledge(knowledge.type, raw_data): logging.info("app_id = %s, update knowledge failed, the type of knowledge is modified, type = %s, raw_data = %s", str(knowledge.app_id), knowledge.type, str(raw_data)) raise Exception("The type of knowledge cannot be modified.") # The module of knowledge need to be redefine. module = ModuleType.ALL if knowledge.type == KnowledgeType.FEW_SHOT: # Currently, only few-shot for text2sql is supported. if self.__is_sql_statement(raw_data.get("answer")): module = ModuleType.TASK_TEXT2SQL # Extract the names of the tables used in the SQL.(if user not specified) meta_data["involved_tables"] = Parser(raw_data.get("answer")).tables else: logging.info("app_id = %s, update knowledge failed, only Text2SQL examples are supported and the SQL is not valid, type = %s, raw_data = %s", str(knowledge.app_id), knowledge.type, str(raw_data)) raise Exception("Only Text2SQL examples are supported, and your SQL is not valid.") if knowledge.type == KnowledgeType.INSTRUCTION: module = self.__guess_insturction_for(raw_data.get("instruction"), knowledge.app_id) knowledge.update_data(session, module, raw_data, meta_data) return knowledge.to_json() def delete_knowledge(self, session, id, knowledge_base_id): knowledge = DomainKnowledge.get_by_id(session, id) # record not exist if not knowledge: raise NotFoundException("The specified knowledge does not exist.") if knowledge.knowledge_base_id != knowledge_base_id: raise PermissionDeniedException("The associated knowledge_base_id does not match the current id.") DomainKnowledge.delete_by_id(session, id) def __validate_knowledge(self, _type: str, raw_data: dict) -> bool: if _type == KnowledgeType.FEW_SHOT: if raw_data.get("question") and raw_data.get("answer"): return True elif _type == KnowledgeType.TERM_SHEET: if raw_data.get("term") and raw_data.get("description"): return True elif _type == KnowledgeType.INSTRUCTION: if raw_data.get("instruction"): return True elif _type == KnowledgeType.CHUNKED_TEXT: if raw_data.get("chunked_text"): return True return False def __is_sql_statement(self, answer) -> bool: is_sql = True try: table_list = Parser(answer).tables if len(table_list) == 0: is_sql = False except Exception as e: is_sql = False return is_sql def __guess_insturction_for(self, instruction, app_id) -> str: """infer the instruction is for which module Returns: str: specified module """ time_before = int(time.time()) task_prompt = """ Act as a data analysis expert, your task is to determine which module the given instruction is most likely used for: Instruction: %(instruction)s The module options provided are: - schemaLink: this module is to map a given task to relevant tables and columns from the provided database schema. - text2sql: this module is to write an SQL query using best practices to achieve the desired results for a given task. - data2chart: this module is to recommend suitable visualization and analysis methods to gain insights for a given task. - all: All of the above 3 stages, the instruction can be used for any. The output should in a json format as described below: { "module": } Do not output anything else. Json: """%{"instruction": instruction} resp_json = ai_completion("text2sql", task_prompt, "json", ai_model=GPT_4) logging.info("app_id = %s, determine module for instruction finished, instruction = %s, resp = %s, time_cost = %s", str(app_id), instruction, str(resp_json), str(int(time.time()) - time_before)) if resp_json.get("module") == "schemaLink": return ModuleType.TASK_RETRIEVAL_TABLES elif resp_json.get("module") == "text2sql": return ModuleType.TASK_TEXT2SQL elif resp_json.get("module") == "data2chart": return ModuleType.TASK_DATA2CHART else: return ModuleType.ALL local_kg_management = LocalKnowledgeManagement() remote_kg_management = RemoteKnowledgeManagement() def get_kg_management(source_type): if source_type == KnowledgeSourceType.LOCAL: return local_kg_management elif source_type == KnowledgeSourceType.REMOTE: return remote_kg_management else: raise BadRequestException("invalid knowledge source type") def get_file_content(file_name: str) -> dict: """ Get build-in domain knowledge. Parameters: db_name: The name of db, eg: lengjing """ if not file_name: logging.info("Get local file content failed, file_name is empty") return {} path = os.path.join(script_dir, "{}.json".format(file_name)) content = {} try: with open(path) as f: content = json.load(f) except FileNotFoundError: logging.info("Get local file content failed, it is not existed, file_name = %s, path = %s", file_name, path) return content def get_local_suggest_questions(key, num, session_id = ""): if num <= 0: return [] domain_infos = get_file_content("local_suggest_questions") candidate_set = domain_infos.get(key, []) if len(candidate_set) == 0: logging.info("session_id = %s, key = %s, local suggest questions is empty", str(session_id), key) return [] if len(candidate_set) <= num: logging.info("session_id = %s, key = %s, get local suggest questions successfully, questions = %s", str(session_id), key, str(candidate_set)) return candidate_set random_nums = random.sample(range(0, len(candidate_set)), num) retrieval_set = [] for i in random_nums: retrieval_set.append(candidate_set[i]) logging.info("session_id = %s, key = %s, get local suggest questions successfully, questions = %s", str(session_id), key, str(retrieval_set)) return retrieval_set def get_local_question_result(key, question, module): result_infos = get_file_content("local_questions_result") if module == ModuleType.TASK_BREAKDOWN: return result_infos.get(key, {}).get(question.strip(), {}).get(module, []) else: return result_infos.get(key, {}).get(question.strip(), {}).get(module, {}) def _is_subset(list1, list2): return all(item in list2 for item in list1)