feat: Vector retrieval matches tables

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -114,7 +114,8 @@ def __init__(self, current_user: CurrentUser, chat_question: ChatQuestion,
                 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 @@ def generate_recommend_questions_task(self):
         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 @@ def select_datasource(self):
                     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 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
                 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

@@ -16,6 +16,7 @@
 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

@@ -9,6 +9,7 @@
 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

@@ -101,5 +101,7 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
     PG_POOL_RECYCLE: int = 3600
     PG_POOL_PRE_PING: bool = True
 
+    TABLE_EMBEDDING_COUNT: int = 30
+
 
 settings = Settings()  # type: ignore