From 86b4d381e1fb5a3118a29051a172c7e92c97cc11 Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Thu, 9 Oct 2025 12:27:17 +0800
Subject: [PATCH] feat: Vector retrieval matches tables

---
 backend/apps/datasource/crud/table.py                | 9 +++++++++
 backend/apps/datasource/embedding/table_embedding.py | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/backend/apps/datasource/crud/table.py b/backend/apps/datasource/crud/table.py
index 55153a32..b7bb4ab0 100644
--- a/backend/apps/datasource/crud/table.py
+++ b/backend/apps/datasource/crud/table.py
@@ -1,3 +1,5 @@
+import json
+import time
 import traceback
 from typing import List
 
@@ -6,6 +8,7 @@ from sqlalchemy import and_, select, update
 from apps.ai_model.embedding import EmbeddingModelCache
 from common.core.config import settings
 from common.core.deps import SessionDep
+from common.utils.utils import SQLBotLogUtil
 from ..models.datasource import CoreTable, CoreField
 
 
@@ -33,6 +36,7 @@ def run_fill_empty_table_embedding(session: SessionDep):
 
     stmt = select(CoreTable.id).where(and_(CoreTable.embedding.is_(None)))
     results = session.execute(stmt).scalars().all()
+    SQLBotLogUtil.info(json.dumps(results))
 
     save_table_embedding(session, results)
 
@@ -77,7 +81,12 @@ def save_table_embedding(session: SessionDep, ids: List[int]):
 
         model = EmbeddingModelCache.get_model()
 
+        SQLBotLogUtil.info(json.dumps(table_schema))
+        SQLBotLogUtil.info('start table embedding')
+        start_time = time.time()
         results = model.embed_documents(table_schema)
+        end_time = time.time()
+        SQLBotLogUtil.info('table embedding finished in:' + str(end_time - start_time) + 'seconds')
 
         for index in range(len(results)):
             item = results[index]
diff --git a/backend/apps/datasource/embedding/table_embedding.py b/backend/apps/datasource/embedding/table_embedding.py
index 7ba96404..ca949a5d 100644
--- a/backend/apps/datasource/embedding/table_embedding.py
+++ b/backend/apps/datasource/embedding/table_embedding.py
@@ -56,7 +56,7 @@ def calc_table_embedding(tables: list[dict], question: str):
             # results = model.embed_documents(text)
             # end_time = time.time()
             # SQLBotLogUtil.info(str(end_time - start_time))
-            results = [item.get('embedding') for item in _list]
+            results = [item.get('embedding') if item.get('embedding') else ' ' for item in _list]
 
             q_embedding = model.embed_query(question)
             for index in range(len(results)):