wehub-resource-sync
133 lines
4.5 KiB
Go
133 lines
4.5 KiB
Go
package mcp
|
|
|
|
import (
|
|
"context"
|
|
"slices"
|
|
|
|
"github.com/mudler/LocalAI/core/config"
|
|
mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
|
|
"github.com/mudler/LocalAI/pkg/functions"
|
|
"github.com/mudler/xlog"
|
|
)
|
|
|
|
// ToolExecutor abstracts MCP tool discovery and execution.
|
|
// Implementations handle local (in-process sessions) vs distributed (NATS) modes.
|
|
type ToolExecutor interface {
|
|
// DiscoverTools returns the tool function schemas available from MCP servers.
|
|
DiscoverTools(ctx context.Context) ([]functions.Function, error)
|
|
// IsTool returns true if the given function name is an MCP tool.
|
|
IsTool(name string) bool
|
|
// ExecuteTool executes an MCP tool by name with the given JSON arguments.
|
|
ExecuteTool(ctx context.Context, toolName, arguments string) (string, error)
|
|
// HasTools returns true if any MCP tools are available.
|
|
HasTools() bool
|
|
}
|
|
|
|
// LocalToolExecutor uses in-process MCP sessions for tool operations.
|
|
type LocalToolExecutor struct {
|
|
tools []MCPToolInfo
|
|
}
|
|
|
|
// NewLocalToolExecutor creates a ToolExecutor from local named sessions.
|
|
// It discovers tools immediately and caches the result.
|
|
func NewLocalToolExecutor(ctx context.Context, sessions []NamedSession) *LocalToolExecutor {
|
|
tools, err := DiscoverMCPTools(ctx, sessions)
|
|
if err != nil {
|
|
xlog.Error("Failed to discover MCP tools (local)", "error", err)
|
|
}
|
|
return &LocalToolExecutor{tools: tools}
|
|
}
|
|
|
|
func (e *LocalToolExecutor) DiscoverTools(_ context.Context) ([]functions.Function, error) {
|
|
var fns []functions.Function
|
|
for _, t := range e.tools {
|
|
fns = append(fns, t.Function)
|
|
}
|
|
return fns, nil
|
|
}
|
|
|
|
func (e *LocalToolExecutor) IsTool(name string) bool {
|
|
return IsMCPTool(e.tools, name)
|
|
}
|
|
|
|
func (e *LocalToolExecutor) ExecuteTool(ctx context.Context, toolName, arguments string) (string, error) {
|
|
return ExecuteMCPToolCall(ctx, e.tools, toolName, arguments)
|
|
}
|
|
|
|
func (e *LocalToolExecutor) HasTools() bool {
|
|
return len(e.tools) > 0
|
|
}
|
|
|
|
// DistributedToolExecutor routes tool operations through NATS to agent workers.
|
|
type DistributedToolExecutor struct {
|
|
natsClient MCPNATSClient
|
|
modelName string
|
|
remote config.MCPGenericConfig[config.MCPRemoteServers]
|
|
stdio config.MCPGenericConfig[config.MCPSTDIOServers]
|
|
toolDefs []mcpRemote.MCPToolDef
|
|
}
|
|
|
|
// NewDistributedToolExecutor creates a ToolExecutor that routes through NATS.
|
|
// It discovers tools immediately via a NATS request-reply to an agent worker.
|
|
func NewDistributedToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string,
|
|
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
|
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
|
) *DistributedToolExecutor {
|
|
e := &DistributedToolExecutor{
|
|
natsClient: natsClient,
|
|
modelName: modelName,
|
|
remote: remote,
|
|
stdio: stdio,
|
|
}
|
|
resp, err := DiscoverMCPToolsRemote(ctx, natsClient, modelName, remote, stdio)
|
|
if err != nil {
|
|
xlog.Error("Failed to discover MCP tools (distributed)", "error", err)
|
|
} else if resp != nil {
|
|
e.toolDefs = resp.Tools
|
|
}
|
|
return e
|
|
}
|
|
|
|
func (e *DistributedToolExecutor) DiscoverTools(_ context.Context) ([]functions.Function, error) {
|
|
var fns []functions.Function
|
|
for _, td := range e.toolDefs {
|
|
fns = append(fns, td.Function)
|
|
}
|
|
return fns, nil
|
|
}
|
|
|
|
func (e *DistributedToolExecutor) IsTool(name string) bool {
|
|
return slices.ContainsFunc(e.toolDefs, func(td mcpRemote.MCPToolDef) bool {
|
|
return td.ToolName == name
|
|
})
|
|
}
|
|
|
|
func (e *DistributedToolExecutor) ExecuteTool(ctx context.Context, toolName, arguments string) (string, error) {
|
|
return ExecuteMCPToolCallRemote(ctx, e.natsClient, e.modelName, e.remote, e.stdio, toolName, arguments)
|
|
}
|
|
|
|
func (e *DistributedToolExecutor) HasTools() bool {
|
|
return len(e.toolDefs) > 0
|
|
}
|
|
|
|
// NewToolExecutor creates the appropriate ToolExecutor based on the current mode.
|
|
// When natsClient is non-nil, returns a DistributedToolExecutor that routes through NATS.
|
|
// When natsClient is nil, creates local sessions and returns a LocalToolExecutor.
|
|
func NewToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string,
|
|
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
|
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
|
enabledServers []string,
|
|
) ToolExecutor {
|
|
if natsClient != nil {
|
|
return NewDistributedToolExecutor(ctx, natsClient, modelName, remote, stdio)
|
|
}
|
|
sessions, err := NamedSessionsFromMCPConfig(modelName, remote, stdio, enabledServers)
|
|
if err != nil || len(sessions) == 0 {
|
|
if err != nil {
|
|
xlog.Error("Failed to create MCP sessions", "error", err)
|
|
}
|
|
return &LocalToolExecutor{} // empty, HasTools() returns false
|
|
}
|
|
return NewLocalToolExecutor(ctx, sessions)
|
|
}
|