From 568afa73996b0c937ff64f838d3f8a0878910c89 Mon Sep 17 00:00:00 2001 From: junjun Date: Thu, 23 Oct 2025 09:54:16 +0800 Subject: [PATCH] feat: Vector retrieval matches datasource --- backend/apps/datasource/embedding/ds_embedding.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/backend/apps/datasource/embedding/ds_embedding.py b/backend/apps/datasource/embedding/ds_embedding.py index 657a3bf0..19feada2 100644 --- a/backend/apps/datasource/embedding/ds_embedding.py +++ b/backend/apps/datasource/embedding/ds_embedding.py @@ -42,13 +42,13 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o _list.sort(key=lambda x: x['cosine_similarity'], reverse=True) # print(len(_list)) + _list = _list[:settings.DS_EMBEDDING_COUNT] SQLBotLogUtil.info(json.dumps( [{"id": ele.get("id"), "name": ele.get("ds").name, "cosine_similarity": ele.get("cosine_similarity")} for ele in _list])) - ds_l = _list[:settings.DS_EMBEDDING_COUNT] return [{"id": obj.get('ds').id, "name": obj.get('ds').name, "description": obj.get('ds').description} - for obj in ds_l] + for obj in _list] except Exception: traceback.print_exc() else: @@ -79,13 +79,13 @@ def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, o # print(len(_list)) end_time = time.time() SQLBotLogUtil.info(str(end_time - start_time)) + _list = _list[:settings.DS_EMBEDDING_COUNT] SQLBotLogUtil.info(json.dumps( [{"id": ele.get("id"), "name": ele.get("ds").name, "cosine_similarity": ele.get("cosine_similarity")} for ele in _list])) - ds_l = _list[:settings.DS_EMBEDDING_COUNT] return [{"id": obj.get('ds').id, "name": obj.get('ds').name, "description": obj.get('ds').description} - for obj in ds_l] + for obj in _list] except Exception: traceback.print_exc() return _list