feat: Adjust the organizational attribution logic of dialogues in MCP #1221

Author: XiaJunjie2020

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)
 
 

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):