feat: MCP supports configuration to disable rendering charts

Author: XiaJunjie2020

backend/apps/chat/api/chat.py

@@ -80,6 +80,7 @@ def inner():
 
     return await asyncio.to_thread(inner)
 
+
 @router.get("/record/{chat_record_id}/data_live", summary=f"{PLACEHOLDER_PREFIX}get_chart_data_live")
 async def chat_record_data_live(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
     def inner():
@@ -88,6 +89,7 @@ def inner():
 
     return await asyncio.to_thread(inner)
 
+
 @router.get("/record/{chat_record_id}/predict_data", summary=f"{PLACEHOLDER_PREFIX}get_chart_predict_data")
 async def chat_predict_data(session: SessionDep, current_user: CurrentUser, chat_record_id: int):
     def inner():
@@ -275,7 +277,8 @@ async def question_answer(session: SessionDep, current_user: CurrentUser, reques
 async def question_answer_inner(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
                                 current_assistant: Optional[CurrentAssistant] = None, in_chat: bool = True,
                                 stream: bool = True,
-                                finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False):
+                                finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False,
+                                return_img: bool = True):
     try:
         command, text_before_command, record_id, warning_info = parse_quick_command(request_question.question)
         if command:
@@ -331,7 +334,7 @@ async def question_answer_inner(session: SessionDep, current_user: CurrentUser,
                 request_question.question = text_before_command
                 request_question.regenerate_record_id = rec_id
                 return await stream_sql(session, current_user, request_question, current_assistant, in_chat, stream,
-                                        finish_step, embedding)
+                                        finish_step, embedding, return_img)
 
             elif command == QuickCommand.ANALYSIS:
                 return await analysis_or_predict(session, current_user, rec_id, 'analysis', current_assistant, in_chat,
@@ -344,7 +347,7 @@ async def question_answer_inner(session: SessionDep, current_user: CurrentUser,
                 raise Exception(f'Unknown command: {command.value}')
         else:
             return await stream_sql(session, current_user, request_question, current_assistant, in_chat, stream,
-                                    finish_step, embedding)
+                                    finish_step, embedding, return_img)
     except Exception as e:
         traceback.print_exc()
 
@@ -366,12 +369,13 @@ def _err(_e: Exception):
 
 async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion,
                      current_assistant: Optional[CurrentAssistant] = None, in_chat: bool = True, stream: bool = True,
-                     finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False):
+                     finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, embedding: bool = False,
+                     return_img: bool = True):
     try:
         llm_service = await LLMService.create(session, current_user, request_question, current_assistant,
                                               embedding=embedding)
         llm_service.init_record(session=session)
-        llm_service.run_task_async(in_chat=in_chat, stream=stream, finish_step=finish_step)
+        llm_service.run_task_async(in_chat=in_chat, stream=stream, finish_step=finish_step, return_img=return_img)
     except Exception as e:
         traceback.print_exc()
 
@@ -536,7 +540,8 @@ async def export_excel(session: SessionDep, current_user: CurrentUser, chat_reco
 
     def inner():
 
-        data_list = DataFormat.convert_large_numbers_in_object_array(obj_array=_data + _predict_data, int_threshold=1e11)
+        data_list = DataFormat.convert_large_numbers_in_object_array(obj_array=_data + _predict_data,
+                                                                     int_threshold=1e11)
 
         md_data, _fields_list = DataFormat.convert_object_array_for_pandas(fields, data_list)
 

backend/apps/chat/models/chat_model.py

@@ -338,6 +338,7 @@ class McpQuestion(BaseModel):
     datasource_id: Optional[int | str] = Body(description='数据源ID，仅当当前对话没有确定数据源时有效', default=None)
     oid: Optional[str] = Body(
         description='组织ID，仅当数据源ID为空时有效，如果不传则为最后一次登录SQLBot时所使用的组织ID', default=None)
+    return_img: Optional[bool] = Body(description='是否返回图表，默认为true开启, 关闭false则仅返回数据', default=True)
 
 
 class AxisObj(BaseModel):

backend/apps/chat/task/llm.py

@@ -1059,18 +1059,18 @@ def await_result(self):
             yield chunk
 
     def run_task_async(self, in_chat: bool = True, stream: bool = True,
-                       finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART):
+                       finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
         if in_chat:
             stream = True
-        self.future = executor.submit(self.run_task_cache, in_chat, stream, finish_step)
+        self.future = executor.submit(self.run_task_cache, in_chat, stream, finish_step, return_img)
 
     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):
+                       finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
+        for chunk in self.run_task(in_chat, stream, finish_step, return_img):
             self.chunk_list.append(chunk)
 
     def run_task(self, in_chat: bool = True, stream: bool = True,
-                 finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART):
+                 finish_step: ChatFinishStep = ChatFinishStep.GENERATE_CHART, return_img: bool = True):
         json_result: Dict[str, Any] = {'success': True}
         _session = None
         try:
@@ -1311,7 +1311,7 @@ def run_task(self, in_chat: bool = True, stream: bool = True,
             else:
                 # generate picture
                 try:
-                    if chart.get('type') != 'table':
+                    if chart.get('type') != 'table' and return_img:
                         # yield '### generated chart picture\n\n'
                         self.current_logs[OperationEnum.GENERATE_PICTURE] = start_log(session=_session,
                                                                                       operate=OperationEnum.GENERATE_PICTURE,

backend/apps/mcp/mcp.py

@@ -151,7 +151,7 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
     mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question, datasource_id=ds_id)
 
     return await question_answer_inner(session=session, current_user=session_user, request_question=mcp_chat,
-                                       in_chat=False, stream=chat.stream)
+                                       in_chat=False, stream=chat.stream, return_img=chat.return_img)
 
 
 # Cordys crm