74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
import argparse
|
||
import os
|
||
import time
|
||
|
||
import mcp.types as mt
|
||
from fastmcp import FastMCP
|
||
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
||
from main import application
|
||
from fastmcp.server.providers.openapi import MCPType
|
||
from fastmcp.tools.base import ToolResult
|
||
from config.main import config
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description='Dcalc MCP server')
|
||
parser.add_argument(
|
||
'--tool-result-mode',
|
||
choices=['structured', 'both'],
|
||
default='structured',
|
||
help='控制工具返回内容：structured 仅保留 structuredContent；both 保留默认双份。',
|
||
)
|
||
args, _ = parser.parse_known_args()
|
||
return args
|
||
|
||
|
||
class ToolResultModeMiddleware(Middleware):
|
||
def __init__(self, mode: str):
|
||
self.mode = mode
|
||
|
||
async def on_call_tool(
|
||
self,
|
||
context: MiddlewareContext[mt.CallToolRequestParams],
|
||
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
|
||
) -> ToolResult:
|
||
result = await call_next(context)
|
||
if result.structured_content is None:
|
||
return result
|
||
if self.mode == 'structured':
|
||
return ToolResult(
|
||
content=[],
|
||
structured_content=result.structured_content,
|
||
meta=result.meta,
|
||
)
|
||
return result
|
||
|
||
|
||
def custom_route_mapper(route, mcp_type: MCPType) -> MCPType | None:
|
||
"""Advanced route type mapping."""
|
||
# Convert all admin routes to tools regardless of HTTP method
|
||
if "/open/" in route.path and route.method in ["GET"]:
|
||
return MCPType.TOOL
|
||
else:
|
||
return MCPType.EXCLUDE
|
||
|
||
|
||
def create_mcp(tool_result_mode: str = 'structured') -> FastMCP:
|
||
mcp = FastMCP.from_fastapi(
|
||
app=application,
|
||
name="Dcalc_MCP",
|
||
instructions="DNF相关的MCP服务",
|
||
route_map_fn=custom_route_mapper,
|
||
)
|
||
mcp.add_middleware(ToolResultModeMiddleware(tool_result_mode))
|
||
return mcp
|
||
|
||
|
||
os.environ['TZ'] = 'Asia/Shanghai'
|
||
time.tzset()
|
||
|
||
if __name__ == '__main__':
|
||
args = parse_args()
|
||
mcp = create_mcp(args.tool_result_mode)
|
||
mcp.run(transport="streamable-http", host="0.0.0.0", port=config.MCPPORT)
|