64 lines
1.8 KiB
Go
64 lines
1.8 KiB
Go
package mcpTool
|
||
|
||
import (
|
||
"context"
|
||
|
||
"github.com/mark3labs/mcp-go/mcp"
|
||
"github.com/mark3labs/mcp-go/server"
|
||
"go.uber.org/zap"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
// McpTool 定义了MCP工具必须实现的接口
|
||
type McpTool interface {
|
||
// Handle 返回工具调用信息
|
||
Handle(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error)
|
||
// New 返回工具注册信息
|
||
New() mcp.Tool
|
||
}
|
||
|
||
// DependencyInjectable 定义了支持依赖注入的工具接口
|
||
type DependencyInjectable interface {
|
||
SetDependencies(db *gorm.DB, logger *zap.Logger)
|
||
}
|
||
|
||
// ConfigInjectable 定义了支持配置注入的工具接口
|
||
type ConfigInjectable interface {
|
||
SetDependencies(db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig)
|
||
}
|
||
|
||
// 工具注册表
|
||
var toolRegister = make(map[string]McpTool)
|
||
|
||
// RegisterTool 供工具在init时调用，将自己注册到工具注册表中
|
||
func RegisterTool(tool McpTool) {
|
||
mcpTool := tool.New()
|
||
toolRegister[mcpTool.Name] = tool
|
||
}
|
||
|
||
// RegisterAllTools 将所有注册的工具注册到MCP服务中
|
||
func RegisterAllTools(mcpServer *server.MCPServer) {
|
||
for _, tool := range toolRegister {
|
||
mcpServer.AddTool(tool.New(), tool.Handle)
|
||
}
|
||
}
|
||
|
||
// RegisterAllToolsWithDependencies 将所有注册的工具注册到MCP服务中，并注入依赖
|
||
func RegisterAllToolsWithDependencies(mcpServer *server.MCPServer, db *gorm.DB, logger *zap.Logger, config *AutoCodeConfig) {
|
||
for _, tool := range toolRegister {
|
||
// 尝试注入配置依赖
|
||
if injectable, ok := tool.(ConfigInjectable); ok {
|
||
injectable.SetDependencies(db, logger, config)
|
||
} else if injectable, ok := tool.(DependencyInjectable); ok {
|
||
// 尝试注入基本依赖
|
||
injectable.SetDependencies(db, logger)
|
||
}
|
||
mcpServer.AddTool(tool.New(), tool.Handle)
|
||
}
|
||
}
|
||
|
||
// GetRegisteredTools 获取所有已注册的工具
|
||
func GetRegisteredTools() map[string]McpTool {
|
||
return toolRegister
|
||
}
|