From eebb601a6dd854733823c17010eb6860465b4a46 Mon Sep 17 00:00:00 2001 From: junjun Date: Fri, 19 Sep 2025 15:04:37 +0800 Subject: [PATCH] feat: Vector retrieval matches datasource --- backend/apps/datasource/embedding/ds_embedding.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/backend/apps/datasource/embedding/ds_embedding.py b/backend/apps/datasource/embedding/ds_embedding.py index afa1ec69..2b637e18 100644 --- a/backend/apps/datasource/embedding/ds_embedding.py +++ b/backend/apps/datasource/embedding/ds_embedding.py @@ -8,6 +8,7 @@ from apps.ai_model.embedding import EmbeddingModelCache from apps.datasource.crud.datasource import get_table_schema from apps.datasource.models.datasource import CoreDatasource from common.core.deps import SessionDep, CurrentUser +from common.utils.utils import SQLBotLogUtil def cosine_similarity(vec_a, vec_b): @@ -50,7 +51,10 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, q _list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item) _list.sort(key=lambda x: x['cosine_similarity'], reverse=True) - print(len(_list)) + # print(len(_list)) + SQLBotLogUtil.info(json.dumps( + [{"id": ele.get("id"), "name": ele.get("name"), "cosine_similarity": ele.get("cosine_similarity")} + for ele in _list])) ds = _list[0].get('ds') return {"id": ds.id, "name": ds.name, "description": ds.description} except Exception: