mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
225 lines
9.5 KiB
Python
225 lines
9.5 KiB
Python
# 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"]:
|
|
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)
|