From 2a74b4d69e28db02f2961c1ef40270dd675512fa Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Wed, 11 Jun 2025 14:26:31 +0800
Subject: [PATCH] feat: mcp

---
 backend/apps/chat/api/chat.py      |  2 +-
 backend/apps/system/api/aimodel.py | 43 ++++++++++++++++++------------
 backend/common/utils/whitelist.py  |  5 +++-
 backend/main.py                    |  5 ++--
 4 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py
index e1f1d4dc..536bb9d3 100644
--- a/backend/apps/chat/api/chat.py
+++ b/backend/apps/chat/api/chat.py
@@ -68,7 +68,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
         )
 
 
-@router.post("/question")
+@router.post("/question", operation_id="question")
 async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion):
     """Stream SQL analysis results
     
diff --git a/backend/apps/system/api/aimodel.py b/backend/apps/system/api/aimodel.py
index 7a039e62..a3bf6f21 100644
--- a/backend/apps/system/api/aimodel.py
+++ b/backend/apps/system/api/aimodel.py
@@ -1,6 +1,6 @@
-from datetime import datetime
 from fastapi import APIRouter
-from sqlmodel import select    
+from sqlmodel import select
+
 from apps.system.models.system_model import AiModelDetail
 from apps.system.schemas.system_schema import model_status
 from common.core.deps import SessionDep
@@ -13,9 +13,9 @@ router = APIRouter(tags=["system/aimodel"], prefix="/system/aimodel")
 
 @router.get("/pager/{pageNum}/{pageSize}", response_model=PaginatedResponse[AiModelDetail])
 async def pager(
-    session: SessionDep,
-    pageNum: int,
-    pageSize: int
+        session: SessionDep,
+        pageNum: int,
+        pageSize: int
 ):
     pagination = PaginationParams(page=pageNum, size=pageSize)
     paginator = Paginator(session)
@@ -24,19 +24,21 @@ async def pager(
         model=AiModelDetail,
         pagination=pagination,
         **filters)
-    
+
+
 @router.get("/{id}", response_model=AiModelDetail)
 async def get_model_by_id(
-    session: SessionDep,
-    id: int
+        session: SessionDep,
+        id: int
 ):
     term = session.get(AiModelDetail, id)
     return term
-    
+
+
 @router.post("", response_model=AiModelDetail)
 async def add_model(
-    session: SessionDep,
-    creator: AiModelDetail
+        session: SessionDep,
+        creator: AiModelDetail
 ):
     data = AiModelDetail.model_validate(creator)
     data.create_time = get_timestamp()
@@ -44,10 +46,11 @@ async def add_model(
     session.commit()
     return creator
 
+
 @router.put("", response_model=AiModelDetail)
 async def update_terminology(
-    session: SessionDep,
-    model: AiModelDetail
+        session: SessionDep,
+        model: AiModelDetail
 ):
     model.id = int(model.id)
     term = session.exec(select(AiModelDetail).where(AiModelDetail.id == model.id)).first()
@@ -58,10 +61,11 @@ async def update_terminology(
     session.commit()
     return model
 
+
 @router.delete("/{id}", response_model=AiModelDetail)
 async def delete_terminology(
-    session: SessionDep,
-    id: int
+        session: SessionDep,
+        id: int
 ):
     term = session.exec(select(AiModelDetail).where(AiModelDetail.id == id)).first()
     session.delete(term)
@@ -69,7 +73,7 @@ async def delete_terminology(
     return {
         "message": f"AiModel with ID {id} deleted successfully."
     }
-    
+
 
 @router.patch("/status", response_model=dict)
 async def status(session: SessionDep, dto: model_status):
@@ -83,4 +87,9 @@ async def status(session: SessionDep, dto: model_status):
         term.status = status
         session.add(term)
     session.commit()
-    return {"message": f"AiModel with IDs {ids} updated successfully."}
\ No newline at end of file
+    return {"message": f"AiModel with IDs {ids} updated successfully."}
+
+
+@router.get("/list", operation_id="get_model_list")
+async def get_model_list(session: SessionDep):
+    return session.query(AiModelDetail).all()
diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py
index 78010184..f3543019 100644
--- a/backend/common/utils/whitelist.py
+++ b/backend/common/utils/whitelist.py
@@ -22,7 +22,10 @@ wlist = [
     "*.eot",
     "*.otf",
     "/mcp",
-    "/mcp/message"
+    "/mcp/message",
+    "/datasource/list",
+    "/system/aimodel/list",
+    "/chat/question"
 ]
 
 class WhitelistChecker:
diff --git a/backend/main.py b/backend/main.py
index bef058cb..df21bfc6 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -49,7 +49,7 @@ mcp = FastApiMCP(
     app,
     name="SQLBot MCP Server",
     description="SQLBot MCP Server",
-    include_operations=["get_datasource_list"]
+    include_operations=["get_datasource_list", "get_model_list", "question"]
 )
 
 mcp.mount(mcp_app)
@@ -76,7 +76,7 @@ frontend_dist = os.path.abspath("../frontend/dist")
 if not os.path.exists(frontend_dist):
     logging.warning(f"The front-end build directory does not exist: {frontend_dist}")
     logging.warning("Please make sure you have built the front-end project")
-    
+
 else:
 
     @app.get("/", include_in_schema=False)
@@ -92,3 +92,4 @@ if __name__ == "__main__":
     import uvicorn
 
     uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)
+    # uvicorn.run("main:mcp_app", host="0.0.0.0", port=8001) # mcp server