mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 09:17:06 +08:00
fix: Fix the data source-related issues in the PR
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user