From dca52ce1b2845ab11c41ff5ac4b93b8f47915dd1 Mon Sep 17 00:00:00 2001 From: ulleo Date: Mon, 22 Sep 2025 18:19:25 +0800 Subject: [PATCH] feat: MCP mcp/mcp_question support param to disable streaming --- backend/apps/chat/models/chat_model.py | 4 +- backend/apps/chat/task/llm.py | 110 ++++++++++++------- backend/apps/mcp/mcp.py | 36 +++++- backend/apps/system/schemas/system_schema.py | 89 +++++++++------ backend/common/core/response_middleware.py | 34 ++++-- 5 files changed, 182 insertions(+), 91 deletions(-) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 2173b539..306aa181 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -136,7 +136,7 @@ class CreateChat(BaseModel): id: int = None question: str = None datasource: int = None - origin: Optional[int] = 0 + origin: Optional[int] = 0 # 0是页面上，mcp是1，小助手是2 class RenameChat(BaseModel): @@ -246,6 +246,7 @@ class McpQuestion(BaseModel): question: str = Body(description='用户提问') chat_id: int = Body(description='会话ID') token: str = Body(description='token') + stream: Optional[bool] = Body(description='是否流式输出，默认为true开启, 关闭false则返回JSON对象', default=True) class AxisObj(BaseModel): @@ -264,3 +265,4 @@ class McpAssistant(BaseModel): question: str = Body(description='用户提问') url: str = Body(description='第三方数据接口') authorization: str = Body(description='第三方接口凭证') + stream: Optional[bool] = Body(description='是否流式输出，默认为true开启, 关闭false则返回JSON对象', default=True) diff --git a/backend/apps/chat/task/llm.py b/backend/apps/chat/task/llm.py index 5f988199..cb0cb268 100644 --- a/backend/apps/chat/task/llm.py +++ b/backend/apps/chat/task/llm.py @@ -933,14 +933,17 @@ class LLMService: break yield chunk - def run_task_async(self, in_chat: bool = True): - self.future = executor.submit(self.run_task_cache, in_chat) + def run_task_async(self, in_chat: bool = True, stream: bool = True): + if in_chat: + stream = True + self.future = executor.submit(self.run_task_cache, in_chat, stream) - def run_task_cache(self, in_chat: bool = True): - for chunk in self.run_task(in_chat): + def run_task_cache(self, in_chat: bool = True, stream: bool = True): + for chunk in self.run_task(in_chat, stream): self.chunk_list.append(chunk) - def run_task(self, in_chat: bool = True): + def run_task(self, in_chat: bool = True, stream: bool = True): + json_result = {'success': True} try: if self.ds: oid = self.ds.oid if isinstance(self.ds, CoreDatasource) else 1 @@ -955,6 +958,8 @@ class LLMService: # return id if in_chat: yield 'data:' + orjson.dumps({'type': 'id', 'id': self.get_record().id}).decode() + '\n\n' + if not stream: + json_result['record_id'] = self.get_record().id # return title if self.change_title: @@ -964,8 +969,10 @@ class LLMService: brief=self.chat_question.question.strip()[:20])) if in_chat: yield 'data:' + orjson.dumps({'type': 'brief', 'brief': brief}).decode() + '\n\n' + if not stream: + json_result['title'] = brief - # select datasource if datasource is none + # select datasource if datasource is none if not self.ds: ds_res = self.select_datasource() @@ -1003,7 +1010,6 @@ class LLMService: 'type': 'sql-result'}).decode() + '\n\n' if in_chat: yield 'data:' + orjson.dumps({'type': 'info', 'msg': 'sql generated'}).decode() + '\n\n' - # filter sql SQLBotLogUtil.info(full_sql_text) @@ -1039,11 +1045,16 @@ class LLMService: sql = self.check_save_sql(res=full_sql_text) SQLBotLogUtil.info(sql) + + if not stream: + json_result['sql'] = sql + format_sql = sqlparse.format(sql, reindent=True) if in_chat: yield 'data:' + orjson.dumps({'content': format_sql, 'type': 'sql'}).decode() + '\n\n' else: - yield f'```sql\n{format_sql}\n```\n\n' + if stream: + yield f'```sql\n{format_sql}\n```\n\n' # execute sql real_execute_sql = sql @@ -1058,6 +1069,8 @@ class LLMService: 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') # generate chart chart_res = self.generate_chart(chart_type) @@ -1075,41 +1088,46 @@ class LLMService: SQLBotLogUtil.info(full_chart_text) chart = self.check_save_chart(res=full_chart_text) SQLBotLogUtil.info(chart) + + if not stream: + json_result['chart'] = chart + if in_chat: yield 'data:' + orjson.dumps( {'content': orjson.dumps(chart).decode(), 'type': 'chart'}).decode() + '\n\n' else: - data = [] - _fields = {} - if chart.get('columns'): - for _column in chart.get('columns'): - if _column: - _fields[_column.get('value')] = _column.get('name') - if chart.get('axis'): - if chart.get('axis').get('x'): - _fields[chart.get('axis').get('x').get('value')] = chart.get('axis').get('x').get('name') - if chart.get('axis').get('y'): - _fields[chart.get('axis').get('y').get('value')] = chart.get('axis').get('y').get('name') - if chart.get('axis').get('series'): - _fields[chart.get('axis').get('series').get('value')] = chart.get('axis').get('series').get( - 'name') - _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 if not _fields.get(field) else _fields.get(field)) - data.append(_row) - _fields_skip = True + if stream: + data = [] + _fields = {} + if chart.get('columns'): + for _column in chart.get('columns'): + if _column: + _fields[_column.get('value')] = _column.get('name') + if chart.get('axis'): + if chart.get('axis').get('x'): + _fields[chart.get('axis').get('x').get('value')] = chart.get('axis').get('x').get('name') + if chart.get('axis').get('y'): + _fields[chart.get('axis').get('y').get('value')] = chart.get('axis').get('y').get('name') + if chart.get('axis').get('series'): + _fields[chart.get('axis').get('series').get('value')] = chart.get('axis').get('series').get( + 'name') + _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 if not _fields.get(field) else _fields.get(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' + 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' if in_chat: yield 'data:' + orjson.dumps({'type': 'finish'}).decode() + '\n\n' @@ -1119,7 +1137,14 @@ class LLMService: yield '### generated chart picture\n\n' image_url = request_picture(self.record.chat_id, self.record.id, chart, result) SQLBotLogUtil.info(image_url) - yield f'![{chart["type"]}]({image_url})' + if stream: + yield f'![{chart["type"]}]({image_url})' + else: + json_result['image_url'] = image_url + + if not stream: + yield json_result + except Exception as e: traceback.print_exc() error_msg: str @@ -1137,7 +1162,12 @@ class LLMService: if in_chat: yield 'data:' + orjson.dumps({'content': error_msg, 'type': 'error'}).decode() + '\n\n' else: - yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。' + if stream: + yield f'> ❌ **ERROR**\n\n> \n\n> {error_msg}。' + else: + json_result['success'] = False + json_result['message'] = error_msg + yield json_result finally: self.finish() diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index b1126962..a1862975 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -1,6 +1,7 @@ # Author: Junjun # Date: 2025/7/1 import json +import traceback from datetime import timedelta import jwt @@ -10,6 +11,7 @@ from fastapi.responses import StreamingResponse from jwt.exceptions import InvalidTokenError from pydantic import ValidationError 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 @@ -106,11 +108,37 @@ async def mcp_question(session: SessionDep, chat: McpQuestion): raise HTTPException(status_code=400, detail="Inactive user") mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question) - # ask - llm_service = await LLMService.create(session_user, mcp_chat) - llm_service.init_record() - return StreamingResponse(llm_service.run_task(False), media_type="text/event-stream") + try: + llm_service = await LLMService.create(session_user, mcp_chat) + llm_service.init_record() + llm_service.run_task_async(False, chat.stream) + 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 {'message': str(e)} + 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, + ) @router.post("/mcp_assistant", operation_id="mcp_assistant") diff --git a/backend/apps/system/schemas/system_schema.py b/backend/apps/system/schemas/system_schema.py index 79eca6f4..728ca480 100644 --- a/backend/apps/system/schemas/system_schema.py +++ b/backend/apps/system/schemas/system_schema.py @@ -1,8 +1,9 @@ +import re from typing import Optional + from pydantic import BaseModel, Field, field_validator from common.core.schemas import BaseCreatorDTO -import re EMAIL_REGEX = re.compile( r"^[a-zA-Z0-9]+([._-][a-zA-Z0-9]+)*@" @@ -14,106 +15,121 @@ PWD_REGEX = re.compile( r"(?=.*[~!@#$%^&*()_+\-={}|:\"<>?`\[\];',./])" r"[A-Za-z\d~!@#$%^&*()_+\-={}|:\"<>?`\[\];',./]{8,20}$" ) + + class UserStatus(BaseCreatorDTO): status: int = 1 - - + class UserLanguage(BaseModel): language: str - + class BaseUser(BaseModel): account: str = Field(min_length=1, max_length=100, description="用户账号") oid: int - + + class BaseUserDTO(BaseUser, BaseCreatorDTO): language: str = Field(pattern=r"^(zh-CN|en)$", default="zh-CN", description="用户语言") password: str status: int = 1 + def to_dict(self): return { "id": self.id, "account": self.account, "oid": self.oid } + @field_validator("language") def validate_language(cls, lang: str) -> str: if not re.fullmatch(r"^(zh-CN|en)$", lang): raise ValueError("Language must be 'zh-CN' or 'en'") return lang + class UserCreator(BaseUser): name: str = Field(min_length=1, max_length=100, description="用户名") email: str = Field(min_length=1, max_length=100, description="用户邮箱") status: int = 1 - oid_list: Optional[list[int]] = None - + oid_list: Optional[list[int]] = None + """ @field_validator("email") def validate_email(cls, lang: str) -> str: if not re.fullmatch(EMAIL_REGEX, lang): raise ValueError("Email format is invalid!") return lang """ + class UserEditor(UserCreator, BaseCreatorDTO): - pass + pass + class UserGrid(UserEditor): create_time: int language: str = "zh-CN" - #space_name: Optional[str] = None + # space_name: Optional[str] = None origin: str = '' - - + + class PwdEditor(BaseModel): pwd: str new_pwd: str - + + class UserWsBase(BaseModel): uid_list: list[int] oid: Optional[int] = None + + class UserWsDTO(UserWsBase): weight: Optional[int] = 0 + class UserWsEditor(BaseModel): uid: int - oid: int + oid: int weight: int = 0 + class UserInfoDTO(UserEditor): language: str = "zh-CN" weight: int = 0 isAdmin: bool = False - + class AssistantBase(BaseModel): name: str domain: str - type: int = 0 + type: int = 0 # 0普通小助手 1高级 4页面嵌入 configuration: Optional[str] = None description: Optional[str] = None + + class AssistantDTO(AssistantBase, BaseCreatorDTO): pass + class AssistantHeader(AssistantDTO): unique: Optional[str] = None certificate: Optional[str] = None online: bool = False - + class AssistantValidator(BaseModel): valid: bool = False id_match: bool = False domain_match: bool = False token: Optional[str] = None - + def __init__( - self, - valid: bool = False, - id_match: bool = False, - domain_match: bool = False, - token: Optional[str] = None, - **kwargs + self, + valid: bool = False, + id_match: bool = False, + domain_match: bool = False, + token: Optional[str] = None, + **kwargs ): super().__init__( valid=valid, @@ -122,23 +138,28 @@ class AssistantValidator(BaseModel): token=token, **kwargs ) - + + class WorkspaceUser(UserEditor): weight: int create_time: int - + + class UserWs(BaseCreatorDTO): name: str - + + class UserWsOption(UserWs): account: str - - + + class AssistantFieldSchema(BaseModel): id: Optional[int] = None name: Optional[str] = None type: Optional[str] = None comment: Optional[str] = None + + class AssistantTableSchema(BaseModel): id: Optional[int] = None name: Optional[str] = None @@ -147,6 +168,7 @@ class AssistantTableSchema(BaseModel): sql: Optional[str] = None fields: Optional[list[AssistantFieldSchema]] = None + class AssistantOutDsBase(BaseModel): id: Optional[int] = None name: str @@ -154,8 +176,8 @@ class AssistantOutDsBase(BaseModel): type_name: Optional[str] = None comment: Optional[str] = None description: Optional[str] = None - - + + class AssistantOutDsSchema(AssistantOutDsBase): host: Optional[str] = None port: Optional[int] = None @@ -165,7 +187,8 @@ class AssistantOutDsSchema(AssistantOutDsBase): db_schema: Optional[str] = None extraParams: Optional[str] = None tables: Optional[list[AssistantTableSchema]] = None - + + class AssistantUiSchema(BaseCreatorDTO): theme: Optional[str] = None header_font_color: Optional[str] = None @@ -179,7 +202,3 @@ class AssistantUiSchema(BaseCreatorDTO): name: Optional[str] = None welcome: Optional[str] = None welcome_desc: Optional[str] = None - - - - \ No newline at end of file diff --git a/backend/common/core/response_middleware.py b/backend/common/core/response_middleware.py index 6842893e..69c49e54 100644 --- a/backend/common/core/response_middleware.py +++ b/backend/common/core/response_middleware.py @@ -1,18 +1,32 @@ import json -from starlette.middleware.base import BaseHTTPMiddleware -from starlette.responses import JSONResponse + from starlette.exceptions import HTTPException +from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request +from starlette.responses import JSONResponse + from common.core.config import settings from common.utils.utils import SQLBotLogUtil + + class ResponseMiddleware(BaseHTTPMiddleware): def __init__(self, app): super().__init__(app) async def dispatch(self, request, call_next): response = await call_next(request) - - if isinstance(response, JSONResponse) or request.url.path == f"{settings.API_V1_STR}/openapi.json": + + direct_paths = [ + f"{settings.API_V1_STR}/openapi.json", + f"{settings.API_V1_STR}/mcp/mcp_question", + f"{settings.API_V1_STR}/mcp/mcp_assistant" + ] + + route = request.scope.get("route") + # 获取定义的路径模式，例如 '/items/{item_id}' + path_pattern = route.path_format + + if isinstance(response, JSONResponse) or path_pattern in direct_paths: return response if response.status_code != 200: return response @@ -21,9 +35,9 @@ class ResponseMiddleware(BaseHTTPMiddleware): body = b"" async for chunk in response.body_iterator: body += chunk - + raw_data = json.loads(body.decode()) - + if isinstance(raw_data, dict) and all(k in raw_data for k in ["code", "data", "msg"]): return JSONResponse( content=raw_data, @@ -33,13 +47,13 @@ class ResponseMiddleware(BaseHTTPMiddleware): if k.lower() not in ("content-length", "content-type") } ) - + wrapped_data = { "code": 0, "data": raw_data, "msg": None } - + return JSONResponse( content=wrapped_data, status_code=response.status_code, @@ -58,7 +72,7 @@ class ResponseMiddleware(BaseHTTPMiddleware): if k.lower() not in ("content-length", "content-type") } ) - + return response @@ -72,7 +86,6 @@ class exception_handler(): headers={"Access-Control-Allow-Origin": "*"} ) - @staticmethod async def global_exception_handler(request: Request, exc: Exception): SQLBotLogUtil.error(f"Unhandled Exception: {str(exc)}", exc_info=True) @@ -81,4 +94,3 @@ class exception_handler(): content=str(exc), headers={"Access-Control-Allow-Origin": "*"} ) -