feat: Vector retrieval matches datasource

This commit is contained in:
junjun
2025-09-19 15:07:21 +08:00
committed by Junjun
parent 336ffcf2a1
commit eebb601a6d
@@ -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: