klinsmaya/SQLBot
1
0
Fork 0
mirror of https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git synced 2026-11-08 01:06:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: MCP mcp/mcp_question support param to disable streaming

Browse Source
This commit is contained in:
ulleo
2025-09-22 18:35:38 +08:00
committed by ulleo
parent f52178f0c4
commit dca52ce1b2
5 changed files with 182 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/chat/models/chat_model.py
+3 -1
View File
@@ -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)
backend/apps/chat/task/llm.py
+70 -40
View File
@@ -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()
backend/apps/mcp/mcp.py
+32 -4
View File
@@ -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")
backend/apps/system/schemas/system_schema.py
+54 -35
View File
@@ -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
backend/common/core/response_middleware.py
+23 -11
View File
@@ -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": "*"}
)