diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 9f55b468..f990fe3c 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -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() diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 95f81a4d..4c87fb97 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -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 diff --git a/backend/apps/datasource/embedding/ds_embedding.py b/backend/apps/datasource/embedding/ds_embedding.py index d37001ae..47dcd606 100644 --- a/backend/apps/datasource/embedding/ds_embedding.py +++ b/backend/apps/datasource/embedding/ds_embedding.py @@ -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 \ No newline at end of file diff --git a/backend/common/core/config.py b/backend/common/core/config.py index 962a8a84..df1bf6f4 100644 --- a/backend/common/core/config.py +++ b/backend/common/core/config.py @@ -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