feat:添加达梦数据库支持以及对接pre环境数据

.env

@@ -9,8 +9,9 @@ EMBEDDING_MODEL_API_KEY=sk-cjiakyfpzamtxgxitcbxvwyvaulnygmyxqpykkgngsvfqhuy
 EMBEDDING_MODEL_NAME=Qwen/Qwen3-Embedding-8B
 
 #mysql ,sqlite,pg等
-DATA_SOURCE_TYPE=mysql
+DATA_SOURCE_TYPE=dameng
-
+#数据库类型
+DB_ENGINE=达梦数据库
 #sqlite 连接信息
 SQLITE_DATABASE_URL=E://db/db_flights.sqlite
 
@@ -20,3 +21,10 @@ MYSQL_DATABASE_PORT=3306
 MYSQL_DATABASE_PASSWORD=Admin1234!
 MYSQL_DATABASE_USER=yu
 MYSQL_DATABASE_DBNAME=test
+
+
+#达梦数据库
+DAMENG_DATABASE_HOST=10.254.192.191
+DAMENG_DATABASE_PORT=5236
+DAMENG_DATABASE_PASSWORD=SYSDBA
+DAMENG_DATABASE_USER=SYSDBA

main_service.py

@@ -13,17 +13,23 @@ def connect_database(vn):
         vn.connect_to_sqlite(config('SQLITE_DATABASE_URL', default=''))
     elif db_type == 'mysql':
         vn.connect_to_mysql(host=config('MYSQL_DATABASE_HOST', default=''),
-                            port=int(config('MYSQL_DATABASE_PORT', default=3306)),
+                            port=config('MYSQL_DATABASE_PORT', default=3306),
                             user=config('MYSQL_DATABASE_USER', default=''),
                             password=config('MYSQL_DATABASE_PASSWORD', default=''),
-                            dbname=config('MYSQL_DATABASE_DBNAME', default=''))
+                            database=config('MYSQL_DATABASE_DBNAME', default=''))
-    elif db_type == 'postgresql':
+    elif db_type == 'dameng':
         # 待补充
-        pass
+        vn.connect_to_dameng(
+            host=config('DAMENG_DATABASE_HOST', default=''),
+            port=config('DAMENG_DATABASE_PORT', default=3306),
+            user=config('DAMENG_DATABASE_USER', default=''),
+            password=config('DAMENG_DATABASE_PASSWORD', default=''),
+        )
     else:
         pass
 
 
+
 def load_train_data_ddl(vn: CustomVanna):
     vn.train()
 

service/cus_vanna_srevice.py

@@ -1,6 +1,6 @@
 from email.policy import default
-from typing import List
+from typing import List, Union
-
+import  dmPython
 import orjson
 import pandas as pd
 from vanna.base import VannaBase
@@ -54,6 +54,46 @@ class OpenAICompatibleLLM(VannaBase):
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
 
+    def connect_to_dameng(
+            self,
+            host: str = None,
+            dbname: str = None,
+            user: str = None,
+            password: str = None,
+            port: int = None,
+            **kwargs
+    ):
+        conn = None
+        try:
+            conn = dmPython.connect(user=user, password=password, server=host, port=port)
+        except Exception as e:
+            raise Exception(f"Failed to connect to dameng database: {e}")
+
+        def run_sql_damengsql(sql: str) -> Union[pd.DataFrame, None]:
+            if conn:
+                try:
+                    # conn.ping(reconnect=True)
+                    cs = conn.cursor()
+                    cs.execute(sql)
+                    results = cs.fetchall()
+
+                    # Create a pandas dataframe from the results
+                    df = pd.DataFrame(
+                        results, columns=[desc[0] for desc in cs.description]
+                    )
+
+                    return df
+
+
+
+                except Exception as e:
+                    conn.rollback()
+                    raise e
+            return None
+
+        self.run_sql_is_set = True
+        self.run_sql = run_sql_damengsql
+
     def user_message(self, message: str) -> any:
         return {"role": "user", "content": message}
 
@@ -148,7 +188,7 @@ class OpenAICompatibleLLM(VannaBase):
         sql_temp = template['template']['sql']
         char_temp = template['template']['chart']
         # --------基于提示词，生成sql以及图表类型
-        sys_temp = sql_temp['system'].format(engine=config("DATA_SOURCE_TYPE",default='mysql'), lang='中文', schema=ddl_list, documentation=doc_list,
+        sys_temp = sql_temp['system'].format(engine=config("DB_ENGINE",default='mysql'), lang='中文', schema=ddl_list, documentation=doc_list,
                                              data_training=question_sql_list)
         print("sys_temp", sys_temp)
         user_temp = sql_temp['user'].format(question=question,
@@ -163,7 +203,7 @@ class OpenAICompatibleLLM(VannaBase):
         # ---------------生成图表
         char_type = get_chart_type_from_sql_answer(llm_response)
         if char_type:
-            sys_char_temp = char_temp['system'].format(engine=config("DATA_SOURCE_TYPE",default='mysql'), lang='中文', sql=sql, chart_type=char_type)
+            sys_char_temp = char_temp['system'].format(engine=config("DB_ENGINE",default='mysql'), lang='中文', sql=sql, chart_type=char_type)
             user_char_temp = char_temp['user'].format(sql=sql, chart_type=char_type, question=question)
             llm_response2 = self.submit_prompt(
                 [{'role': 'system', 'content': sys_char_temp}, {'role': 'user', 'content': user_char_temp}], **kwargs)

util/load_ddl_doc.py

@@ -1,131 +1,19 @@
 from service.cus_vanna_srevice import CustomVanna
-# table_ddls = [
+from util import train_ddl
-#     """
+
-#     create table db_user
-#         (
-#             id        integer not null
-#                 constraint db_user_pk
-#                     primary key autoincrement,
-#             user_name TEXT    not null,
-#             age       integer not null,
-#             address   TEXT,
-#             gender    integer not null,
-#             email     TEXT
-#                  )
-#     """,
-# ]
-# list_documentions = [
-#     """
-#     gender 字段 0代表女性，1代表男性;
-#     查询address时,尽量使用like查询，如:select * from db_user where address like '%北京%';
-#     语法为sqlite语法;
-#     """,
-# ]
 table_ddls = [
-    """
+    train_ddl.ddl_sql,
-    CREATE TABLE 人员库表 (
-    id VARCHAR(22) PRIMARY KEY COMMENT '主键',
-    name VARCHAR(600) DEFAULT NULL COMMENT '姓名',
-    gender VARCHAR(108) DEFAULT NULL COMMENT '性别',
-    id_card_type VARCHAR(108) DEFAULT NULL COMMENT '身份证件类型',
-    id_card VARCHAR(600) DEFAULT NULL COMMENT '身份证号码',
-    birthday VARCHAR(30) DEFAULT NULL COMMENT '出生日期',
-    native_place TEXT DEFAULT NULL COMMENT '籍贯',
-    nation TEXT DEFAULT NULL COMMENT '民族',
-    country TEXT DEFAULT NULL COMMENT '国籍',
-    residence_address TEXT DEFAULT NULL COMMENT '户籍地址',
-    highest_education VARCHAR(108) DEFAULT NULL COMMENT '最高学历',
-    highest_degree TEXT DEFAULT NULL COMMENT '最高学位',
-    graduate_school TEXT DEFAULT NULL COMMENT '毕业院校',
-    political_status TEXT DEFAULT NULL COMMENT '政治面貌',
-    phone_number TEXT DEFAULT NULL COMMENT '手机号',
-    email VARCHAR(600) DEFAULT NULL COMMENT '电子邮箱',
-    worker_id VARCHAR(200) DEFAULT NULL COMMENT '工号',
-    post TEXT DEFAULT NULL COMMENT '职务',
-    engage_post TEXT DEFAULT NULL COMMENT '现从事岗位',
-    work_unit TEXT DEFAULT NULL COMMENT '工作单位全称',
-    work_content TEXT DEFAULT NULL COMMENT '工作内容',
-    engage_contract_no VARCHAR(600) DEFAULT NULL COMMENT '从事项目合同编号',
-    engage_contract_name VARCHAR(600) DEFAULT NULL COMMENT '从事项目合同名称',
-    is_subcontractor VARCHAR(108) DEFAULT NULL COMMENT '是否分包商',
-    general_contractor_unit VARCHAR(600) DEFAULT NULL COMMENT '总包单位全称',
-    office_city TEXT DEFAULT NULL COMMENT '办公城市',
-    office_address TEXT DEFAULT NULL COMMENT '办公地点',
-    person_type TEXT DEFAULT NULL COMMENT '人员类型',
-    person_status VARCHAR(108) DEFAULT NULL COMMENT '人员状态',
-    is_internal VARCHAR(108) DEFAULT NULL COMMENT '是否内部员工',
-    internal_unit VARCHAR(108) DEFAULT NULL COMMENT '内部单位',
-    internal_dept VARCHAR(108) DEFAULT NULL COMMENT '内部部门',
-    external_unit VARCHAR(600) DEFAULT NULL COMMENT '外部单位',
-    external_dept VARCHAR(600) DEFAULT NULL COMMENT '外部部门',
-    to_dept VARCHAR(600) DEFAULT NULL COMMENT '所属处室',
-    pass_type VARCHAR(108) DEFAULT NULL COMMENT '通行证类型',
-    entry_date VARCHAR(30) DEFAULT NULL COMMENT '入场日期',
-    expected_departure_date VARCHAR(30) DEFAULT NULL COMMENT '预计离场日期',
-    expire_time DATETIME DEFAULT NULL COMMENT '失效时间',
-    verifystate INT DEFAULT NULL COMMENT '单据状态',
-    auditor VARCHAR(180) DEFAULT NULL COMMENT '终审审批人',
-    auditor1 VARCHAR(36) DEFAULT NULL COMMENT '处室负责人',
-    auditnote VARCHAR(200) DEFAULT NULL COMMENT '当前审批人',
-    procinst_id VARCHAR(36) DEFAULT NULL COMMENT '流程实例ID',
-    bizflow_id VARCHAR(36) DEFAULT NULL COMMENT '业务流id',
-    bizflowname VARCHAR(200) DEFAULT NULL COMMENT '流程名称',
-    bizflow_makebillcode VARCHAR(200) DEFAULT NULL COMMENT '单据转换规则编码',
-    bizflowinstance_id VARCHAR(36) DEFAULT NULL COMMENT '业务流实例id',
-    sourcegrand_id VARCHAR(108) DEFAULT NULL COMMENT '来源孙表id',
-    first_id VARCHAR(108) DEFAULT NULL COMMENT '来源单据主表id',
-    firstchild_id VARCHAR(108) DEFAULT NULL COMMENT '来源单据子表id',
-    firstbusiobj VARCHAR(108) DEFAULT NULL COMMENT '来源业务对象',
-    firstcode TEXT DEFAULT NULL COMMENT '来源单据号',
-    source_id VARCHAR(36) DEFAULT NULL COMMENT '上游单据主表id',
-    sourcechild_id VARCHAR(36) DEFAULT NULL COMMENT '上游单据子表id',
-    sourcebusiobj VARCHAR(36) DEFAULT NULL COMMENT '上游业务对象',
-    sourcecode VARCHAR(200) DEFAULT NULL COMMENT '上游单据号',
-    code TEXT DEFAULT NULL COMMENT '编码',
-    ytenant_id VARCHAR(64) DEFAULT NULL COMMENT '租户id',
-    photo TEXT DEFAULT NULL COMMENT '照片',
-    input_time DATETIME DEFAULT NULL COMMENT '录入时间',
-    create_time DATETIME DEFAULT NULL COMMENT '创建时间',
-    modify_time DATETIME DEFAULT NULL COMMENT '修改时间',
-    audit_time DATETIME DEFAULT NULL COMMENT '审批日期',
-    input_user VARCHAR(108) DEFAULT NULL COMMENT '录入人',
-    input_dept VARCHAR(108) DEFAULT NULL COMMENT '录入部门',
-    creator VARCHAR(60) DEFAULT NULL COMMENT '创建人',
-    modifier VARCHAR(60) DEFAULT NULL COMMENT '修改人',
-    sort INT DEFAULT NULL COMMENT '排序',
-    dr INT DEFAULT 0 COMMENT '逻辑删除：0-未删除，1-已删除',
-    DHDATASTA INT DEFAULT NULL COMMENT '推送状态',
-    pubts DATETIME DEFAULT NULL COMMENT '发布时间戳（或其他时间戳）'
-    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='人员信息表';
-    """,
 ]
 list_documentions = [
-    """
+    train_ddl.train_document,
-    <人员库表注意事项>
-        <rule>
-            person_status 字段 1代表草稿，2代表审批中，3代表制卡中，4代表已入库，5代表停用;
-            gender 字段 1代表男，2代表女
-            is_internal 字段 0代表否，1代表是
-            pass_type 字段 1代表集团公司员工，2代表借调员工，3代表借用人员，4代表外部监管人员，5代表外协服务人员，6代表工勤人员，7代表来访人员
-            person_type 字段 YG代表正式员工，PQ代表劳务派遣人员，QT代表其他柔性引进人员，WHZ代表合作单位，WLS代表临时访客，WQT代表其他外部人员
-            dr 字段 0代表否，1代表是
-            id_card_type 字段 1代表居民身份证，2代表护照，3代表港澳通行证
-            highest_education 字段 1代表初中，2代表高中，3代表中专，4代表技校，5代表职高，6代表大专，7代表本科，8代表硕士，9代表博士
-            highest_degree 字段 1代表学士学位，2代表硕士学位，3代表博士学位，4代表无
-            is_subcontractor 字段 0代表否，1代表是
-            is_sign_confidentiality_agreement 字段 0代表否，1代表是
-            DHDATASTA 字段 0代表新增 1代表更新
-            查询address时,尽量使用like查询，如:select * from 人员库 where address like '%张三%';
-            语法为mysql语法;
-        </rule>
-    </人员库表注意事项>
-    """,
 ]
+
+
 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)
-

util/train_ddl.py

@@ -0,0 +1,124 @@
+ddl_sql = """
+          create table YJOA_APPSERVICE_DB."t_pr3rl2oj_yj_person_database"
+          (
+              "person_status"                     VARCHAR(108) default 'NULL' comment '人员状态',
+              "gender"                            VARCHAR(108) default 'NULL' comment '性别(1:男,2:女)',
+              "id_card"                           VARCHAR(600) default 'NULL' comment '身份证号码',
+              "input_time"                        DATETIME(39) default NULL comment '录入时间',
+              "external_unit"                     VARCHAR(600) default 'NULL' comment '外部单位',
+              "photo"                             TEXT(max) default NULL comment '照片(二进制文本)',
+              "input_user"                        VARCHAR(108) default 'NULL' comment '录入人',
+              "is_internal"                       VARCHAR(108) default 'NULL' comment '是否内部员工(1代表是，0代表否)',
+              "phone_number"                      VARCHAR(600) default 'NULL' comment '手机号',
+              "external_dept"                     VARCHAR(600) default 'NULL' comment '福建亿榕信息技术有限公司',
+              "pass_type"                         VARCHAR(108) default 'NULL' comment '通行证类型',
+              "expire_time"                       DATETIME(39) default NULL comment '失效时间',
+              "internal_dept"                     VARCHAR(108) default 'NULL' comment '内部部门',
+              "input_dept"                        VARCHAR(108) default 'NULL' comment '录入部门',
+              "name"                              VARCHAR(600) default 'NULL' comment '姓名',
+              "internal_unit"                     VARCHAR(108) default 'NULL' comment '内部单位',
+              "person_type"                       VARCHAR(108) default 'NULL' comment '人员类型',
+              "id"                                VARCHAR(22) not null comment 'id',
+              "pubts"                             DATETIME(39) default CURRENT_TIMESTAMP,
+              "dr"                                INT          default 0 comment '逻辑删除，1代表已删除',
+              "code"                              VARCHAR(600) default 'NULL' comment '编码',
+              "ytenant_id"                        VARCHAR(64) not null
+                  constraint T_PR3RL2OJ_YJ_PERSON_DATABASE_I_I_YTENANT_ID
+                      unique comment '租户ID',
+              "sourcegrand_id"                    VARCHAR(108) default 'NULL' comment '来源表id',
+              "first_id"                          VARCHAR(108) default 'NULL' comment '来源单据主表id',
+              "firstchild_id"                     VARCHAR(108) default 'NULL' comment '来源单据子表id',
+              "firstbusiobj"                      VARCHAR(108) default 'NULL' comment '来源业务对象',
+              "firstcode"                         VARCHAR(600) default 'NULL' comment '来源单据号',
+              "verifystate"                       INT          default NULL comment '单据状态',
+              "auditor"                           VARCHAR(180) default 'NULL' comment '终审审批人',
+              "audit_time"                        DATETIME(39) default NULL comment '审批日期',
+              "auditnote"                         VARCHAR(600) default 'NULL' comment '当前审批人',
+              "procinst_id"                       VARCHAR(108) default 'NULL' comment '流程实例ID',
+              "bizflow_id"                        VARCHAR(108) default 'NULL' comment '业务流id',
+              "bizflowname"                       VARCHAR(600) default 'NULL' comment '流程名称',
+              "source_id"                         VARCHAR(108) default 'NULL' comment '上游单据主表id',
+              "sourcechild_id"                    VARCHAR(108) default 'NULL' comment '上游单据子表id',
+              "bizflowinstance_id"                VARCHAR(108) default 'NULL' comment '业务流实例id',
+              "sourcebusiobj"                     VARCHAR(108) default 'NULL' comment '上游业务对象',
+              "sourcecode"                        VARCHAR(600) default 'NULL' comment '上游单据号',
+              "bizflow_makebillcode"              VARCHAR(600) default 'NULL' comment '单据转换规则编码',
+              "create_time"                       DATETIME(39) default NULL comment '创建时间',
+              "modify_time"                       DATETIME(39) default NULL comment '修改时间',
+              "creator"                           VARCHAR(180) default 'NULL' comment '创建人',
+              "modifier"                          VARCHAR(180) default 'NULL' comment '修改人',
+              "worker_id"                         VARCHAR(200) default 'NULL' comment '工号',
+              "to_dept"                           VARCHAR(600) default 'NULL' comment '所属处室',
+              "id_card_type"                      VARCHAR(108) default 'NULL' comment '身份证件类型',
+              "native_place"                      VARCHAR(600) default 'NULL' comment '籍贯',
+              "nation"                            VARCHAR(600) default 'NULL' comment '民族',
+              "country"                           VARCHAR(600) default 'NULL' comment '国籍',
+              "residence_address"                 VARCHAR(600) default 'NULL' comment '户籍地址',
+              "highest_education"                 VARCHAR(108) default 'NULL' comment '最高学历',
+              "highest_degree"                    VARCHAR(108) default 'NULL' comment '最高学位',
+              "graduate_school"                   VARCHAR(600) default 'NULL' comment '毕业院校',
+              "birthday"                          VARCHAR(30)  default 'NULL' comment '出生日期',
+              "work_date"                         VARCHAR(30)  default 'NULL' comment '参加工作日期',
+              "political_status"                  VARCHAR(600) default 'NULL' comment '政治面貌',
+              "work_unit"                         VARCHAR(600) default 'NULL' comment '工作单位全称',
+              "post"                              VARCHAR(600) default 'NULL' comment '职务',
+              "engage_post"                       VARCHAR(600) default 'NULL' comment '现从事岗位',
+              "work_content"                      TEXT(max) default NULL comment '工作内容',
+              "engage_contract_no"                VARCHAR(600) default 'NULL' comment '从事项目合同编号',
+              "engage_contract_name"              VARCHAR(600) default 'NULL' comment '从事项目合同名称',
+              "is_subcontractor"                  VARCHAR(108) default 'NULL' comment '是否分包商',
+              "general_contractor_unit"           VARCHAR(600) default 'NULL' comment '总包单位全称',
+              "entry_date"                        VARCHAR(30)  default 'NULL' comment '入场日期',
+              "expected_departure_date"           VARCHAR(30)  default 'NULL' comment '预计离场日期',
+              "office_city"                       VARCHAR(600) default 'NULL' comment '办公城市',
+              "office_address"                    VARCHAR(600) default 'NULL' comment '办公地点',
+              "is_sign_confidentiality_agreement" VARCHAR(108) default 'NULL' comment '是否签署保密协议(1:是,0:否)',
+              "email"                             VARCHAR(600) default 'NULL' comment '电子邮箱',
+              DHDATASTA                           INT          default 0 comment '推送状态(0新增，1更新)',
+              "auditor1"                          VARCHAR(36) comment '处室负责人'
+          )
+          """
+
+
+train_document='''
+                语法为达梦数据库语法;
+                查询地址，籍贯，公司，单位时，尽量使用like查询;
+                查询人员信息时,由于数据表字段过多。只需要展示人员关键信息字段，id，工号，姓名，单位以及用户问题中需要查询的字段;
+                表字段信息以及字段枚举信息:
+                <table>
+                YJOA_APPSERVICE_DB."t_pr3rl2oj_yj_person_database"(人员信息表);
+                    <filed>
+                        person_status: 人员状态 (1:草稿,2:审批中,3:制卡中,4:已入库,5:停用);
+                        pass_type: 通行证类型 (
+                            1：集团公司员工,
+                            2：借调人员,
+                            3：借用人员,
+                            4：外部监管人员,
+                            5：外协服务人员,
+                            6：工勤人员,
+                            7：来访人员
+                        );
+                        person_type: 人员类型 (
+                            YG：正式员工,
+                            PQ：劳务派遣人员,
+                            QT：其他柔性引进人员,
+                            WHZ：合作单位,
+                            WLS：临时访客,
+                            WQT：其他外部人员
+                        );
+                        id_card_type: 身份证件类型 (1:身份证,2:护照,3:港澳通行证);
+                        highest_education: 最高学历 (
+                            1：初中
+                            2：高中
+                            3：中专
+                            4：技校
+                            5：职高
+                            6：大专
+                            7：本科
+                            8：硕士
+                            9：博士
+                        );
+                        highest_degree: 最高学位 (1：学士学位,2：硕士学位,3：博士学位,4：无);
+                    </filed>
+                </table>
+        '''