From 60ee7338ae46187eb068db91b09209a41afa1c26 Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Fri, 26 Sep 2025 17:02:07 +0800
Subject: [PATCH] feat: Vector retrieval matches tables

---
 backend/apps/chat/task/llm.py              | 4 ++--
 backend/apps/datasource/crud/datasource.py | 3 ++-
 backend/common/core/config.py              | 1 +
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py
index f36bca316..a30a741dd 100644
--- a/backend/apps/chat/task/llm.py
+++ b/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,
diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py
index 7c7aad381..4380b3cf4 100644
--- a/backend/apps/datasource/crud/datasource.py
+++ b/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:
diff --git a/backend/common/core/config.py b/backend/common/core/config.py
index e528d1cfd..3527d35e2 100644
--- a/backend/common/core/config.py
+++ b/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