Files
lang-agent/lang_agent/mcp_server.py
2025-10-14 20:50:33 +08:00

96 lines
2.9 KiB
Python

# https://gofastmcp.com/patterns/decorating-methods
from dataclasses import dataclass, field, is_dataclass
from typing import Type, Literal
import tyro
from fastmcp import FastMCP
from fastapi.middleware.cors import CORSMiddleware
from fastmcp.tools.tool import FunctionTool
from loguru import logger
from lang_agent.rag.simple import SimpleRagConfig
from lang_agent.base import LangToolBase
from lang_agent.config import InstantiateConfig, ToolConfig
from lang_agent.dummy.calculator import Calculator, CalculatorConfig
from catering_end.lang_tool import CartToolConfig, CartTool
@tyro.conf.configure(tyro.conf.SuppressFixed)
@dataclass
class MCPServerConfig(InstantiateConfig):
_target: Type = field(default_factory=lambda: MCPServer)
server_name:str = "langserver"
host: str = "6.6.6.78"
"""host of server"""
port: int = 50051
"""port"""
transport:Literal["stdio", "sse", "streamable-http"] = "streamable-http"
"""transport method"""
# tool configs here;
rag_config: SimpleRagConfig = field(default_factory=SimpleRagConfig)
cart_config: CartToolConfig = field(default_factory=CartToolConfig)
calc_config: CalculatorConfig = field(default_factory=CalculatorConfig)
class MCPServer:
def __init__(self, config: MCPServerConfig):
self.config = config
self.mcp = FastMCP(self.config.server_name)
self.register_mcp_functions()
def _register_tool_fnc(self, tool:LangToolBase):
for fnc in tool.get_tool_fnc():
if isinstance(fnc, FunctionTool):
fnc = fnc.fn
self.mcp.tool(fnc)
def _get_tool_config(self):
tool_confs = []
for e in dir(self.config):
el = getattr(self.config, e)
if ("calc_config" in e) and is_dataclass(el):
tool_confs.append(el)
return tool_confs
def register_mcp_functions(self):
# NOTE: add config here for new tools; too stupid to do this automatically
tool_configs = self._get_tool_config()
for tool_conf in tool_configs:
if tool_conf.use_tool:
logger.info(f"using tool:{tool_conf._target}")
self._register_tool_fnc(tool_conf.setup())
else:
logger.info(f"skipping tool:{tool_conf._target}")
def run(self):
# 获取FastAPI应用实例
app = self.mcp.http_app()
# 配置CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
self.mcp.run(transport=self.config.transport,
host=self.config.host,
port=self.config.port)
if __name__ == "__main__":
conf:MCPServer = MCPServerConfig().setup()
tool_conf = conf._get_tool_config()
for e in tool_conf:
print(e)