feat: MCP mcp/mcp_assistant support param to disable streaming

Author: ulleo

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
 

backend/apps/chat/task/llm.py

@@ -28,7 +28,8 @@
     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 @@ def await_result(self):
                 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,13 +1070,47 @@ def run_task(self, in_chat: bool = True, stream: bool = True):
                                                                           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:
                 yield 'data:' + orjson.dumps({'content': 'execute-success', 'type': 'sql-data'}).decode() + '\n\n'
             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 = ''

backend/apps/mcp/mcp.py

@@ -14,7 +14,8 @@
 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 @@ def _err(_e: Exception):
 
             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,
+        )