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()) |