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

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
@@ -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})
@@ -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: