From 69d719e04ee5cba72e3c4e3b0a9fc126fa8d0b0a Mon Sep 17 00:00:00 2001 From: junjun Date: Mon, 29 Jun 2026 15:24:53 +0800 Subject: [PATCH] feat: Adjust the organizational attribution logic of dialogues in MCP #1221 --- backend/apps/chat/models/chat_model.py | 12 +++- backend/apps/mcp/mcp.py | 84 ++++++++++++++++---------- 2 files changed, 60 insertions(+), 36 deletions(-) diff --git a/backend/apps/chat/models/chat_model.py b/backend/apps/chat/models/chat_model.py index 3fa50ab8..dbe6d058 100644 --- a/backend/apps/chat/models/chat_model.py +++ b/backend/apps/chat/models/chat_model.py @@ -353,11 +353,19 @@ class McpDs(BaseModel): oid: Optional[str] = Body(description='组织ID，如果不传则为最后一次登录SQLBot时所使用的组织ID', default=None) -class ChatStart(BaseModel): +class ChatToken(BaseModel): username: str = Body(description='用户名') password: str = Body(description='密码') +class ChatStart(BaseModel): + username: str = Body(description='用户名', default=None) + password: str = Body(description='密码', default=None) + token: str = Body(description='token', default=None) + oid: Optional[str] = Body( + description='组织ID，仅当数据源ID为空时有效，如果不传则为最后一次登录SQLBot时所使用的组织ID', default=None) + + class ChatQuestionBase(BaseModel): question: str = Body(description='用户提问') chat_id: int = Body(description='会话ID') @@ -368,8 +376,6 @@ class McpQuestion(ChatQuestionBase): stream: Optional[bool] = Body(description='是否流式输出，默认为true开启, 关闭false则返回JSON对象', default=True) lang: Optional[str] = Body(description='语言：zh-CN|zh-TW|en|ko-KR', default='zh-CN') 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) diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 6185ff96..e018c672 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -13,7 +13,7 @@ from apps.chat.api.chat import create_chat, question_answer_inner from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart, McpQuestion, McpAssistant, ChatQuestion, \ - ChatFinishStep, McpDs + ChatFinishStep, McpDs, ChatToken from apps.datasource.crud.datasource import get_datasource_list from apps.system.crud.user import authenticate, user_ws_options from apps.system.crud.user import get_db_user @@ -34,19 +34,21 @@ router = APIRouter(tags=["mcp"], prefix="/mcp") -# @router.post("/access_token", operation_id="access_token") -# def local_login( -# session: SessionDep, -# form_data: Annotated[OAuth2PasswordRequestForm, Depends()] -# ) -> Token: -# user = authenticate(session=session, account=form_data.username, password=form_data.password) -# if not user: -# raise HTTPException(status_code=400, detail="Incorrect account or password") -# access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) -# user_dict = user.to_dict() -# return Token(access_token=create_access_token( -# user_dict, expires_delta=access_token_expires -# )) +@router.post("/access_token", operation_id="access_token") +async def access_token(session: SessionDep, chat: ChatToken): + user: BaseUserDTO = authenticate(session=session, account=chat.username, password=chat.password) + if not user: + raise HTTPException(status_code=400, detail="Incorrect account or password") + + if not user.oid or user.oid == 0: + raise HTTPException(status_code=400, detail="No associated workspace, Please contact the administrator") + access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) + user_dict = user.to_dict() + t = Token(access_token=create_access_token( + user_dict, expires_delta=access_token_expires + )) + # c = create_chat(session, user, CreateChat(origin=1), False) + return {"access_token": t.access_token} def get_user(session: SessionDep, token: str): @@ -82,20 +84,36 @@ def get_user(session: SessionDep, token: str): @router.post("/mcp_start", operation_id="mcp_start") -async def mcp_start(session: SessionDep, chat: ChatStart): - user: BaseUserDTO = authenticate(session=session, account=chat.username, password=chat.password) - if not user: - raise HTTPException(status_code=400, detail="Incorrect account or password") +async def mcp_start(session: SessionDep, trans: Trans, chat: ChatStart): + res_token = None + user = None + if chat.token: + res_token = chat.token + user = get_user(session, chat.token) + else: + user = authenticate(session=session, account=chat.username, password=chat.password) + if not user: + raise HTTPException(status_code=400, detail="Incorrect account or password") + + if not user.oid or user.oid == 0: + raise HTTPException(status_code=400, detail="No associated workspace, Please contact the administrator") + access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) + user_dict = user.to_dict() + t = Token(access_token=create_access_token( + user_dict, expires_delta=access_token_expires + )) + res_token = t.access_token + + if chat.oid: + w_list = await user_ws_options(session, user.id, trans) + oid_list = [item.id for item in w_list] + if int(chat.oid) not in oid_list: + raise HTTPException(status_code=400, detail="The current user is not in the selected workspace") + + user.oid = int(chat.oid) - if not user.oid or user.oid == 0: - raise HTTPException(status_code=400, detail="No associated workspace, Please contact the administrator") - access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) - user_dict = user.to_dict() - t = Token(access_token=create_access_token( - user_dict, expires_delta=access_token_expires - )) c = create_chat(session, user, CreateChat(origin=1), False) - return {"access_token": t.access_token, "chat_id": c.id} + return {"access_token": res_token, "chat_id": c.id} @router.post("/mcp_ws_list", operation_id="mcp_ws_list") @@ -139,13 +157,13 @@ async def mcp_question(session: SessionDep, trans: Trans, chat: McpQuestion): lang = chat.lang if lang in ["zh-CN", "zh-TW", "en", "ko-KR"]: session_user.language = lang - if chat.oid: - w_list = await user_ws_options(session, session_user.id, trans) - oid_list = [item.id for item in w_list] - if int(chat.oid) not in oid_list: - raise HTTPException(status_code=400, detail="The current user is not in the selected workspace") - - session_user.oid = int(chat.oid) + # if chat.oid: + # w_list = await user_ws_options(session, session_user.id, trans) + # oid_list = [item.id for item in w_list] + # if int(chat.oid) not in oid_list: + # raise HTTPException(status_code=400, detail="The current user is not in the selected workspace") + # + # session_user.oid = int(chat.oid) ds_id: Optional[int] = None if chat.datasource_id: if isinstance(chat.datasource_id, str):