(session: SessionDep, current_user: CurrentUser, _ds_list, out_ds: AssistantOutDs,
question: str,
current_assistant: Optional[CurrentAssistant] = None)
| 16 | |
| 17 | |
| 18 | def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, out_ds: AssistantOutDs, |
| 19 | question: str, |
| 20 | current_assistant: Optional[CurrentAssistant] = None): |
| 21 | _list = [] |
| 22 | if current_assistant and current_assistant.type == 1: |
| 23 | if out_ds.ds_list: |
| 24 | for _ds in out_ds.ds_list: |
| 25 | ds = out_ds.get_ds(_ds.id) |
| 26 | table_schema, tables = out_ds.get_db_schema(_ds.id, question, embedding=False) |
| 27 | ds_info = f"{ds.name}, {ds.description}\n" |
| 28 | ds_schema = ds_info + table_schema |
| 29 | _list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds}) |
| 30 | |
| 31 | if _list: |
| 32 | try: |
| 33 | text = [s.get('ds_schema') for s in _list] |
| 34 | |
| 35 | model = EmbeddingModelCache.get_model() |
| 36 | results = model.embed_documents(text) |
| 37 | |
| 38 | q_embedding = model.embed_query(question) |
| 39 | for index in range(len(results)): |
| 40 | item = results[index] |
| 41 | _list[index]['cosine_similarity'] = cosine_similarity(q_embedding, item) |
| 42 | |
| 43 | _list.sort(key=lambda x: x['cosine_similarity'], reverse=True) |
| 44 | # print(len(_list)) |
| 45 | _list = _list[:settings.DS_EMBEDDING_COUNT] |
| 46 | SQLBotLogUtil.info(json.dumps( |
| 47 | [{"id": ele.get("id"), "name": ele.get("ds").name, |
| 48 | "cosine_similarity": ele.get("cosine_similarity")} |
| 49 | for ele in _list])) |
| 50 | return [{"id": obj.get('ds').id, "name": obj.get('ds').name, "description": obj.get('ds').description} |
| 51 | for obj in _list] |
| 52 | except Exception: |
| 53 | traceback.print_exc() |
| 54 | else: |
| 55 | for _ds in _ds_list: |
| 56 | if _ds.get('id'): |
| 57 | ds = session.get(CoreDatasource, _ds.get('id')) |
| 58 | # table_schema = get_table_schema(session, current_user, ds, question, embedding=False) |
| 59 | # ds_info = f"{ds.name}, {ds.description}\n" |
| 60 | # ds_schema = ds_info + table_schema |
| 61 | _list.append({"id": ds.id, "cosine_similarity": 0.0, "ds": ds, "embedding": ds.embedding}) |
| 62 | |
| 63 | if _list: |
| 64 | try: |
| 65 | # text = [s.get('ds_schema') for s in _list] |
| 66 | |
| 67 | model = EmbeddingModelCache.get_model() |
| 68 | start_time = time.time() |
| 69 | # results = model.embed_documents(text) |
| 70 | results = [item.get('embedding') for item in _list] |
| 71 | |
| 72 | q_embedding = model.embed_query(question) |
| 73 | for index in range(len(results)): |
| 74 | item = results[index] |
| 75 | if item: |
no test coverage detected