feat: mcp

Author: XiaJunjie2020

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
     

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.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,30 +24,33 @@ 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()
     session.add(data)
     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,18 +61,19 @@ 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)
     session.commit()
     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."}
+    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()

backend/common/utils/whitelist.py

@@ -22,7 +22,10 @@
     "*.eot",
     "*.otf",
     "/mcp",
-    "/mcp/message"
+    "/mcp/message",
+    "/datasource/list",
+    "/system/aimodel/list",
+    "/chat/question"
 ]
 
 class WhitelistChecker:

backend/main.py

@@ -49,7 +49,7 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     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 @@ def custom_generate_unique_id(route: APIRoute) -> str:
 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 @@ async def read_index():
     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