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

feat: Vector retrieval matches tables

Browse Source
This commit is contained in:
junjun
2025-09-23 14:53:30 +08:00
committed by Junjun
parent 5f77b82f1e
commit 487c8e358d
4 changed files with 57 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/task/llm.py
+8 -4
View File
@@ -114,7 +114,8 @@ class LLMService:
if not ds:
raise SingleMessageError("No available datasource configuration found")
chat_question.engine = (ds.type_name if ds.type != 'excel' else 'PostgreSQL') + get_version(ds)
chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds)
chat_question.db_schema = get_table_schema(session=self.session, current_user=current_user, ds=ds,
question=self.chat_question.question)
self.generate_sql_logs = list_generate_sql_logs(session=self.session, chart_id=chat_id)
self.generate_chart_logs = list_generate_chart_logs(session=self.session, chart_id=chat_id)
@@ -346,7 +347,8 @@ class LLMService:
if self.ds and not self.chat_question.db_schema:
self.chat_question.db_schema = self.out_ds_instance.get_db_schema(
self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session,
current_user=self.current_user, ds=self.ds)
current_user=self.current_user, ds=self.ds,
question=self.chat_question.question)
guess_msg: List[Union[BaseMessage, dict[str, Any]]] = []
guess_msg.append(SystemMessage(content=self.chat_question.guess_sys_question()))
@@ -506,7 +508,8 @@ class LLMService:
self.chat_question.engine = (_ds.type_name if _ds.type != 'excel' else 'PostgreSQL') + get_version(
self.ds)
self.chat_question.db_schema = get_table_schema(session=self.session,
current_user=self.current_user, ds=self.ds)
current_user=self.current_user, ds=self.ds,
question=self.chat_question.question)
_engine_type = self.chat_question.engine
_chat.engine_type = _ds.type_name
# save chat
@@ -995,7 +998,8 @@ class LLMService:
self.chat_question.db_schema = self.out_ds_instance.get_db_schema(
self.ds.id) if self.out_ds_instance else get_table_schema(session=self.session,
current_user=self.current_user,
ds=self.ds)
ds=self.ds,
question=self.chat_question.question)
else:
self.validate_history_ds()
backend/apps/datasource/crud/datasource.py
+17 -6
View File
@@ -16,6 +16,7 @@ from common.utils.utils import deepcopy_ignore_extra
from .table import get_tables_by_ds_id
from ..crud.field import delete_field_by_ds_id, update_field
from ..crud.table import delete_table_by_ds_id, update_table
from ..embedding.ds_embedding import get_table_embedding
from ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema, TableObj, \
DatasourceConf, TableAndFields
@@ -344,22 +345,25 @@ def get_table_obj_by_ds(session: SessionDep, current_user: CurrentUser, ds: Core
return _list
def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource) -> str:
def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, question: str,
embedding: bool = True) -> str:
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
db_name = table_objs[0].schema
schema_str += f"【DB_ID】 {db_name}\n【Schema】\n"
tables = []
for obj in table_objs:
schema_str += f"# Table: {db_name}.{obj.table.table_name}" if ds.type != "mysql" and ds.type != "es" else f"# Table: {obj.table.table_name}"
schema_table = ''
schema_table += f"# Table: {db_name}.{obj.table.table_name}" if ds.type != "mysql" and ds.type != "es" else f"# Table: {obj.table.table_name}"
table_comment = ''
if obj.table.custom_comment:
table_comment = obj.table.custom_comment.strip()
if table_comment == '':
schema_str += '\n[\n'
schema_table += '\n[\n'
else:
schema_str += f", {table_comment}\n[\n"
schema_table += f", {table_comment}\n[\n"
field_list = []
for field in obj.fields:
@@ -370,7 +374,14 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
field_list.append(f"({field.field_name}:{field.field_type})")
else:
field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})")
schema_str += ",\n".join(field_list)
schema_str += '\n]\n'
schema_table += ",\n".join(field_list)
schema_table += '\n]\n'
tables.append(schema_table)
# do table embedding
if embedding:
tables = get_table_embedding(session, current_user, tables, question)
# todo 外键
for s in tables:
schema_str += s
return schema_str
backend/apps/datasource/embedding/ds_embedding.py
+30 -1
View File
@@ -9,6 +9,7 @@ from apps.ai_model.embedding import EmbeddingModelCache
from apps.datasource.crud.datasource import get_table_schema
from apps.datasource.models.datasource import CoreDatasource
from apps.system.crud.assistant import AssistantOutDs
from common.core.config import settings
from common.core.deps import CurrentAssistant
from common.core.deps import SessionDep, CurrentUser
from common.utils.utils import SQLBotLogUtil
@@ -45,7 +46,7 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o
for _ds in _ds_list:
if _ds.get('id'):
ds = session.get(CoreDatasource, _ds.get('id'))
table_schema = get_table_schema(session, current_user, ds)
table_schema = get_table_schema(session, current_user, ds, 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})
@@ -71,3 +72,31 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o
return {"id": ds.id, "name": ds.name, "description": ds.description}
except Exception:
traceback.print_exc()
return _list
def get_table_embedding(session: SessionDep, current_user: CurrentUser, tables: list[str], question: str):
_list = []
for table_schema in tables:
_list.append({"table_schema": table_schema, "cosine_similarity": 0.0})
if _list:
try:
text = [s.get('table_schema') for s in _list]
model = EmbeddingModelCache.get_model()
results = model.embed_documents(text)
q_embedding = model.embed_query(question)
for index in range(len(results)):
item = results[index]
_list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item)
_list.sort(key=lambda x: x['cosine_similarity'], reverse=True)
_list = _list[:settings.TABLE_EMBEDDING_COUNT]
# print(len(_list))
SQLBotLogUtil.info(json.dumps(_list))
return [t.get("table_schema") for t in _list]
except Exception:
traceback.print_exc()
return _list
backend/common/core/config.py
+2
View File
@@ -101,5 +101,7 @@ class Settings(BaseSettings):
PG_POOL_RECYCLE: int = 3600
PG_POOL_PRE_PING: bool = True
TABLE_EMBEDDING_COUNT: int = 30
settings = Settings() # type: ignore