feat: Vector retrieval matches tables

Author: XiaJunjie2020

backend/apps/datasource/crud/datasource.py

@@ -8,13 +8,14 @@
 from sqlmodel import select
 
 from apps.datasource.crud.permission import get_column_permission_fields, get_row_permission_filters, is_normal_user
-from apps.datasource.embedding.table_embedding import get_table_embedding
+from apps.datasource.embedding.table_embedding import get_table_embedding, calc_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.embedding_threads import run_save_table_embeddings
 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
@@ -194,6 +195,9 @@ def sync_table(session: SessionDep, ds: CoreDatasource, tables: List[CoreTable])
         session.query(CoreField).filter(CoreField.ds_id == ds.id).delete(synchronize_session=False)
         session.commit()
 
+    # do table embedding
+    run_save_table_embeddings(id_list)
+
 
 def sync_fields(session: SessionDep, ds: CoreDatasource, table: CoreTable, fields: List[ColumnSchema]):
     id_list = []
@@ -232,14 +236,23 @@ def update_table_and_fields(session: SessionDep, data: TableObj):
     for field in data.fields:
         update_field(session, field)
 
+    # do table embedding
+    run_save_table_embeddings([data.table.id])
+
 
 def updateTable(session: SessionDep, table: CoreTable):
     update_table(session, table)
 
+    # do table embedding
+    run_save_table_embeddings([table.id])
+
 
 def updateField(session: SessionDep, field: CoreField):
     update_field(session, field)
 
+    # do table embedding
+    run_save_table_embeddings([field.table_id])
+
 
 def preview(session: SessionDep, current_user: CurrentUser, id: int, data: TableObj):
     ds = session.query(CoreDatasource).filter(CoreDatasource.id == id).first()
@@ -398,13 +411,13 @@ def get_table_schema(session: SessionDep, current_user: CurrentUser, ds: CoreDat
             schema_table += ",\n".join(field_list)
         schema_table += '\n]\n'
 
-        t_obj = {"id": obj.table.id, "schema_table": schema_table}
+        t_obj = {"id": obj.table.id, "schema_table": schema_table, "embedding": obj.table.embedding}
         tables.append(t_obj)
         all_tables.append(t_obj)
 
     # do table embedding
     if embedding and tables and settings.TABLE_EMBEDDING_ENABLED:
-        tables = get_table_embedding(tables, question)
+        tables = calc_table_embedding(tables, question)
     # splice schema
     if tables:
         for s in tables:

backend/apps/datasource/crud/table.py

@@ -1,6 +1,12 @@
+import traceback
+from typing import List
+
+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 ..models.datasource import CoreDatasource, CreateDatasource, CoreTable, CoreField, ColumnSchema
-from sqlalchemy import and_
+from ..models.datasource import CoreTable, CoreField
 
 
 def delete_table_by_ds_id(session: SessionDep, id: int):
@@ -19,3 +25,65 @@ def update_table(session: SessionDep, item: CoreTable):
     record.custom_comment = item.custom_comment
     session.add(record)
     session.commit()
+
+
+def run_fill_empty_table_embedding(session: SessionDep):
+    if not settings.EMBEDDING_ENABLED:
+        return
+
+    stmt = select(CoreTable.id).where(and_(CoreTable.embedding.is_(None)))
+    results = session.execute(stmt).scalars().all()
+
+    save_table_embedding(session, results)
+
+
+def save_table_embedding(session: SessionDep, ids: List[int]):
+    if not settings.EMBEDDING_ENABLED:
+        return
+
+    if not ids or len(ids) == 0:
+        return
+    try:
+
+        _list = session.query(CoreTable).filter(and_(CoreTable.id.in_(ids))).all()
+
+        table_schema = []
+        for item in _list:
+            fields = session.query(CoreField).filter(CoreField.table_id == item.id).all()
+
+            schema_table = ''
+            schema_table += f"# Table: {item.table_name}"
+            table_comment = ''
+            if item.custom_comment:
+                table_comment = item.custom_comment.strip()
+            if table_comment == '':
+                schema_table += '\n[\n'
+            else:
+                schema_table += f", {table_comment}\n[\n"
+
+            if fields:
+                field_list = []
+                for field in fields:
+                    field_comment = ''
+                    if field.custom_comment:
+                        field_comment = field.custom_comment.strip()
+                    if field_comment == '':
+                        field_list.append(f"({field.field_name}:{field.field_type})")
+                    else:
+                        field_list.append(f"({field.field_name}:{field.field_type}, {field_comment})")
+                schema_table += ",\n".join(field_list)
+            schema_table += '\n]\n'
+            table_schema.append(schema_table)
+
+        model = EmbeddingModelCache.get_model()
+
+        results = model.embed_documents(table_schema)
+
+        for index in range(len(results)):
+            item = results[index]
+            stmt = update(CoreTable).where(and_(CoreTable.id == _list[index].id)).values(embedding=item)
+            session.execute(stmt)
+            session.commit()
+
+    except Exception:
+        traceback.print_exc()

backend/apps/datasource/embedding/table_embedding.py

@@ -38,3 +38,36 @@ def get_table_embedding(tables: list[dict], question: str):
         except Exception:
             traceback.print_exc()
     return _list
+
+
+def calc_table_embedding(tables: list[dict], question: str):
+    _list = []
+    for table in tables:
+        _list.append(
+            {"id": table.get('id'), "schema_table": table.get('schema_table'), "embedding": table.get('embedding'),
+             "cosine_similarity": 0.0})
+
+    if _list:
+        try:
+            # text = [s.get('schema_table') for s in _list]
+            #
+            model = EmbeddingModelCache.get_model()
+            # start_time = time.time()
+            # results = model.embed_documents(text)
+            # end_time = time.time()
+            # SQLBotLogUtil.info(str(end_time - start_time))
+            results = [item.get('embedding') for item in _list]
+
+            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 _list
+        except Exception:
+            traceback.print_exc()
+    return _list

backend/apps/datasource/models/datasource.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import List, Optional
 
+from pgvector.sqlalchemy import VECTOR
 from pydantic import BaseModel
 from sqlalchemy import Column, Text, BigInteger, DateTime, Identity
 from sqlalchemy.dialects.postgresql import JSONB
@@ -31,6 +32,7 @@ class CoreTable(SQLModel, table=True):
     table_name: str = Field(sa_column=Column(Text))
     table_comment: str = Field(sa_column=Column(Text))
     custom_comment: str = Field(sa_column=Column(Text))
+    embedding: Optional[List[float]] = Field(sa_column=Column(VECTOR(), nullable=True))
 
 
 class CoreField(SQLModel, table=True):

backend/common/utils/embedding_threads.py

@@ -31,3 +31,13 @@ def run_save_data_training_embeddings(ids: List[int]):
 def fill_empty_data_training_embeddings():
     from apps.data_training.curd.data_training import run_fill_empty_embeddings
     executor.submit(run_fill_empty_embeddings, session)
+
+
+def run_save_table_embeddings(ids: List[int]):
+    from apps.datasource.crud.table import save_table_embedding
+    executor.submit(save_table_embedding, session, ids)
+
+
+def fill_empty_table_embeddings():
+    from apps.datasource.crud.table import run_fill_empty_table_embedding
+    executor.submit(run_fill_empty_table_embedding, session)

backend/main.py

@@ -18,7 +18,8 @@
 from common.core.config import settings
 from common.core.response_middleware import ResponseMiddleware, exception_handler
 from common.core.sqlbot_cache import init_sqlbot_cache
-from common.utils.embedding_threads import fill_empty_terminology_embeddings, fill_empty_data_training_embeddings
+from common.utils.embedding_threads import fill_empty_terminology_embeddings, fill_empty_data_training_embeddings, \
+    fill_empty_table_embeddings
 from common.utils.utils import SQLBotLogUtil
 
 
@@ -35,13 +36,18 @@ def init_data_training_embedding_data():
     fill_empty_data_training_embeddings()
 
 
+def init_table_embedding():
+    fill_empty_table_embeddings()
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     run_migrations()
     init_sqlbot_cache()
     init_dynamic_cors(app)
     init_terminology_embedding_data()
     init_data_training_embedding_data()
+    init_table_embedding()
     SQLBotLogUtil.info("✅ SQLBot 初始化完成")
     await sqlbot_xpack.core.clean_xpack_cache()
     await async_model_info()  # 异步加密已有模型的密钥和地址