diff --git a/.env b/.env
index 018084d..8a8988e 100644
--- a/.env
+++ b/.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
diff --git a/main_service.py b/main_service.py
index 5c18e13..ce36856 100644
--- a/main_service.py
+++ b/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=''))
- elif db_type == 'postgresql':
+ database=config('MYSQL_DATABASE_DBNAME', default=''))
+ 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()
diff --git a/service/cus_vanna_srevice.py b/service/cus_vanna_srevice.py
index 22375cd..19c7eea 100644
--- a/service/cus_vanna_srevice.py
+++ b/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)
diff --git a/util/load_ddl_doc.py b/util/load_ddl_doc.py
index a1e2d5a..1ac1fed 100644
--- a/util/load_ddl_doc.py
+++ b/util/load_ddl_doc.py
@@ -1,131 +1,19 @@
from service.cus_vanna_srevice import CustomVanna
-# table_ddls = [
-# """
-# 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语法;
-# """,
-# ]
+from util import train_ddl
+
table_ddls = [
- """
- 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='人员信息表';
- """,
+ train_ddl.ddl_sql,
]
list_documentions = [
- """
- <人员库表注意事项>
-
- 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语法;
-
-
- """,
+ 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)
-
diff --git a/util/train_ddl.py b/util/train_ddl.py
new file mode 100644
index 0000000..23adb04
--- /dev/null
+++ b/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，工号，姓名，单位以及用户问题中需要查询的字段;
+ 表字段信息以及字段枚举信息:
+
+ YJOA_APPSERVICE_DB."t_pr3rl2oj_yj_person_database"(人员信息表);
+
+ 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：无);
+
+
+ '''
\ No newline at end of file