feat: mcp provide datasource

Author: XiaJunjie2020

backend/apps/mcp/mcp.py

@@ -13,6 +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
+from apps.datasource.crud.datasource import get_datasource_list
 from apps.system.crud.user import authenticate
 from apps.system.crud.user import get_db_user
 from apps.system.models.system_model import UserWsModel
@@ -47,38 +48,10 @@
 #     ))
 
 
-# @router.get("/ds_list", operation_id="get_datasource_list")
-# async def datasource_list(session: SessionDep):
-#     return get_datasource_list(session=session)
-#
-#
-# @router.get("/model_list", operation_id="get_model_list")
-# async def get_model_list(session: SessionDep):
-#     return session.query(AiModelDetail).all()
-
-
-@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")
-
-    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}
-
-
-@router.post("/mcp_question", operation_id="mcp_question")
-async def mcp_question(session: SessionDep, chat: McpQuestion):
+def get_user(session: SessionDep, token: str):
     try:
         payload = jwt.decode(
-            chat.token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+            token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
         )
         token_data = TokenPayload(**payload)
     except (InvalidTokenError, ValidationError):
@@ -91,7 +64,7 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
     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 = chat.lang
+    session_user.language = 'zh'
     if session_user.isAdmin:
         session_user = session_user
     ws_model: UserWsModel = session.exec(
@@ -104,6 +77,42 @@ async def mcp_question(session: SessionDep, chat: McpQuestion):
 
     if session_user.status != 1:
         raise HTTPException(status_code=400, detail="Inactive user")
+    return session_user
+
+
+@router.post("/mcp_ds_list", operation_id="mcp_get_datasource_list")
+async def datasource_list(session: SessionDep, token: str):
+    session_user = get_user(session, token)
+    return get_datasource_list(session=session, user=session_user)
+
+
+#
+#
+# @router.get("/model_list", operation_id="get_model_list")
+# async def get_model_list(session: SessionDep):
+#     return session.query(AiModelDetail).all()
+
+
+@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")
+
+    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}
+
+
+@router.post("/mcp_question", operation_id="mcp_question")
+async def mcp_question(session: SessionDep, chat: McpQuestion):
+    session_user = get_user(session, chat.token)
 
     mcp_chat = ChatMcp(token=chat.token, chat_id=chat.chat_id, question=chat.question)