(improvement)(text2sql) add text2sql feature that only call LLM once, and add correponding configs and docs. (#102)

Co-authored-by: shaoweigong <shaoweigong@tencent.com>

chat/core/src/main/python/run_config.py

@@ -15,6 +15,7 @@ CHROMA_DB_PERSIST_DIR = 'chm_db'
 PRESET_QUERY_COLLECTION_NAME = "preset_query_collection"
 TEXT2DSL_COLLECTION_NAME = "text2dsl_collection"
 TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM = 15
+TEXT2DSL_IS_SHORTCUT = False
 
 CHROMA_DB_PERSIST_PATH = os.path.join(PROJECT_DIR_PATH, CHROMA_DB_PERSIST_DIR)
 

chat/core/src/main/python/sql/constructor.py

@@ -22,10 +22,8 @@ from util.text2vec import Text2VecEmbeddingFunction, hg_embedding
 from util.chromadb_instance import client as chromadb_client, empty_chroma_collection_2
 from run_config import TEXT2DSL_COLLECTION_NAME, TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM
 
-
 def reload_sql_example_collection(vectorstore:Chroma, 
                                 sql_examplars:List[Mapping[str, str]],
-                                schema_linking_example_selector:SemanticSimilarityExampleSelector,
                                 sql_example_selector:SemanticSimilarityExampleSelector, 
                                 example_nums:int
                                 ):
@@ -35,20 +33,16 @@ def reload_sql_example_collection(vectorstore:Chroma,
 
     print("emptied sql_examples_collection size:", vectorstore._collection.count())
 
-    schema_linking_example_selector = SemanticSimilarityExampleSelector(vectorstore=sql_examples_vectorstore, k=example_nums,
-                                                        input_keys=["question"], 
-                                                        example_keys=["table_name", "fields_list", "prior_schema_links", "question", "analysis", "schema_links"])
-
     sql_example_selector = SemanticSimilarityExampleSelector(vectorstore=sql_examples_vectorstore, k=example_nums,
-                                                        input_keys=["question"],
-                                                        example_keys=["question", "current_date", "table_name", "schema_links", "sql"])
+                                                        input_keys=["question"], 
+                                                        example_keys=["table_name", "fields_list", "prior_schema_links", "question", "analysis", "schema_links", "current_date", "sql"])
 
     for example in sql_examplars:
-        schema_linking_example_selector.add_example(example)
+        sql_example_selector.add_example(example)
 
     print("reloaded sql_examples_collection size:", vectorstore._collection.count())
 
-    return vectorstore, schema_linking_example_selector, sql_example_selector
+    return vectorstore, sql_example_selector
 
 
 sql_examples_vectorstore = Chroma(collection_name=TEXT2DSL_COLLECTION_NAME, 
@@ -57,22 +51,14 @@ sql_examples_vectorstore = Chroma(collection_name=TEXT2DSL_COLLECTION_NAME,
 
 example_nums = TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM
 
-schema_linking_example_selector = SemanticSimilarityExampleSelector(vectorstore=sql_examples_vectorstore, k=example_nums,
-                                                        input_keys=["question"], 
-                                                        example_keys=["table_name", "fields_list", "prior_schema_links", "question", "analysis", "schema_links"])
-
 sql_example_selector = SemanticSimilarityExampleSelector(vectorstore=sql_examples_vectorstore, k=example_nums,
-                                                        input_keys=["question"],
-                                                        example_keys=["question", "current_date", "table_name", "schema_links", "sql"])
+                                                        input_keys=["question"], 
+                                                        example_keys=["table_name", "fields_list", "prior_schema_links", "question", "analysis", "schema_links", "current_date", "sql"])
 
 if sql_examples_vectorstore._collection.count() > 0:
     print("examples already in sql_vectorstore")
     print("init sql_vectorstore size:", sql_examples_vectorstore._collection.count()) 
-    if sql_examples_vectorstore._collection.count() < len(sql_examplars):
-        print("sql_examplars size:", len(sql_examplars))
-        sql_examples_vectorstore, schema_linking_example_selector, sql_example_selector = reload_sql_example_collection(sql_examples_vectorstore, sql_examplars, schema_linking_example_selector, sql_example_selector, example_nums)
-        print("added sql_vectorstore size:", sql_examples_vectorstore._collection.count())
-else:
-    sql_examples_vectorstore, schema_linking_example_selector, sql_example_selector = reload_sql_example_collection(sql_examples_vectorstore, sql_examplars, schema_linking_example_selector, sql_example_selector, example_nums)
-    print("added sql_vectorstore size:", sql_examples_vectorstore._collection.count())
 
+print("sql_examplars size:", len(sql_examplars))
+sql_examples_vectorstore, sql_example_selector = reload_sql_example_collection(sql_examples_vectorstore, sql_examplars, sql_example_selector, example_nums)
+print("added sql_vectorstore size:", sql_examples_vectorstore._collection.count())

chat/core/src/main/python/sql/examples_reload_run.py

@@ -8,24 +8,22 @@ import json
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-from run_config import TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM
+from run_config import TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM, TEXT2DSL_IS_SHORTCUT
 from few_shot_example.sql_exampler import examplars as sql_examplars
-from run_config import LLMPARSER_HOST
-from run_config import LLMPARSER_PORT
+from run_config import LLMPARSER_HOST, LLMPARSER_PORT
 
 
 def text2dsl_setting_update(llm_parser_host:str, llm_parser_port:str,
-        sql_examplars:List[Mapping[str, str]], example_nums:int):
+        sql_examplars:List[Mapping[str, str]], example_nums:int, is_shortcut:bool):
 
     url = f"http://{llm_parser_host}:{llm_parser_port}/query2sql_setting_update/"
     print("url: ", url)
-    payload = {"sqlExamplars":sql_examplars, "exampleNums":example_nums}
+    payload = {"sqlExamplars":sql_examplars, "exampleNums":example_nums, "isShortcut":is_shortcut}
     headers = {'content-type': 'application/json'}
     response = requests.post(url, data=json.dumps(payload), headers=headers)
     print(response.text)
 
 
 if __name__ == "__main__":
-    arguments = sys.argv
     text2dsl_setting_update(LLMPARSER_HOST, LLMPARSER_PORT,
-                            sql_examplars, TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM)
+                            sql_examplars, TEXT2DSL_FEW_SHOTS_EXAMPLE_NUM, TEXT2DSL_IS_SHORTCUT)

chat/core/src/main/python/sql/output_parser.py

@@ -10,4 +10,36 @@ def schema_link_parse(schema_link_output):
         print(e)
         schema_link_output = None
 
-    return schema_link_output
+    return schema_link_output
+
+def combo_schema_link_parse(schema_linking_sql_combo_output: str):
+    try:
+        schema_linking_sql_combo_output = schema_linking_sql_combo_output.strip()
+        pattern = r'Schema_links:(\[.*?\])'
+        schema_links_match = re.search(pattern, schema_linking_sql_combo_output)
+
+        if schema_links_match:
+            schema_links = schema_links_match.group(1)
+        else:
+            schema_links = None
+    except Exception as e:
+        print(e)
+        schema_links = None
+
+    return schema_links
+
+def combo_sql_parse(schema_linking_sql_combo_output: str):
+    try:
+        schema_linking_sql_combo_output = schema_linking_sql_combo_output.strip()
+        pattern = r'SQL:(.*)'
+        sql_match = re.search(pattern, schema_linking_sql_combo_output)
+
+        if sql_match:
+            sql = sql_match.group(1)
+        else:
+            sql = None
+    except Exception as e:
+        print(e)
+        sql = None
+
+    return sql

chat/core/src/main/python/sql/prompt_maker.py

@@ -73,3 +73,38 @@ def sql_exampler(user_query: str,
                                                             schema_links=schema_link_str)
 
     return sql_example_prompt
+
+
+def schema_linking_sql_combo_examplar(user_query: str,
+                            domain_name: str,
+                            data_date : str,
+                            fields_list: List[str],
+                            prior_schema_links: Mapping[str,str],
+                            example_selector: SemanticSimilarityExampleSelector) -> str:
+    
+    prior_schema_links_str = '['+ ','.join(["""'{}'->{}""".format(k,v) for k,v in prior_schema_links.items()]) + ']'
+
+    example_prompt_template = PromptTemplate(input_variables=["table_name", "fields_list", "prior_schema_links", "current_date", "question", "analysis", "schema_links", "sql"],
+                                template="Table {table_name}, columns = {fields_list}, prior_schema_links = {prior_schema_links}\nCurrent_date:{current_date}\n问题:{question}\n分析:{analysis} 所以Schema_links是:\nSchema_links:{schema_links}\nSQL:{sql}")
+
+    instruction = "# 根据数据库的表结构,参考先验信息,找出为每个问题生成SQL查询语句的schema_links,再根据schema_links为每个问题生成SQL查询语句"
+
+    schema_linking_sql_combo_prompt = "Table {table_name}, columns = {fields_list}, prior_schema_links = {prior_schema_links}\nCurrent_date:{current_date}\n问题:{question}\n分析: 让我们一步一步地思考。"
+
+    schema_linking_sql_combo_example_prompt_template = FewShotPromptTemplate(
+        example_selector=example_selector,
+        example_prompt=example_prompt_template,
+        example_separator="\n\n", 
+        prefix=instruction,
+        input_variables=["table_name", "fields_list", "prior_schema_links", "current_date", "question"],
+        suffix=schema_linking_sql_combo_prompt
+        )
+
+    schema_linking_sql_combo_example_prompt = schema_linking_sql_combo_example_prompt_template.format(table_name=domain_name,
+                                                                                    fields_list=fields_list,
+                                                                                    prior_schema_links=prior_schema_links_str,
+                                                                                    current_date=data_date,
+                                                                                    question=user_query)
+    return schema_linking_sql_combo_example_prompt
+
+

chat/core/src/main/python/sql/run.py

@@ -7,32 +7,37 @@ import sys
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 
-from sql.prompt_maker import schema_linking_exampler, sql_exampler
-from sql.constructor import schema_linking_example_selector, sql_example_selector,sql_examples_vectorstore, reload_sql_example_collection
-from sql.output_parser import schema_link_parse
+from sql.prompt_maker import schema_linking_exampler, sql_exampler, schema_linking_sql_combo_examplar
+from sql.constructor import sql_examples_vectorstore, sql_example_selector, reload_sql_example_collection
+from sql.output_parser import schema_link_parse, combo_schema_link_parse, combo_sql_parse
 
 from util.llm_instance import llm
-
+from run_config import TEXT2DSL_IS_SHORTCUT
 
 class Text2DSLAgent(object):
         def __init__(self):
             self.schema_linking_exampler = schema_linking_exampler
             self.sql_exampler = sql_exampler
 
+            self.schema_linking_sql_combo_exampler = schema_linking_sql_combo_examplar
+
             self.sql_examples_vectorstore = sql_examples_vectorstore
-            self.schema_linking_example_selector = schema_linking_example_selector
             self.sql_example_selector = sql_example_selector
 
             self.schema_link_parse = schema_link_parse
+            self.combo_schema_link_parse = combo_schema_link_parse
+            self.combo_sql_parse = combo_sql_parse
 
             self.llm = llm
 
-        def update_examples(self, sql_examplars, example_nums):
-            self.sql_examples_vectorstore, self.schema_linking_example_selector, self.sql_example_selector = reload_sql_example_collection(self.sql_examples_vectorstore, 
-                                                                                                                                           sql_examplars, 
-                                                                                                                                           self.schema_linking_example_selector, 
-                                                                                                                                           self.sql_example_selector, 
-                                                                                                                                           example_nums)
+            self.is_shortcut = TEXT2DSL_IS_SHORTCUT
+
+        def update_examples(self, sql_examples, example_nums, is_shortcut):
+            self.sql_examples_vectorstore,  self.sql_example_selector = reload_sql_example_collection(self.sql_examples_vectorstore, 
+                                                                                                    sql_examples, 
+                                                                                                    self.sql_example_selector, 
+                                                                                                    example_nums)
+            self.is_shortcut = is_shortcut
 
         def query2sql(self, query_text: str, 
             schema : Union[dict, None] = None,
@@ -53,14 +58,14 @@ class Text2DSLAgent(object):
             model_name = schema['modelName']
             fields_list = schema['fieldNameList']
 
-            schema_linking_prompt = self.schema_linking_exampler(query_text, model_name, fields_list, prior_schema_links, self.schema_linking_example_selector)
+            schema_linking_prompt = self.schema_linking_exampler(query_text, model_name, fields_list, prior_schema_links, self.sql_example_selector)
             print("schema_linking_prompt->", schema_linking_prompt)
             schema_link_output = self.llm(schema_linking_prompt)
             schema_link_str = self.schema_link_parse(schema_link_output)
             
             sql_prompt = self.sql_exampler(query_text, model_name, schema_link_str, current_date, self.sql_example_selector)
             print("sql_prompt->", sql_prompt)
-            sql_output = llm(sql_prompt)
+            sql_output = self.llm(sql_prompt)
 
             resp = dict()
             resp['query'] = query_text
@@ -69,7 +74,7 @@ class Text2DSLAgent(object):
             resp['priorSchemaLinking'] = linking
             resp['dataDate'] = current_date
 
-            resp['schemaLinkingOutput'] = schema_link_output
+            resp['analysisOutput'] = schema_link_output
             resp['schemaLinkStr'] = schema_link_str
             
             resp['sqlOutput'] = sql_output
@@ -78,5 +83,57 @@ class Text2DSLAgent(object):
 
             return resp
 
+        def query2sqlcombo(self, query_text: str, 
+            schema : Union[dict, None] = None,
+            current_date: str = None,
+            linking: Union[List[Mapping[str, str]], None] = None
+            ):
+
+            print("query_text: ", query_text)
+            print("schema: ", schema)
+            print("current_date: ", current_date)
+            print("prior_schema_links: ", linking)
+
+            if linking is not None:
+                  prior_schema_links = {item['fieldValue']:item['fieldName'] for item in linking}
+            else:
+                  prior_schema_links = {}
+
+            model_name = schema['modelName']
+            fields_list = schema['fieldNameList']
+
+            schema_linking_sql_combo_prompt = self.schema_linking_sql_combo_exampler(query_text, model_name, current_date, fields_list,
+                                                                                    prior_schema_links, self.sql_example_selector)
+            print("schema_linking_sql_combo_prompt->", schema_linking_sql_combo_prompt)
+            schema_linking_sql_combo_output = self.llm(schema_linking_sql_combo_prompt)
+
+            schema_linking_str = self.combo_schema_link_parse(schema_linking_sql_combo_output)
+            sql_str = self.combo_sql_parse(schema_linking_sql_combo_output)
+
+            resp = dict()
+            resp['query'] = query_text
+            resp['model'] = model_name
+            resp['fields'] = fields_list
+            resp['priorSchemaLinking'] = prior_schema_links
+            resp['dataDate'] = current_date
+
+            resp['analysisOutput'] = schema_linking_sql_combo_output
+            resp['schemaLinkStr'] = schema_linking_str
+            resp['sqlOutput'] = sql_str
+
+            print("resp: ", resp)
+
+            return resp
+
+        def query2sql_run(self, query_text: str, 
+            schema : Union[dict, None] = None,
+            current_date: str = None,
+            linking: Union[List[Mapping[str, str]], None] = None):
+
+            if self.is_shortcut:
+                  return self.query2sqlcombo(query_text, schema, current_date, linking)
+            else:
+                  return self.query2sql(query_text, schema, current_date, linking)
+
 text2sql_agent = Text2DSLAgent()
 

chat/core/src/main/python/supersonic_llmparser.py

@@ -51,7 +51,7 @@ async def din_query2sql(query_body: Mapping[str, Any]):
     else:
         linking = query_body['linking']
 
-    resp = text2sql_agent.query2sql(query_text=query_text,  
+    resp = text2sql_agent.query2sql_run(query_text=query_text,  
                      schema=schema, current_date=current_date, linking=linking)
 
     return resp
@@ -70,7 +70,12 @@ async def query2sql_setting_update(query_body: Mapping[str, Any]):
     else:
         example_nums = query_body['exampleNums']
 
-    text2sql_agent.update_examples(sql_examplars=sql_examplars, example_nums=example_nums)
+    if 'isShortcut' not in query_body:
+        raise HTTPException(status_code=400, detail="isShortcut is not in query_body")
+    else:
+        is_shortcut = query_body['isShortcut']
+
+    text2sql_agent.update_examples(sql_examples=sql_examplars, example_nums=example_nums, is_shortcut=is_shortcut)
 
     return "success"