diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index f7f812bb..0e6ff7ee 100644 --- a/backend/apps/chat/api/chat.py +++ b/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) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 66bb79e1..dfa2ea01 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/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): diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 6b6865a8..ba6540a0 100644 --- a/backend/apps/chat/task/llm.py +++ b/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, diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 5bf76a01..93ac5ad3 100644 --- a/backend/apps/mcp/mcp.py +++ b/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