mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: Vector retrieval matches datasource
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user