From 7f641f98c0e43db507f2c8d385ba0aaf5c8c9753 Mon Sep 17 00:00:00 2001 From: junjun Date: Tue, 1 Jul 2025 17:31:03 +0800 Subject: [PATCH] feat: mcp server --- backend/apps/api.py | 11 +++- backend/apps/chat/api/chat.py | 30 +--------- backend/apps/datasource/api/datasource.py | 2 +- backend/apps/mcp/__init__.py | 2 + backend/apps/mcp/mcp.py | 71 +++++++++++++++++++++++ backend/apps/system/api/aimodel.py | 5 -- backend/common/utils/whitelist.py | 8 +-- backend/main.py | 2 +- 8 files changed, 88 insertions(+), 43 deletions(-) create mode 100644 backend/apps/mcp/__init__.py create mode 100644 backend/apps/mcp/mcp.py diff --git a/backend/apps/api.py b/backend/apps/api.py index c20d842c..9702c247 100644 --- a/backend/apps/api.py +++ b/backend/apps/api.py @@ -1,10 +1,11 @@ from fastapi import APIRouter -from apps.system.api import login, user, aimodel -from apps.settings.api import terminology -from apps.datasource.api import datasource from apps.chat.api import chat from apps.dashboard.api import dashboard_api +from apps.datasource.api import datasource +from apps.settings.api import terminology +from apps.system.api import login, user, aimodel +from apps.mcp import mcp api_router = APIRouter() api_router.include_router(login.router) @@ -12,5 +13,9 @@ api_router.include_router(user.router) api_router.include_router(aimodel.router) api_router.include_router(terminology.router) api_router.include_router(datasource.router) +# api_router.include_router(row_permission.router) +# api_router.include_router(column_permission.router) api_router.include_router(chat.router) api_router.include_router(dashboard_api.router) +api_router.include_router(mcp.router) + diff --git a/backend/apps/chat/api/chat.py b/backend/apps/chat/api/chat.py index 7ca40aee..674494c4 100644 --- a/backend/apps/chat/api/chat.py +++ b/backend/apps/chat/api/chat.py @@ -57,34 +57,6 @@ async def delete(session: SessionDep, chart_id: int): ) -@router.post("/mcp_start", operation_id="mcp_start") -async def mcp_start(session: SessionDep, chat: ChatMcp): - user = await get_current_user(session, chat.token) - return create_chat(session, user, CreateChat(), False) - - -@router.post("/mcp_question", operation_id="mcp_question") -async def mcp_question(session: SessionDep, chat: ChatMcp): - user = await get_current_user(session, chat.token) - # return await stream_sql(session, user, chat) - 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("/start") async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat_obj: CreateChat): try: @@ -96,7 +68,7 @@ async def start_chat(session: SessionDep, current_user: CurrentUser, create_chat ) -@router.post("/question", operation_id="question") +@router.post("/question") async def stream_sql(session: SessionDep, current_user: CurrentUser, request_question: ChatQuestion): """Stream SQL analysis results diff --git a/backend/apps/datasource/api/datasource.py b/backend/apps/datasource/api/datasource.py index 08f2e6f7..b0921673 100644 --- a/backend/apps/datasource/api/datasource.py +++ b/backend/apps/datasource/api/datasource.py @@ -18,7 +18,7 @@ router = APIRouter(tags=["datasource"], prefix="/datasource") path = "/opt/sqlbot/data/excel" -@router.get("/list", operation_id="get_datasource_list") +@router.get("/list") async def datasource_list(session: SessionDep): return get_datasource_list(session=session) diff --git a/backend/apps/mcp/__init__.py b/backend/apps/mcp/__init__.py new file mode 100644 index 00000000..a14a79d9 --- /dev/null +++ b/backend/apps/mcp/__init__.py @@ -0,0 +1,2 @@ +# Author: Junjun +# Date: 2025/7/1 diff --git a/backend/apps/mcp/mcp.py b/backend/apps/mcp/mcp.py new file mode 100644 index 00000000..5f33a0b0 --- /dev/null +++ b/backend/apps/mcp/mcp.py @@ -0,0 +1,71 @@ +# Author: Junjun +# Date: 2025/7/1 + +from typing import Annotated +from fastapi import APIRouter, Depends, HTTPException +from fastapi.security import OAuth2PasswordRequestForm +from common.core.deps import SessionDep, get_current_user +from apps.system.crud.user import authenticate +from common.core.security import create_access_token +from datetime import timedelta +from common.core.config import settings +from common.core.schemas import Token +from apps.datasource.crud.datasource import get_datasource_list +from apps.system.models.system_model import AiModelDetail +from apps.chat.models.chat_model import ChatMcp, CreateChat +from apps.chat.api.chat import create_chat, stream_sql + +router = APIRouter(tags=["mcp"], prefix="/mcp") + + +@router.post("/access_token", operation_id="access_token") +def local_login( + session: SessionDep, + form_data: Annotated[OAuth2PasswordRequestForm, Depends()] +) -> Token: + user = authenticate(session=session, account=form_data.username, password=form_data.password) + if not user: + raise HTTPException(status_code=400, detail="Incorrect account or password") + access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES) + user_dict = user.to_dict() + return Token(access_token=create_access_token( + user_dict, expires_delta=access_token_expires + )) + + +@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: ChatMcp): + user = await get_current_user(session, chat.token) + return create_chat(session, user, CreateChat(), False) + + +@router.post("/mcp_question", operation_id="mcp_question") +async def mcp_question(session: SessionDep, chat: ChatMcp): + user = await get_current_user(session, chat.token) + # return await stream_sql(session, user, chat) + 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 \"城市\" = '上海';"} + ``` + """} \ No newline at end of file diff --git a/backend/apps/system/api/aimodel.py b/backend/apps/system/api/aimodel.py index a3bf6f21..eabc1fd6 100644 --- a/backend/apps/system/api/aimodel.py +++ b/backend/apps/system/api/aimodel.py @@ -88,8 +88,3 @@ async def status(session: SessionDep, dto: model_status): session.add(term) session.commit() 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 07598e75..d3f23d33 100644 --- a/backend/common/utils/whitelist.py +++ b/backend/common/utils/whitelist.py @@ -23,10 +23,10 @@ wlist = [ "*.otf", "/mcp", "/mcp/message", - "/datasource/list", - "/system/aimodel/list", - "/chat/mcp_question", - "/chat/mcp_start", + "/mcp/ds_list", + "/mcp/model_list", + "/mcp/mcp_question", + "/mcp/mcp_start", "/system/license", "/system/config/key", "/images/*" diff --git a/backend/main.py b/backend/main.py index 0066aba6..7b8310d7 100644 --- a/backend/main.py +++ b/backend/main.py @@ -51,7 +51,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", "mcp_start"] + include_operations=["get_datasource_list", "get_model_list", "mcp_question", "mcp_start", "access_token"] ) mcp.mount(mcp_app)