| from abc import ABC, abstractmethod |
| from typing import Dict, Any, List, Tuple, Optional, Union |
|
|
|
|
| class AlpacaTemplate: |
| DEFAULT_SYSTEM = ( |
| "Below is an instruction that describes a task. " |
| "Write a response that appropriately completes the request." |
| ) |
|
|
| @classmethod |
| def template( |
| cls, |
| user_content: str, |
| system_content: Union[str, None] = None, |
| response: Union[str, None] = None, |
| ) -> str: |
| template: str = "" |
|
|
| if system_content: |
| template += f"{system_content}\n\n" |
| else: |
| template += f"{cls.DEFAULT_SYSTEM}\n\n" |
|
|
| template += f"### Instruction:\n{user_content}\n\n### Response:\n" |
| if response: |
| template += response |
|
|
| return template |
|
|
|
|
| class SftPrompt(ABC): |
| @classmethod |
| @abstractmethod |
| def prompt_user_content(cls, data: Dict[str, Any]) -> str: |
| pass |
|
|
| @classmethod |
| @abstractmethod |
| def prompt_system_content(cls, data: Dict[str, Any]) -> Optional[str]: |
| pass |
|
|
| @classmethod |
| @abstractmethod |
| def prompt_target(cls, sql: str) -> str: |
| pass |
|
|
| @classmethod |
| def prompt(cls, data: Dict[str, Any]) -> Tuple[str, Optional[str], str]: |
| user_content = cls.prompt_user_content(data) |
| system_content = cls.prompt_system_content(data) |
| if "spark_sql" in data: |
| target = cls.prompt_target(data["spark_sql"]) |
| else: |
| target = None |
| return user_content, system_content, target |
|
|
|
|
| class SQLGeneratePrompt(SftPrompt): |
| @classmethod |
| def split_res( |
| cls, res: Union[str, List[str], None] |
| ) -> Union[str, List[Union[str, None]], None]: |
| if res == None: |
| return res |
| if isinstance(res, list): |
| for i in range(len(res)): |
| res[i] = cls.split_res(res[i]) |
| return res |
| else: |
| res = res.strip() |
| if not (res.startswith("```sql") and res.endswith("```")): |
| return None |
| return res[6:-3].strip() |
|
|
| @classmethod |
| def compose_extra_task_desc( |
| cls, |
| hint: Union[str, None], |
| related_question_sqls: Union[List[Dict[str, str]], None], |
| ) -> str: |
| if not hint and not related_question_sqls: |
| return "" |
| elif hint and not related_question_sqls: |
| return ( |
| "I provide generation hint, you can refer to it to help you generate.\n" |
| ) |
| elif not hint and related_question_sqls: |
| return "I provide other user queries related to the user query with the corresponding Spark SQL queries, you can refer to them to help you generate.\n" |
| else: |
| return "I provide generation hint and other user queries related to the user query with the corresponding Spark SQL queries, you can refer to them to help you generate.\n" |
|
|
| @classmethod |
| def compose_hint_content(cls, hint: Union[str, None]) -> str: |
| if hint == None: |
| return "" |
| if len(hint) == 0: |
| return "" |
| hint = hint.strip() |
| content = ( |
| "[BEGIN OF GENERATION HINT]\n" f"{hint}\n" "[END OF GENERATION HINT]\n" "\n" |
| ) |
| return content |
|
|
| @classmethod |
| def compose_related_question_sqls_content( |
| cls, related_question_sqls: Union[List[Dict[str, str]], None] |
| ) -> str: |
| if related_question_sqls == None: |
| return "" |
| if len(related_question_sqls) == 0: |
| return "" |
| content = "[BEGIN OF RELATED QUERIES]\n" |
| for i, question_sql in enumerate(related_question_sqls): |
| question, sql = question_sql["question"], question_sql["spark_sql"] |
|
|
| question = question.strip() |
| question = question.replace("\n", " ") |
| sql = sql.strip() |
|
|
| sub_content = ( |
| f"# Related Query {i+1}\n" |
| "## User Query\n" |
| f"`{question}`\n" |
| "## Spark SQL Query\n" |
| "```sql\n" |
| f"{sql}\n" |
| "```\n" |
| "\n" |
| ) |
| content = content + sub_content |
| content = content.strip() + "\n" |
| content += "[END OF RELATED QUERIES]\n\n" |
| return content |
|
|
| @classmethod |
| def extract_table_schema(cls, user_content: str) -> str: |
| start_idx = user_content.find("[BEGIN OF TABLE SCHEMAS]") + len( |
| "[BEGIN OF TABLE SCHEMAS]" |
| ) |
| end_idx = user_content.find("[END OF TABLE SCHEMAS]") |
| return user_content[start_idx:end_idx].strip() |
|
|
| @classmethod |
| def extract_user_query(cls, user_content: str) -> str: |
| start_idx = user_content.find("[BEGIN OF QUERY]\nUser Query: ") + len( |
| "[BEGIN OF QUERY]\nUser Query: " |
| ) |
| end_idx = user_content.find("[END OF QUERY]") |
| return user_content[start_idx:end_idx].strip() |
|
|
| @classmethod |
| def extract_hint(cls, user_content: str) -> Union[str, None]: |
| start_idx = user_content.find("[BEGIN OF GENERATION HINT]") + len( |
| "[BEGIN OF GENERATION HINT]" |
| ) |
| end_idx = user_content.find("[END OF GENERATION HINT]") |
| if end_idx == -1: |
| return None |
| return user_content[start_idx:end_idx].strip() |
|
|
| @classmethod |
| def extract_related_question_sqls( |
| cls, user_content: str |
| ) -> Union[List[Dict[str, str]], None]: |
| start_idx = user_content.find("[BEGIN OF RELATED QUERIES]") + len( |
| "[BEGIN OF RELATED QUERIES]" |
| ) |
| end_idx = user_content.find("[END OF RELATED QUERIES]") |
| if end_idx == -1: |
| return None |
| related_question_sqls = user_content[start_idx:end_idx].strip().split("\n\n") |
| res = [] |
| for question_sql in related_question_sqls: |
| question_start_idx = question_sql.find("`") + 1 |
| question_end_idx = question_sql.find("`\n##") |
| question = question_sql[question_start_idx:question_end_idx] |
|
|
| sql_start_idx = question_sql.find("```sql\n") + 7 |
| sql_end_idx = -3 |
| sql = question_sql[sql_start_idx:sql_end_idx] |
| res.append({"question": question, "spark_sql": sql}) |
| return res |
|
|
| @classmethod |
| def prompt_user_content(cls, data: Dict[str, Any]) -> str: |
| question, schema = data["question"].strip(), data["schema"].strip() |
| hint, related_question_sqls = ( |
| data["hint"], |
| data["related_question_sqls"], |
| ) |
| extra_task_desc = cls.compose_extra_task_desc(hint, related_question_sqls) |
| hint_content = cls.compose_hint_content(hint) |
| related_question_sqls_content = cls.compose_related_question_sqls_content( |
| related_question_sqls |
| ) |
|
|
| user_content = ( |
| "[BEGIN OF TASK INSTRUCTION]\n" |
| "You are an expert in composing Spark SQL queries. You are given a user query and a set of table schemas.\n" |
| "Based on the user query, you need to generate one Spark SQL query to achieve the purpose.\n" |
| f"{extra_task_desc}" |
| "[END OF TASK INSTRUCTION]\n" |
| "\n" |
| "[BEGIN OF TABLE SCHEMAS]\n" |
| f"{schema}\n" |
| "[END OF TABLE SCHEMAS]\n" |
| "\n" |
| f"{hint_content}" |
| f"{related_question_sqls_content}" |
| "[BEGIN OF FORMAT INSTRUCTION]\n" |
| "The output MUST strictly adhere to the following format, and NO other text MUST be included.\n" |
| "```sql\n" |
| "your output Spark SQL query\n" |
| "```\n" |
| "[END OF FORMAT INSTRUCTION]\n" |
| "\n" |
| "[BEGIN OF QUERY]\n" |
| f"User Query: {question}\n" |
| "[END OF QUERY]\n" |
| ) |
|
|
| user_content = AlpacaTemplate.template(user_content) |
| return user_content |
|
|
| @classmethod |
| def prompt_system_content(cls, data: Dict[str, Any]) -> Union[str, None]: |
| return None |
|
|
| @classmethod |
| def prompt_target(cls, sql: Union[str, List[str]]) -> Union[str, List[str]]: |
| def map_func(tmp_sql: str) -> str: |
| tmp_sql = tmp_sql.strip() |
| tmp_sql = tmp_sql.strip(";") |
| tmp_sql = tmp_sql.strip() |
| tmp_sql += ";" |
| return "```sql\n" f"{tmp_sql}\n" "```" |
|
|
| if isinstance(sql, str): |
| target = map_func(sql) |
| else: |
| target = [map_func(_sql) for _sql in sql] |
|
|
| return target |
|
|
