mcp接口flask改成fastapi

api/apps/mcp_server_app.py

@@ -13,8 +13,9 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-from flask import Response, request
+from typing import List, Optional, Dict, Any
-from flask_login import current_user, login_required
+from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 
 from api.db import VALID_MCP_SERVER_TYPES
 from api.db.db_models import MCPServer
@@ -23,42 +24,161 @@ from api.db.services.user_service import TenantService
 from api.settings import RetCode
 
 from api.utils import get_uuid
-from api.utils.api_utils import get_data_error_result, get_json_result, server_error_response, validate_request, \
+from api.utils.api_utils import get_data_error_result, get_json_result, server_error_response, \
     get_mcp_tools
 from api.utils.web_utils import get_float, safe_json_parse
 from rag.utils.mcp_tool_call_conn import MCPToolCallSession, close_multiple_mcp_toolcall_sessions
+from pydantic import BaseModel
 
+# Security
+security = HTTPBearer()
 
-@manager.route("/list", methods=["POST"])  # noqa: F821
+# Pydantic models for request/response
-@login_required
+class ListMCPRequest(BaseModel):
-def list_mcp() -> Response:
+    mcp_ids: List[str] = []
-    keywords = request.args.get("keywords", "")
+
-    page_number = int(request.args.get("page", 0))
+class CreateMCPRequest(BaseModel):
-    items_per_page = int(request.args.get("page_size", 0))
+    name: str
-    orderby = request.args.get("orderby", "create_time")
+    url: str
-    if request.args.get("desc", "true").lower() == "false":
+    server_type: str
-        desc = False
+    headers: Optional[Dict[str, Any]] = {}
+    variables: Optional[Dict[str, Any]] = {}
+    timeout: Optional[float] = 10
+
+class UpdateMCPRequest(BaseModel):
+    mcp_id: str
+    name: Optional[str] = None
+    url: Optional[str] = None
+    server_type: Optional[str] = None
+    headers: Optional[Dict[str, Any]] = None
+    variables: Optional[Dict[str, Any]] = None
+    timeout: Optional[float] = None
+
+class RemoveMCPRequest(BaseModel):
+    mcp_ids: List[str]
+
+class ImportMCPRequest(BaseModel):
+    mcpServers: Dict[str, Dict[str, Any]]
+    timeout: Optional[float] = 10
+
+class ExportMCPRequest(BaseModel):
+    mcp_ids: List[str]
+
+class ListToolsRequest(BaseModel):
+    mcp_ids: List[str]
+    timeout: Optional[float] = 10
+
+class TestToolRequest(BaseModel):
+    mcp_id: str
+    tool_name: str
+    arguments: Dict[str, Any]
+    timeout: Optional[float] = 10
+
+class CacheToolsRequest(BaseModel):
+    mcp_id: str
+    tools: List[Dict[str, Any]]
+
+class TestMCPRequest(BaseModel):
+    url: str
+    server_type: str
+    headers: Optional[Dict[str, Any]] = {}
+    variables: Optional[Dict[str, Any]] = {}
+    timeout: Optional[float] = 10
+
+# Dependency injection
+async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """获取当前用户"""
+    from api.db import StatusEnum
+    from api.db.services.user_service import UserService
+    from fastapi import HTTPException, status
+    import logging
+    
+    try:
+        from itsdangerous.url_safe import URLSafeTimedSerializer as Serializer
+    except ImportError:
+        # 如果没有itsdangerous，使用jwt作为替代
+        import jwt
+        Serializer = jwt
+    
+    jwt = Serializer(secret_key=settings.SECRET_KEY)
+    authorization = credentials.credentials
+    
+    if authorization:
+        try:
+            access_token = str(jwt.loads(authorization))
+            
+            if not access_token or not access_token.strip():
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Authentication attempt with empty access token"
+                )
+            
+            # Access tokens should be UUIDs (32 hex characters)
+            if len(access_token.strip()) < 32:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail=f"Authentication attempt with invalid token format: {len(access_token)} chars"
+                )
+            
+            user = UserService.query(
+                access_token=access_token, status=StatusEnum.VALID.value
+            )
+            if user:
+                if not user[0].access_token or not user[0].access_token.strip():
+                    raise HTTPException(
+                        status_code=status.HTTP_401_UNAUTHORIZED,
+                        detail=f"User {user[0].email} has empty access_token in database"
+                    )
+                return user[0]
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Invalid access token"
+                )
+        except Exception as e:
+            logging.warning(f"load_user got exception {e}")
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid access token"
+            )
     else:
-        desc = True
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Authorization header required"
+        )
 
-    req = request.get_json()
+# Create router
-    mcp_ids = req.get("mcp_ids", [])
+router = APIRouter()
+
+
+@router.post("/list")
+async def list_mcp(
+    keywords: str = Query(""),
+    page: int = Query(0),
+    page_size: int = Query(0),
+    orderby: str = Query("create_time"),
+    desc: bool = Query(True),
+    req: ListMCPRequest = None,
+    current_user = Depends(get_current_user)
+):
+    mcp_ids = req.mcp_ids if req else []
     try:
         servers = MCPServerService.get_servers(current_user.id, mcp_ids, 0, 0, orderby, desc, keywords) or []
         total = len(servers)
 
-        if page_number and items_per_page:
+        if page and page_size:
-            servers = servers[(page_number - 1) * items_per_page : page_number * items_per_page]
+            servers = servers[(page - 1) * page_size : page * page_size]
 
         return get_json_result(data={"mcp_servers": servers, "total": total})
     except Exception as e:
         return server_error_response(e)
 
 
-@manager.route("/detail", methods=["GET"])  # noqa: F821
+@router.get("/detail")
-@login_required
+async def detail(
-def detail() -> Response:
+    mcp_id: str = Query(...),
-    mcp_id = request.args["mcp_id"]
+    current_user = Depends(get_current_user)
+):
     try:
         mcp_server = MCPServerService.get_or_none(id=mcp_id, tenant_id=current_user.id)
 
@@ -70,17 +190,16 @@ def detail() -> Response:
         return server_error_response(e)
 
 
-@manager.route("/create", methods=["POST"])  # noqa: F821
+@router.post("/create")
-@login_required
+async def create(
-@validate_request("name", "url", "server_type")
+    req: CreateMCPRequest,
-def create() -> Response:
+    current_user = Depends(get_current_user)
-    req = request.get_json()
+):
-
+    server_type = req.server_type
-    server_type = req.get("server_type", "")
     if server_type not in VALID_MCP_SERVER_TYPES:
         return get_data_error_result(message="Unsupported MCP server type.")
 
-    server_name = req.get("name", "")
+    server_name = req.name
     if not server_name or len(server_name.encode("utf-8")) > 255:
         return get_data_error_result(message=f"Invalid MCP name or length is {len(server_name)} which is large than 255.")
 
@@ -88,20 +207,20 @@ def create() -> Response:
     if e:
         return get_data_error_result(message="Duplicated MCP server name.")
 
-    url = req.get("url", "")
+    url = req.url
     if not url:
         return get_data_error_result(message="Invalid url.")
 
-    headers = safe_json_parse(req.get("headers", {}))
+    headers = safe_json_parse(req.headers)
-    req["headers"] = headers
+    variables = safe_json_parse(req.variables)
-    variables = safe_json_parse(req.get("variables", {}))
     variables.pop("tools", None)
 
-    timeout = get_float(req, "timeout", 10)
+    timeout = req.timeout
 
     try:
-        req["id"] = get_uuid()
+        req_dict = req.dict()
-        req["tenant_id"] = current_user.id
+        req_dict["id"] = get_uuid()
+        req_dict["tenant_id"] = current_user.id
 
         e, _ = TenantService.get_by_id(current_user.id)
         if not e:
@@ -115,49 +234,47 @@ def create() -> Response:
         tools = server_tools[server_name]
         tools = {tool["name"]: tool for tool in tools if isinstance(tool, dict) and "name" in tool}
         variables["tools"] = tools
-        req["variables"] = variables
+        req_dict["variables"] = variables
 
-        if not MCPServerService.insert(**req):
+        if not MCPServerService.insert(**req_dict):
             return get_data_error_result("Failed to create MCP server.")
 
-        return get_json_result(data=req)
+        return get_json_result(data=req_dict)
     except Exception as e:
         return server_error_response(e)
 
 
-@manager.route("/update", methods=["POST"])  # noqa: F821
+@router.post("/update")
-@login_required
+async def update(
-@validate_request("mcp_id")
+    req: UpdateMCPRequest,
-def update() -> Response:
+    current_user = Depends(get_current_user)
-    req = request.get_json()
+):
-
+    mcp_id = req.mcp_id
-    mcp_id = req.get("mcp_id", "")
     e, mcp_server = MCPServerService.get_by_id(mcp_id)
     if not e or mcp_server.tenant_id != current_user.id:
         return get_data_error_result(message=f"Cannot find MCP server {mcp_id} for user {current_user.id}")
 
-    server_type = req.get("server_type", mcp_server.server_type)
+    server_type = req.server_type if req.server_type is not None else mcp_server.server_type
     if server_type and server_type not in VALID_MCP_SERVER_TYPES:
         return get_data_error_result(message="Unsupported MCP server type.")
-    server_name = req.get("name", mcp_server.name)
+    server_name = req.name if req.name is not None else mcp_server.name
     if server_name and len(server_name.encode("utf-8")) > 255:
         return get_data_error_result(message=f"Invalid MCP name or length is {len(server_name)} which is large than 255.")
-    url = req.get("url", mcp_server.url)
+    url = req.url if req.url is not None else mcp_server.url
     if not url:
         return get_data_error_result(message="Invalid url.")
 
-    headers = safe_json_parse(req.get("headers", mcp_server.headers))
+    headers = safe_json_parse(req.headers if req.headers is not None else mcp_server.headers)
-    req["headers"] = headers
+    variables = safe_json_parse(req.variables if req.variables is not None else mcp_server.variables)
-
-    variables = safe_json_parse(req.get("variables", mcp_server.variables))
     variables.pop("tools", None)
 
-    timeout = get_float(req, "timeout", 10)
+    timeout = req.timeout if req.timeout is not None else 10
 
     try:
-        req["tenant_id"] = current_user.id
+        req_dict = req.dict(exclude_unset=True)
-        req.pop("mcp_id", None)
+        req_dict["tenant_id"] = current_user.id
-        req["id"] = mcp_id
+        req_dict.pop("mcp_id", None)
+        req_dict["id"] = mcp_id
 
         mcp_server = MCPServer(id=server_name, name=server_name, url=url, server_type=server_type, variables=variables, headers=headers)
         server_tools, err_message = get_mcp_tools([mcp_server], timeout)
@@ -167,12 +284,12 @@ def update() -> Response:
         tools = server_tools[server_name]
         tools = {tool["name"]: tool for tool in tools if isinstance(tool, dict) and "name" in tool}
         variables["tools"] = tools
-        req["variables"] = variables
+        req_dict["variables"] = variables
 
-        if not MCPServerService.filter_update([MCPServer.id == mcp_id, MCPServer.tenant_id == current_user.id], req):
+        if not MCPServerService.filter_update([MCPServer.id == mcp_id, MCPServer.tenant_id == current_user.id], req_dict):
             return get_data_error_result(message="Failed to updated MCP server.")
 
-        e, updated_mcp = MCPServerService.get_by_id(req["id"])
+        e, updated_mcp = MCPServerService.get_by_id(req_dict["id"])
         if not e:
             return get_data_error_result(message="Failed to fetch updated MCP server.")