SQLBot/backend/apps/mcp/mcp.py

# Author: Junjun
# Date: 2025/7/1
import json
from datetime import timedelta
from typing import Optional

import jwt
from fastapi import HTTPException, status, APIRouter
# from fastapi.security import OAuth2PasswordBearer
from jwt.exceptions import InvalidTokenError
from pydantic import ValidationError
from sqlmodel import select

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, ChatToken, WsMcp
from apps.datasource.crud.datasource import get_datasource_list
from apps.system.crud.aimodel_manage import get_ai_model_list_by_workspace
from apps.system.crud.user import authenticate, user_ws_options
from apps.system.crud.user import get_db_user
from apps.system.models.system_model import UserWsModel
from apps.system.models.user import UserModel
from apps.system.schemas.system_schema import BaseUserDTO, AssistantHeader
from apps.system.schemas.system_schema import UserInfoDTO
from common.audit.models.log_model import OperationType, OperationModules
from common.audit.schemas.logger_decorator import LogConfig, system_log
from common.audit.schemas.request_context import RequestContext
from common.core import security
from common.core.config import settings
from common.core.deps import SessionDep, Trans
from common.core.schemas import TokenPayload, XOAuth2PasswordBearer, Token
from common.core.security import create_access_token

reusable_oauth2 = XOAuth2PasswordBearer(
    tokenUrl=f"{settings.API_V1_STR}/login/access-token"
)

router = APIRouter(tags=["mcp"], prefix="/mcp")


@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):
    try:
        payload = jwt.decode(
            token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
        )
        token_data = TokenPayload(**payload)
    except (InvalidTokenError, ValidationError):
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Could not validate credentials",
        )
    # session_user = await get_user_info(session=session, user_id=token_data.id)

    db_user: UserModel = get_db_user(session=session, user_id=token_data.id)
    session_user = UserInfoDTO.model_validate(db_user.model_dump())
    session_user.isAdmin = session_user.id == 1 and session_user.account == 'admin'
    session_user.language = 'zh-CN'
    if session_user.isAdmin:
        session_user = session_user
    ws_model: UserWsModel = session.exec(
        select(UserWsModel).where(UserWsModel.uid == session_user.id, UserWsModel.oid == session_user.oid)).first()
    session_user.weight = ws_model.weight if ws_model else -1

    session_user = UserInfoDTO.model_validate(session_user)
    if not session_user:
        raise HTTPException(status_code=404, detail="User not found")

    if session_user.status != 1:
        raise HTTPException(status_code=400, detail="Inactive user")
    return session_user


@router.post("/mcp_start", operation_id="mcp_start")
@system_log(LogConfig(
    operation_type=OperationType.CREATE,
    module=OperationModules.CHAT,
    result_id_expr="id",
    save_on_success_only=True
))
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)

    request = RequestContext.get_request()
    request.state.current_user = user

    c = create_chat(session, user, CreateChat(origin=1), False)
    return {"access_token": res_token, "chat_id": c.id}


@router.post("/mcp_ws_list", operation_id="mcp_ws_list")
async def ws_list(session: SessionDep, trans: Trans, token: str):
    session_user = get_user(session, token)
    return await user_ws_options(session, session_user.id, trans)


@router.post("/mcp_ds_list", operation_id="mcp_datasource_list")
async def datasource_list(session: SessionDep, trans: Trans, mcp_ds: McpDs):
    session_user = get_user(session, mcp_ds.token)
    if mcp_ds.oid:
        w_list = await user_ws_options(session, session_user.id, trans)
        oid_list = [item.id for item in w_list]
        if int(mcp_ds.oid) not in oid_list:
            raise HTTPException(status_code=400, detail="The current user is not in the selected workspace")

        session_user.oid = int(mcp_ds.oid)
    ds_list = get_datasource_list(session=session, user=session_user)
    result = []
    for item in ds_list:
        dic = item.__dict__
        dic.pop('embedding', None)
        dic.pop('table_relation', None)
        dic.pop('recommended_config', None)
        dic.pop('configuration', None)
        result.append(dic)
    return result


@router.post("/mcp_model_list", operation_id="mcp_model_list")
async def get_model_by_ws(session: SessionDep, mcp_oid: WsMcp):
    return get_ai_model_list_by_workspace(session, int(mcp_oid.oid), False)


@router.post("/mcp_question", operation_id="mcp_question")
async def mcp_question(session: SessionDep, trans: Trans, chat: McpQuestion):
    session_user = get_user(session, chat.token)
    lang = chat.lang
    if lang in ["zh-CN", "zh-TW", "en", "ko-KR", "th-TH"]:
        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)
    ds_id: Optional[int] = None
    if chat.datasource_id:
        if isinstance(chat.datasource_id, str):
            if chat.datasource_id.strip() == "":
                ds_id = None
            else:
                try:
                    ds_id = int(chat.datasource_id.strip())
                except ValueError:
                    raise HTTPException(status_code=400, detail="Invalid datasource ID")
        elif isinstance(chat.datasource_id, int):
            ds_id = chat.datasource_id
        else:
            raise HTTPException(status_code=400, detail="Invalid datasource ID")

    mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question, datasource_id=ds_id,
                       custom_model=chat.custom_model)

    return await question_answer_inner(session=session, current_user=session_user, request_question=mcp_chat,
                                       in_chat=False, stream=chat.stream, return_img=chat.return_img)


# Cordys crm
@router.post("/mcp_assistant", operation_id="mcp_assistant")
async def mcp_assistant(session: SessionDep, chat: McpAssistant):
    session_user = BaseUserDTO(**{
        "id": -1, "account": 'sqlbot-mcp-assistant', "oid": 1, "assistant_id": -1, "password": '', "language": "zh-CN"
    })
    # session_user: UserModel = get_db_user(session=session, user_id=1)
    # session_user.oid = 1
    c = create_chat(session, session_user, CreateChat(origin=1), False)

    # build assistant param
    configuration = {"endpoint": chat.url}
    # authorization = [{"key": "x-de-token",
    #                 "value": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJ1aWQiOjEsIm9pZCI6MSwiZXhwIjoxNzU4NTEyMDA2fQ.3NR-pgnADLdXZtI3dXX5-LuxfGYRvYD9kkr2de7KRP0",
    #                 "target": "header"}]
    mcp_assistant_header = AssistantHeader(id=1, name='mcp_assist', domain='', type=1,
                                           configuration=json.dumps(configuration),
                                           certificate=chat.authorization)

    # assistant question
    mcp_chat = ChatQuestion(chat_id=c.id, question=chat.question)
    # ask
    return await question_answer_inner(session=session, current_user=session_user, request_question=mcp_chat,
                                       current_assistant=mcp_assistant_header,
                                       in_chat=False, stream=chat.stream, finish_step=ChatFinishStep.QUERY_DATA)