klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 09:17:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: Fix the data source-related issues in the PR

Browse Source
This commit is contained in:
ulleo
2026-05-11 15:07:15 +08:00
committed by ulleo
parent b0fd4398ff
commit e72c0bedda
6 changed files with 65 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/datasource/crud/datasource.py
+15 -7
View File
@@ -65,6 +65,7 @@ def check_name(session: SessionDep, trans: Trans, user: CurrentUser, ds: CoreDat
if ds_list is not None and len(ds_list) > 0:
raise HTTPException(status_code=500, detail=trans('i18n_ds_name_exist'))
@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.DS_ID_LIST, keyExpression="user.oid")
async def create_ds(session: SessionDep, trans: Trans, user: CurrentUser, create_ds: CreateDatasource):
ds = CoreDatasource()
@@ -490,7 +491,7 @@ def get_table_sample_data(ds: CoreDatasource, table_name: str, fields: list) ->
def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource,
table_list: list[str] = None) -> str:
table_list: list[str] = None) -> str:
"""Get sample data (3 rows) for all tables to help AI understand the data"""
table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
if len(table_objs) == 0:
@@ -508,15 +509,16 @@ def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: C
def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str,
embedding: bool = True, table_list: list[str] = None) -> str:
embedding: bool = True, table_list: list[str] = None) -> tuple[str, list]:
schema_str = ""
table_objs = get_table_obj_by_ds(session=session, current_user=current_user, ds=ds)
if len(table_objs) == 0:
return schema_str
return schema_str, []
db_name = table_objs[0].schema
schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
tables = []
all_tables = [] # temp save all tables
table_name_list = []
for obj in table_objs:
# 如果传入了table_list，则只处理在列表中的表
if table_list is not None and obj.table.table_name not in table_list:
@@ -546,13 +548,14 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
schema_table += ",\n".join(field_list)
schema_table += '\n]\n'
t_obj = {"id": obj.table.id, "schema_table": schema_table, "embedding": obj.table.embedding}
t_obj = {"id": obj.table.id, "table_name": obj.table.table_name, "schema_table": schema_table,
"embedding": obj.table.embedding}
tables.append(t_obj)
all_tables.append(t_obj)
# 如果没有符合过滤条件的表，直接返回
if not tables:
return schema_str
return schema_str, []
# do table embedding
if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
@@ -561,6 +564,7 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
if tables:
for s in tables:
schema_str += s.get('schema_table')
table_name_list.append(s.get('table_name'))
# field relation
if tables and ds.table_relation:
@@ -592,6 +596,7 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
if lost_tables:
for s in lost_tables:
schema_str += s.get('schema_table')
table_name_list.append(s.get('table_name'))
# get field dict
relation_field_ids = []
@@ -609,13 +614,16 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
for ele in all_relations:
schema_str += f"{table_dict.get(int(ele.get('source').get('cell')))}.{field_dict.get(int(ele.get('source').get('port')))}={table_dict.get(int(ele.get('target').get('cell')))}.{field_dict.get(int(ele.get('target').get('port')))}\n"
return schema_str
return schema_str, table_name_list
@cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.DS_ID_LIST, keyExpression="oid")
async def get_ws_ds(session, oid) -> list:
stmt = select(CoreDatasource.id).distinct().where(CoreDatasource.oid == oid)
db_list = session.exec(stmt).all()
return db_list
@clear_cache(namespace=CacheNamespace.AUTH_INFO, cacheName=CacheName.DS_ID_LIST, keyExpression="oid")
async def clear_ws_ds_cache(oid):
SQLBotLogUtil.info(f"ds cache for ws [{oid}] has been cleaned")
SQLBotLogUtil.info(f"ds cache for ws [{oid}] has been cleaned")
backend/apps/datasource/embedding/ds_embedding.py
+1 -1
View File
@@ -23,7 +23,7 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o
if out_ds.ds_list:
for _ds in out_ds.ds_list:
ds = out_ds.get_ds(_ds.id)
table_schema = out_ds.get_db_schema(_ds.id, question, embedding=False)
table_schema, tables = out_ds.get_db_schema(_ds.id, question, embedding=False)
ds_info = f"{ds.name}, {ds.description}\n"
ds_schema = ds_info + table_schema
_list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds})
backend/apps/datasource/embedding/table_embedding.py
+2 -2
View File
@@ -45,7 +45,7 @@ def calc_table_embedding(tables: list[dict], question: str):
for table in tables:
_list.append(
{"id": table.get('id'), "schema_table": table.get('schema_table'), "embedding": table.get('embedding'),
"cosine_similarity": 0.0})
"cosine_similarity": 0.0, "table_name": table.get('table_name')})
if _list:
try:
@@ -70,7 +70,7 @@ def calc_table_embedding(tables: list[dict], question: str):
end_time = time.time()
SQLBotLogUtil.info(str(end_time - start_time))
SQLBotLogUtil.info(json.dumps([{"id": ele.get('id'), "schema_table": ele.get('schema_table'),
"cosine_similarity": ele.get('cosine_similarity')}
"cosine_similarity": ele.get('cosine_similarity'), "table_name": ele.get('table_name')}
for ele in _list]))
return _list
except Exception: