sqlbot_agent/util/load_ddl_doc.py

from service.cus_vanna_srevice import CustomVanna
from util import train_ddl
from  util import q_and_a_dict,q_and_a_test1,q_a_next

import random
table_ddls = [
    train_ddl.person_database_ddl,train_ddl.person_status_ddl,
    train_ddl.person_attendance_ddl,train_ddl.person_ac_area,
    train_ddl.org_orgs_ddl

]
list_documentions = [
    train_ddl.train_document,
]


def add_ddl(vn: CustomVanna):
    for ddl in table_ddls:
        vn.add_ddl(ddl)


def add_documentation(vn: CustomVanna):
    for doc in list_documentions:
        vn.add_documentation(doc)
    dd=[]
    dd.extend(q_and_a_test1.question_and_answer)
    dd.extend(q_a_next.org_and_attend_q_a)
    random.shuffle(dd)
    # for d in dd:
    #     vn.add_question_sql(question=d['question'], sql=d['answer'])
    ## tag  待定，容易过拟合
    for d in dd:
        vn.add_question_sql(
            question=f"{d['question']}",
            sql=d['answer']
        )