lei_y601/sqlbot_agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'dev' of gitlab-devt.yced.com.cn:lei_y601/sqlbot_agent into dev

Browse Source 
# Conflicts:
#	main_service.py
#	service/cus_vanna_srevice.py
#	util/load_ddl_doc.py
This commit is contained in:
yujj128
2025-09-24 14:49:37 +08:00
parent 2533126e36 c85db61070
commit d5decebe33
6 changed files with 195 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
.env
View File

@@ -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

12
main_service.py
View File

@@ -1,5 +1,4 @@
from email.policy import default
import dmPython
import logging
from logging_config import LOGGING_CONFIG
@@ -21,12 +20,15 @@ def connect_database(vn):
port=int(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 == 'dameng':
vn.connect_to_dameng( )
elif db_type == 'postgresql':
# 待补充
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

3
requirement.txt
View File

@@ -1,4 +1,5 @@
vanna ==0.7.9
vanna[openai]
vanna[qdrant]
python-decouple==3.8
python-decouple==3.8
dmPython==2.5.22

65
service/cus_vanna_srevice.py
View File

@@ -1,6 +1,6 @@
from email.policy import default
from typing import List
import dmPython
from typing import List, Union
import dmPython
import orjson
import pandas as pd
from vanna.base import VannaBase
@@ -22,7 +22,7 @@ class OpenAICompatibleLLM(VannaBase):
# default parameters - can be overrided using config
self.temperature = 0.5
self.max_tokens = 5000
self.conn = None
if "temperature" in config_file:
self.temperature = config_file["temperature"]
@@ -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}
@@ -140,21 +180,6 @@ class OpenAICompatibleLLM(VannaBase):
return response.choices[0].message.content
# def connect_to_dameng(self, host, port, username, password, database):
# try:
# self.conn = dmPython.connect(
# user=username,
# password=password,
# server=host,
# port=port, # 达梦默认端口5236
# autoCommit=True
# )
# print("达梦数据库连接成功")
# return True
# except Exception as e:
# print(f"连接失败: {e}")
# return False
def generate_sql_2(self, question: str, allow_llm_to_see_data=False, **kwargs) -> dict:
try:
question_sql_list = self.get_similar_question_sql(question, **kwargs)
@@ -164,7 +189,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='中文',
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)
@@ -180,7 +205,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'),
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(

130
util/load_ddl_doc.py
View File

@@ -1,135 +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 = [
"""
<人员库表注意事项>
<rule>
查询address时,尽量使用like查询，如:select * from 人员库 where address like '%张三%';
语法为mysql语法;
如果涉及下面<info>中的字段需要展示给用户看时请替换成相关代表
birthday 字段涉及计算时，请转化为合理格式计算
</rule>
<info>
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代表更新
</info>
</人员库表注意事项>
""",
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)

124
util/train_ddl.py Normal file
View File

@@ -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>
'''