diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index cb0cb268..c5576510 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -426,7 +426,8 @@ class LLMService: full_text = '' if not ignore_auto_select: if settings.EMBEDDING_ENABLED: - ds = get_ds_embedding(self.session, self.current_user, _ds_list, self.chat_question.question) + ds = get_ds_embedding(self.session, self.current_user, _ds_list, self.out_ds_instance, + self.chat_question.question, self.current_assistant) yield {'content': '{"id":' + str(ds.get('id')) + '}'} else: _ds_list_dict = [] diff --git a/backend/apps/datasource/embedding/ds_embedding.py b/backend/apps/datasource/embedding/ds_embedding.py index a8875744..d37001ae 100644 --- a/backend/apps/datasource/embedding/ds_embedding.py +++ b/backend/apps/datasource/embedding/ds_embedding.py @@ -3,10 +3,13 @@ import json import math import traceback +from typing import Optional from apps.ai_model.embedding import EmbeddingModelCache from apps.datasource.crud.datasource import get_table_schema from apps.datasource.models.datasource import CoreDatasource +from apps.system.crud.assistant import AssistantOutDs +from common.core.deps import CurrentAssistant from common.core.deps import SessionDep, CurrentUser from common.utils.utils import SQLBotLogUtil @@ -26,17 +29,26 @@ def cosine_similarity(vec_a, vec_b): return dot_product / (norm_a * norm_b) -def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, question: str): +def get_ds_embedding(session: SessionDep, current_user: CurrentUser, _ds_list, out_ds: AssistantOutDs, + question: str, + current_assistant: Optional[CurrentAssistant] = None): _list = [] - for _ds in _ds_list: - if _ds.get('id'): - ds = session.get(CoreDatasource, _ds.get('id')) - - table_schema = get_table_schema(session, current_user, ds) - ds_info = f"{ds.name}, {ds.description}\n" - ds_schema = ds_info + table_schema - - _list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds}) + if current_assistant and current_assistant.type != 4: + if out_ds.ds_list: + for _ds in out_ds.ds_list: + ds = out_ds.get_ds(_ds.id) + table_schema = out_ds.get_db_schema(_ds.id) + ds_info = f"{ds.name}, {ds.description}\n" + ds_schema = ds_info + table_schema + _list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds}) + else: + for _ds in _ds_list: + if _ds.get('id'): + ds = session.get(CoreDatasource, _ds.get('id')) + table_schema = get_table_schema(session, current_user, ds) + ds_info = f"{ds.name}, {ds.description}\n" + ds_schema = ds_info + table_schema + _list.append({"id": ds.id, "ds_schema": ds_schema, "cosine_similarity": 0.0, "ds": ds}) if _list: try: