feat: Vector retrieval matches tables

Author: XiaJunjie2020

backend/apps/chat/task/llm.py

@@ -406,7 +406,7 @@ def select_datasource(self):
         full_thinking_text = ''
         full_text = ''
         if not ignore_auto_select:
-            if settings.EMBEDDING_ENABLED:
+            if settings.TABLE_EMBEDDING_ENABLED:
                 ds = get_ds_embedding(self.session, self.current_user, _ds_list, self.out_ds_instance,
                                       self.chat_question.question, self.current_assistant)
                 yield {'content': '{"id":' + str(ds.get('id')) + '}'}
@@ -496,7 +496,7 @@ def select_datasource(self):
         except Exception as e:
             _error = e
 
-        if not ignore_auto_select and not settings.EMBEDDING_ENABLED:
+        if not ignore_auto_select and not settings.TABLE_EMBEDDING_ENABLED:
             self.record = save_select_datasource_answer(session=self.session, record_id=self.record.id,
                                                         answer=orjson.dumps({'content': full_text}).decode(),
                                                         datasource=_datasource,

backend/apps/datasource/crud/datasource.py

@@ -13,6 +13,7 @@
 from apps.db.constant import DB
 from apps.db.db import get_tables, get_fields, exec_sql, check_connection
 from apps.db.engine import get_engine_config, get_engine_conn
+from common.core.config import settings
 from common.core.deps import SessionDep, CurrentUser, Trans
 from common.utils.utils import deepcopy_ignore_extra
 from .table import get_tables_by_ds_id
@@ -402,7 +403,7 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
         all_tables.append(t_obj)
 
     # do table embedding
-    if embedding and tables:
+    if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
         tables = get_table_embedding(session, current_user, tables, question)
     # splice schema
     if tables:

backend/common/core/config.py

@@ -105,6 +105,7 @@ def SQLALCHEMY_DATABASE_URI(self) -> PostgresDsn | str:
     PG_POOL_RECYCLE: int = 3600
     PG_POOL_PRE_PING: bool = True
 
+    TABLE_EMBEDDING_ENABLED = False
     TABLE_EMBEDDING_COUNT: int = 10