diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index 92efe3c6..f9339bdc 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/backend/apps/datasource/crud/datasource.py @@ -8,13 +8,14 @@ from sqlbot_xpack.permissions.models.ds_rules import DsRules 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: diff --git a/backend/apps/datasource/crud/table.py b/backend/apps/datasource/crud/table.py index f76d5fa2..55153a32 100644 --- a/backend/apps/datasource/crud/table.py +++ b/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() diff --git a/backend/apps/datasource/embedding/table_embedding.py b/backend/apps/datasource/embedding/table_embedding.py index e8efc405..7ba96404 100644 --- a/backend/apps/datasource/embedding/table_embedding.py +++ b/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 diff --git a/backend/apps/datasource/models/datasource.py b/backend/apps/datasource/models/datasource.py index 78f23916..fd8f3d5b 100644 --- a/backend/apps/datasource/models/datasource.py +++ b/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): diff --git a/backend/common/utils/embedding_threads.py b/backend/common/utils/embedding_threads.py index a38b66f0..182c525e 100644 --- a/backend/common/utils/embedding_threads.py +++ b/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) diff --git a/backend/main.py b/backend/main.py index 34e1ad67..3486df56 100644 --- a/backend/main.py +++ b/backend/main.py @@ -18,7 +18,8 @@ from apps.system.middleware.auth import TokenMiddleware 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,6 +36,10 @@ 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() @@ -42,6 +47,7 @@ async def lifespan(app: FastAPI): 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() # 异步加密已有模型的密钥和地址