96 lines
2.7 KiB
Python
96 lines
2.7 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
MCP服务器示例脚本
|
|||
|
|
---------------
|
|||
|
|
这个脚本演示了如何启动和使用MCP服务器
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import logging
|
|||
|
|
import asyncio
|
|||
|
|
from typing import Dict, Any
|
|||
|
|
|
|||
|
|
# 设置日志级别
|
|||
|
|
logging.basicConfig(
|
|||
|
|
level=logging.INFO,
|
|||
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 导入MCP服务器
|
|||
|
|
from mcp_server import MCPServer
|
|||
|
|
|
|||
|
|
# 自定义资源和工具示例
|
|||
|
|
async def setup_custom_handlers(server: MCPServer):
|
|||
|
|
"""设置自定义资源和工具"""
|
|||
|
|
|
|||
|
|
# 注册自定义资源
|
|||
|
|
@server.register_resource("comfyui://status")
|
|||
|
|
def get_comfy_status() -> str:
|
|||
|
|
"""获取ComfyUI状态信息"""
|
|||
|
|
return "ComfyUI 正在运行中"
|
|||
|
|
|
|||
|
|
# 注册自定义工具
|
|||
|
|
@server.register_tool()
|
|||
|
|
def generate_image(prompt: str, width: int = 512, height: int = 512) -> Dict[str, Any]:
|
|||
|
|
"""生成图像的工具（示例）"""
|
|||
|
|
# 这里只是示例，实际实现需要与ComfyUI集成
|
|||
|
|
return {
|
|||
|
|
"status": "success",
|
|||
|
|
"prompt": prompt,
|
|||
|
|
"dimensions": f"{width}x{height}",
|
|||
|
|
"message": "图像生成请求已提交"
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 注册自定义提示
|
|||
|
|
@server.register_prompt("图像生成")
|
|||
|
|
def image_generation_prompt(style: str, subject: str) -> str:
|
|||
|
|
"""创建图像生成提示"""
|
|||
|
|
return f"请创建一个{style}风格的{subject}图像。"
|
|||
|
|
|
|||
|
|
async def register_models(server: MCPServer):
|
|||
|
|
"""注册示例模型"""
|
|||
|
|
# 注册示例模型
|
|||
|
|
model1 = server.register_model({
|
|||
|
|
"name": "stable-diffusion-v1.5",
|
|||
|
|
"type": "diffusion",
|
|||
|
|
"description": "Stable Diffusion v1.5图像生成模型"
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
model2 = server.register_model({
|
|||
|
|
"name": "llama-3-8b",
|
|||
|
|
"type": "llm",
|
|||
|
|
"description": "Llama 3 8B语言模型"
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
logging.info(f"已注册模型: {model1}, {model2}")
|
|||
|
|
return [model1, model2]
|
|||
|
|
|
|||
|
|
async def main():
|
|||
|
|
"""主函数"""
|
|||
|
|
# 创建MCP服务器实例
|
|||
|
|
server = MCPServer(server_name="ComfyUI-MCP示例")
|
|||
|
|
|
|||
|
|
# 设置自定义处理器
|
|||
|
|
await setup_custom_handlers(server)
|
|||
|
|
|
|||
|
|
# 注册模型
|
|||
|
|
models = await register_models(server)
|
|||
|
|
|
|||
|
|
# 打印服务器信息
|
|||
|
|
logging.info(f"服务器统计信息: {server.get_stats()}")
|
|||
|
|
|
|||
|
|
# 启动服务器（默认地址为127.0.0.1:8189）
|
|||
|
|
try:
|
|||
|
|
logging.info("正在启动MCP服务器...")
|
|||
|
|
await server.start()
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
logging.info("服务器被用户中断")
|
|||
|
|
finally:
|
|||
|
|
# 清理资源
|
|||
|
|
for model_id in models:
|
|||
|
|
server.unregister_model(model_id)
|
|||
|
|
logging.info("服务器已关闭")
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 运行主函数
|
|||
|
|
asyncio.run(main())
|