diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index fdff8872..14cfdbf4 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -7,7 +7,7 @@ from apps.chat.curd.chat import list_chats, get_chat_with_records, create_chat, delete_chat, get_chat_chart_data, get_chat_predict_data from apps.chat.models.chat_model import CreateChat, ChatRecord, RenameChat, ChatQuestion from apps.chat.task.llm import LLMService, run_task, run_analysis_or_predict_task, run_recommend_questions_task -from common.core.deps import SessionDep, CurrentUser +from common.core.deps import CurrentAssistant, SessionDep, CurrentUser router = APIRouter(tags=["Data Q&A"], prefix="/chat") @@ -94,7 +94,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser): @router.post("/recommend_questions/{chat_record_id}") -async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int): +async def recommend_questions(session: SessionDep, current_user: CurrentUser, chat_record_id: int, current_assistant: CurrentAssistant): try: record = session.query(ChatRecord).get(chat_record_id) if not record: @@ -104,7 +104,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch ) request_question = ChatQuestion(chat_id=record.chat_id, question=record.question if record.question else '') - llm_service = LLMService(session, current_user, request_question) + llm_service = LLMService(session, current_user, request_question, current_assistant) llm_service.set_record(record) except Exception as e: traceback.print_exc() @@ -117,7 +117,7 @@ async def recommend_questions(session: SessionDep, current_user: CurrentUser, ch @router.post("/question") -async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion): +async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion, current_assistant: CurrentAssistant): """Stream SQL analysis results Args: @@ -130,7 +130,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que """ try: - llm_service = LLMService(session, current_user, request_question) + llm_service = LLMService(session, current_user, request_question, current_assistant) llm_service.init_record() except Exception as e: traceback.print_exc() @@ -143,7 +143,7 @@ async def stream_sql(session: SessionDep, current_user: CurrentUser, request_que @router.post("/record/{chat_record_id}/{action_type}") -async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str): +async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, chat_record_id: int, action_type: str, current_assistant: CurrentAssistant): if action_type != 'analysis' and action_type != 'predict': raise HTTPException( status_code=404, @@ -166,7 +166,7 @@ async def analysis_or_predict(session: SessionDep, current_user: CurrentUser, ch request_question = ChatQuestion(chat_id=record.chat_id, question='') try: - llm_service = LLMService(session, current_user, request_question) + llm_service = LLMService(session, current_user, request_question, current_assistant) except Exception as e: traceback.print_exc() raise HTTPException( diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index d8e23569..78d5d222 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -353,11 +353,21 @@ class LLMService: if self.current_assistant: _ds_list = get_assistant_ds(llm_service=self) else: - _ds_list = self.session.exec(select(CoreDatasource).options( - load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all() + oid: str = self.current_user.oid + stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid) + _ds_list = [ + { + "id": ds.id, + "name": ds.name, + "description": ds.description + } + for ds in self.session.exec(stmt) + ] + """ _ds_list = self.session.exec(select(CoreDatasource).options( + load_only(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description))).all() """ _ds_list_dict = [] for _ds in _ds_list: - _ds_list_dict.append({'id': _ds[0].id, 'name': _ds[0].name, 'description': _ds[0].description}) + _ds_list_dict.append(_ds) datasource_msg.append( HumanMessage(self.chat_question.datasource_user_question(orjson.dumps(_ds_list_dict).decode()))) diff --git a/backend/apps/system/crud/assistant.py b/backend/apps/system/crud/assistant.py index 5aadc8f0..66b674a9 100644 --- a/backend/apps/system/crud/assistant.py +++ b/backend/apps/system/crud/assistant.py @@ -33,12 +33,22 @@ def get_assistant_ds(llm_service) -> list[dict]: config: dict[any] = json.loads(configuration) oid: str = config['oid'] stmt = select(CoreDatasource.id, CoreDatasource.name, CoreDatasource.description).where(CoreDatasource.oid == oid) - private_list:list[int] = config['private_list'] + private_list:list[int] = config.get('private_list') or None if private_list: stmt.where(~CoreDatasource.id.in_(private_list)) - db_ds_list = session.exec(stmt).all() + db_ds_list = session.exec(stmt) + + result_list = [ + { + "id": ds.id, + "name": ds.name, + "description": ds.description + } + for ds in db_ds_list + ] + # filter private ds if offline - return db_ds_list + return result_list out_ds_instance: AssistantOutDs = AssistantOutDsFactory.get_instance(assistant, llm_service.assistant_certificate) llm_service.out_ds_instance = out_ds_instance dslist = out_ds_instance.get_simple_ds_list()