From 35699551b07197e357f92e98f1ac04c494d8c52b Mon Sep 17 00:00:00 2001
From: junjun <junjie.xia@fit2cloud.com>
Date: Fri, 27 Jun 2025 11:16:07 +0800
Subject: [PATCH] feat: mcp tools

---
 backend/apps/chat/api/chat.py     | 31 ++++++++++++-------------------
 backend/common/utils/whitelist.py |  3 ++-
 backend/main.py                   |  2 +-
 3 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py
index d1971558..733b5f70 100644
--- a/backend/apps/chat/api/chat.py
+++ b/backend/apps/chat/api/chat.py
@@ -57,6 +57,18 @@ async def delete(session: SessionDep, chart_id: int):
         )
 
 
+@router.post("/mcp_start", operation_id="mcp_start")
+async def mcp_start(session: SessionDep, token: str, create_chat_obj: CreateChat):
+    user = await get_current_user(session, token)
+    return await start_chat(session, user, create_chat_obj)
+
+
+@router.post("/mcp_question", operation_id="mcp_question")
+async def mcp_question(session: SessionDep, token: str, request_question: ChatQuestion):
+    user = await get_current_user(session, token)
+    return await stream_sql(session, user, request_question)
+
+
 @router.post("/start")
 async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat):
     try:
@@ -68,25 +80,6 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat
         )
 
 
-@router.post("/mcp_question", operation_id="mcp_question")
-async def mcp_question(session: SessionDep, token: str, request_question: ChatQuestion):
-    user = await get_current_user(session, token)
-    # return await stream_sql(session, user, request_question)
-    return {"content": """步骤1: 确定需要查询的字段。
-我们需要统计上海的订单总数，因此需要从"城市"字段中筛选出值为"上海"的记录，并使用COUNT函数计算这些记录的数量。
-
-步骤2: 确定筛选条件。
-问题要求统计上海的订单总数，所以我们需要在SQL语句中添加WHERE "城市" = '上海'来筛选出符合条件的记录。
-
-步骤3: 避免关键字冲突。
-因为这个Excel/CSV数据库是 PostgreSQL 类型，所以在schema、表名、字段名和别名外层加双引号。
-
-最终答案:
-```json
-{"success":true,"sql":"SELECT COUNT(*) AS \"TotalOrders\" FROM \"public\".\"Sheet1_c27345b66e\" WHERE \"城市\" = '上海';"}
-```"""}
-
-
 @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/common/utils/whitelist.py b/backend/common/utils/whitelist.py
index a8b0f42c..c699c5b0 100644
--- a/backend/common/utils/whitelist.py
+++ b/backend/common/utils/whitelist.py
@@ -25,7 +25,8 @@ wlist = [
     "/mcp/message",
     "/datasource/list",
     "/system/aimodel/list",
-    "/chat/mcp_question"
+    "/chat/mcp_question",
+    "/chat/mcp_start"
 ]
 
 class WhitelistChecker:
diff --git a/backend/main.py b/backend/main.py
index a3daf538..198d8b26 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -53,7 +53,7 @@ mcp = FastApiMCP(
     description="SQLBot MCP Server",
     describe_all_responses=True,
     describe_full_response_schema=True,
-    include_operations=["get_datasource_list", "get_model_list", "mcp_question"]
+    include_operations=["get_datasource_list", "get_model_list", "mcp_question", "mcp_start"]
 )
 
 mcp.mount(mcp_app)