From eb82d4fb56b464b12a7394e25bda621eece986bf Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 23 Sep 2025 15:11:08 +0800 Subject: [PATCH] feat: Vector retrieval matches tables --- backend/apps/datasource/crud/datasource.py | 33 +---------------- .../apps/datasource/embedding/ds_embedding.py | 19 +--------- .../datasource/embedding/table_embedding.py | 37 +++++++++++++++++++ backend/apps/datasource/embedding/utils.py | 18 +++++++++ 4 files changed, 58 insertions(+), 49 deletions(-) create mode 100644 backend/apps/datasource/embedding/table_embedding.py create mode 100644 backend/apps/datasource/embedding/utils.py diff --git a/backend/apps/datasource/crud/datasource.py b/backend/apps/datasource/crud/datasource.py index de527712..b43e3277 100644 --- a/backend/apps/datasource/crud/datasource.py +++ b/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 diff --git a/backend/apps/datasource/embedding/ds_embedding.py b/backend/apps/datasource/embedding/ds_embedding.py index b024fd6e..7aad7d29 100644 --- a/backend/apps/datasource/embedding/ds_embedding.py +++ b/backend/apps/datasource/embedding/ds_embedding.py @@ -1,12 +1,12 @@ # 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 @@ -14,21 +14,6 @@ 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}) diff --git a/backend/apps/datasource/embedding/table_embedding.py b/backend/apps/datasource/embedding/table_embedding.py new file mode 100644 index 00000000..e1827174 --- /dev/null +++ b/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 diff --git a/backend/apps/datasource/embedding/utils.py b/backend/apps/datasource/embedding/utils.py new file mode 100644 index 00000000..3f6ddced --- /dev/null +++ b/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)