diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 306aa181..6cc30b83 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -40,6 +40,10 @@ class OperationEnum(Enum): CHOOSE_DATASOURCE = '6' GENERATE_DYNAMIC_SQL = '7' +class ChatFinishStep(Enum): + GENERATE_SQL = 1 + QUERY_DATA = 2 + GENERATE_CHART = 3 # TODO choose table / check connection / generate description diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index c5576510..9f55b468 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -28,7 +28,8 @@ from apps.chat.curd.chat import save_question, save_sql_answer, save_sql, \ get_old_questions, save_analysis_predict_record, rename_chat, get_chart_config, \ get_chat_chart_data, list_generate_sql_logs, list_generate_chart_logs, start_log, end_log, \ get_last_execute_sql_error -from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum +from apps.chat.models.chat_model import ChatQuestion, ChatRecord, Chat, RenameChat, ChatLog, OperationEnum, \ + ChatFinishStep from apps.data_training.curd.data_training import get_training_template from apps.datasource.crud.datasource import get_table_schema from apps.datasource.crud.permission import get_row_permission_filters, is_normal_user @@ -934,17 +935,20 @@ class LLMService: break yield chunk - def run_task_async(self, in_chat: bool = True, stream: bool = True): + def run_task_async(self, in_chat: bool = True, stream: bool = True, + finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART): if in_chat: stream = True - self.future = executor.submit(self.run_task_cache, in_chat, stream) + self.future = executor.submit(self.run_task_cache, in_chat, stream, finish_step) - def run_task_cache(self, in_chat: bool = True, stream: bool = True): - for chunk in self.run_task(in_chat, stream): + def run_task_cache(self, in_chat: bool = True, stream: bool = True, + finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART): + for chunk in self.run_task(in_chat, stream, finish_step): self.chunk_list.append(chunk) - def run_task(self, in_chat: bool = True, stream: bool = True): - json_result = {'success': True} + def run_task(self, in_chat: bool = True, stream: bool = True, + finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART): + json_result: Dict[str, Any] = {'success': True} try: if self.ds: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 @@ -1066,6 +1070,13 @@ class LLMService: subsql) real_execute_sql = assistant_dynamic_sql + if finish_step.value <= ChatFinishStep.GENERATE_SQL.value: + if in_chat: + yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n' + if not stream: + yield json_result + return + result = self.execute_sql(sql=real_execute_sql) self.save_sql_data(data_obj=result) if in_chat: @@ -1073,6 +1084,33 @@ class LLMService: if not stream: json_result['data'] = result.get('data') + if finish_step.value <= ChatFinishStep.QUERY_DATA.value: + if stream: + if in_chat: + yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n' + else: + data = [] + _fields_list = [] + _fields_skip = False + for _data in result.get('data'): + _row = [] + for field in result.get('fields'): + _row.append(_data.get(field)) + if not _fields_skip: + _fields_list.append(field) + data.append(_row) + _fields_skip = True + + if not data or not _fields_list: + yield 'The SQL execution result is empty.\n\n' + else: + df = pd.DataFrame(np.array(data), columns=_fields_list) + markdown_table = df.to_markdown(index=False) + yield markdown_table + '\n\n' + else: + yield json_result + return + # generate chart chart_res = self.generate_chart(chart_type) full_chart_text = '' diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index a1862975..433e4727 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -14,7 +14,8 @@ from sqlmodel import select from starlette.responses import JSONResponse from apps.chat.api.chat import create_chat -from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion, McpAssistant, ChatQuestion +from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion, McpAssistant, ChatQuestion, \ + ChatFinishStep from apps.chat.task.llm import LLMService from apps.system.crud.user import authenticate from apps.system.crud.user import get_db_user @@ -122,7 +123,10 @@ async def mcp_question(session: SessionDep, chat: McpQuestion): return StreamingResponse(_err(e), media_type="text/event-stream") else: - return {'message': str(e)} + return JSONResponse( + content={'message': str(e)}, + status_code=500, + ) if chat.stream: return StreamingResponse(llm_service.await_result(), media_type="text/event-stream") else: @@ -162,6 +166,36 @@ async def mcp_assistant(session: SessionDep, chat: McpAssistant): # assistant question mcp_chat = ChatQuestion(chat_id=c.id, question=chat.question) # ask - llm_service = await LLMService.create(session_user, mcp_chat, mcp_assistant_header) - llm_service.init_record() - return llm_service.run_task(False) + try: + llm_service = await LLMService.create(session_user, mcp_chat, mcp_assistant_header) + llm_service.init_record() + llm_service.run_task_async(False, chat.stream, ChatFinishStep.QUERY_DATA) + except Exception as e: + traceback.print_exc() + + if chat.stream: + def _err(_e: Exception): + yield str(_e) + '\n\n' + + return StreamingResponse(_err(e), media_type="text/event-stream") + else: + return JSONResponse( + content={'message': str(e)}, + status_code=500, + ) + if chat.stream: + return StreamingResponse(llm_service.await_result(), media_type="text/event-stream") + else: + res = llm_service.await_result() + raw_data = {} + for chunk in res: + if chunk: + raw_data = chunk + status_code = 200 + if not raw_data.get('success'): + status_code = 500 + + return JSONResponse( + content=raw_data, + status_code=status_code, + )