From 93798b7c58e15ef8f903042d841297b5000c45b8 Mon Sep 17 00:00:00 2001 From: junjun Date: Fri, 18 Jul 2025 10:14:41 +0800 Subject: [PATCH] feat: mcp tools --- backend/apps/mcp/mcp.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py index 42f945e5..d5b2dd94 100644 --- a/backend/apps/mcp/mcp.py +++ b/backend/apps/mcp/mcp.py @@ -9,11 +9,15 @@ from fastapi.responses import StreamingResponse # 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 from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart from apps.chat.task.llm import LLMService, run_task -from apps.system.crud.user import authenticate, get_user_info +from apps.system.crud.user import authenticate +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 from apps.system.schemas.system_schema import UserInfoDTO from common.core import security @@ -83,7 +87,17 @@ async def mcp_question(session: SessionDep, chat: ChatMcp): status_code=status.HTTP_403_FORBIDDEN, detail="Could not validate credentials", ) - session_user = await get_user_info(session=session, user_id=token_data.id) + # 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' + 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")