feat: mcp tools

Author: XiaJunjie2020

backend/apps/mcp/mcp.py

@@ -2,14 +2,12 @@
 # Date: 2025/7/1
 
 from datetime import timedelta
-from typing import Annotated
 
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
-from fastapi.security import OAuth2PasswordRequestForm
 
 from apps.chat.api.chat import create_chat
-from apps.chat.models.chat_model import ChatMcp, CreateChat
+from apps.chat.models.chat_model import ChatMcp, CreateChat, ChatStart
 from apps.chat.task.llm import LLMService, run_task
 from apps.datasource.crud.datasource import get_datasource_list
 from apps.system.crud.user import authenticate
@@ -22,19 +20,19 @@
 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.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")
@@ -48,9 +46,17 @@ async def get_model_list(session: SessionDep):
 
 
 @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)
+async def mcp_start(session: SessionDep, chat: ChatStart):
+    user = authenticate(session=session, account=chat.username, password=chat.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()
+    t = Token(access_token=create_access_token(
+        user_dict, expires_delta=access_token_expires
+    ))
+    c = create_chat(session, user, CreateChat(), False)
+    return {"access_token": t.access_token, "chat_id": c.id}
 
 
 @router.post("/mcp_question", operation_id="mcp_question")

backend/main.py

@@ -15,6 +15,7 @@
 from fastapi.staticfiles import StaticFiles
 import sqlbot_xpack
 
+
 def run_migrations():
     alembic_cfg = Config("alembic.ini")
     command.upgrade(alembic_cfg, "head")
@@ -51,7 +52,7 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     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", "access_token"]
+    include_operations=["get_datasource_list", "get_model_list", "mcp_question", "mcp_start"]
 )
 
 mcp.mount(mcp_app)
@@ -74,8 +75,6 @@ def custom_generate_unique_id(route: APIRoute) -> str:
 app.add_exception_handler(StarletteHTTPException, exception_handler.http_exception_handler)
 app.add_exception_handler(Exception, exception_handler.global_exception_handler)
 
-
-
 mcp.setup_server()
 
 sqlbot_xpack.init_fastapi_app(app)