feat: Vector retrieval matches tables

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

@@ -1,22 +1,18 @@
 import datetime
 import json
-import traceback
 from typing import List, Optional
 
 from fastapi import HTTPException
 from sqlalchemy import and_, text
 from sqlmodel import select
 
-from apps.ai_model.embedding import EmbeddingModelCache
 from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
-from apps.datasource.embedding.ds_embedding import cosine_similarity
+from apps.datasource.embedding.table_embedding import get_table_embedding
 from apps.datasource.utils.utils import aes_decrypt
 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 SQLBotLogUtil
 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
@@ -389,30 +385,3 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
     for s in tables:
         schema_str += s
     return schema_str
-
-
-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/apps/datasource/embedding/ds_embedding.py

@@ -1,34 +1,19 @@
 # Author: Junjun
 # Date: 2025/9/18
 import json
-import math
 import traceback
 from typing import Optional
 
 from apps.ai_model.embedding import EmbeddingModelCache
 from apps.datasource.crud.datasource import get_table_schema
+from apps.datasource.embedding.utils import cosine_similarity
 from apps.datasource.models.datasource import CoreDatasource
 from apps.system.crud.assistant import AssistantOutDs
 from common.core.deps import CurrentAssistant
 from common.core.deps import SessionDep, CurrentUser
 from common.utils.utils import SQLBotLogUtil
 
 
-def cosine_similarity(vec_a, vec_b):
-    if len(vec_a) != len(vec_b):
-        raise ValueError("The vector dimension must be the same")
-
-    dot_product = sum(a * b for a, b in zip(vec_a, vec_b))
-
-    norm_a = math.sqrt(sum(a * a for a in vec_a))
-    norm_b = math.sqrt(sum(b * b for b in vec_b))
-
-    if norm_a == 0 or norm_b == 0:
-        return 0.0
-
-    return dot_product / (norm_a * norm_b)
-
-
 def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, out_ds: AssistantOutDs,
                      question: str,
                      current_assistant: Optional[CurrentAssistant] = None):
@@ -45,7 +30,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, embedding=False)
+                table_schema = get_table_schema(session, current_user, ds, question, 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})

backend/apps/datasource/embedding/table_embedding.py

@@ -0,0 +1,37 @@
+# Author: Junjun
+# Date: 2025/9/23
+import json
+import traceback
+
+from apps.ai_model.embedding import EmbeddingModelCache
+from apps.datasource.embedding.utils import cosine_similarity
+from common.core.config import settings
+from common.core.deps import SessionDep, CurrentUser
+from common.utils.utils import SQLBotLogUtil
+
+
+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/apps/datasource/embedding/utils.py

@@ -0,0 +1,18 @@
+# Author: Junjun
+# Date: 2025/9/23
+import math
+
+
+def cosine_similarity(vec_a, vec_b):
+    if len(vec_a) != len(vec_b):
+        raise ValueError("The vector dimension must be the same")
+
+    dot_product = sum(a * b for a, b in zip(vec_a, vec_b))
+
+    norm_a = math.sqrt(sum(a * a for a in vec_a))
+    norm_b = math.sqrt(sum(b * b for b in vec_b))
+
+    if norm_a == 0 or norm_b == 0:
+        return 0.0
+
+    return dot_product / (norm_a * norm_b)