mirror/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2026-01-06 10:39:55 -06:00
Code Issues Packages Projects Releases Wiki Activity

chore: update cogito and simplify MCP logics (#6413)

Browse Source 
* chore: update cogito and simplify MCP logics

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Refine signal handling

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto
2025-10-09 12:36:45 +02:00
committed by GitHub
parent 459b6ab86d
commit 27c4161401
13 changed files with 144 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
core/http/endpoints/mcp/tools.go
View File

@@ -2,43 +2,66 @@ package mcp
import (
"context"
"encoding/json"
"errors"
"net/http"
"os"
"os/exec"
"os/signal"
"syscall"
"sync"
"time"
"github.com/mudler/LocalAI/core/config"
"github.com/sashabaranov/go-openai"
"github.com/tmc/langchaingo/jsonschema"
"github.com/mudler/LocalAI/pkg/signals"
"github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/rs/zerolog/log"
)
func ToolsFromMCPConfig(ctx context.Context, remote config.MCPGenericConfig[config.MCPRemoteServers], stdio config.MCPGenericConfig[config.MCPSTDIOServers]) ([]*MCPTool, error) {
allTools := []*MCPTool{}
type sessionCache struct {
mu sync.Mutex
cache map[string][]*mcp.ClientSession
}
var (
cache = sessionCache{
cache: make(map[string][]*mcp.ClientSession),
}
client = mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil)
)
func SessionsFromMCPConfig(
name string,
remote config.MCPGenericConfig[config.MCPRemoteServers],
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
) ([]*mcp.ClientSession, error) {
cache.mu.Lock()
defer cache.mu.Unlock()
sessions, exists := cache.cache[name]
if exists {
return sessions, nil
}
allSessions := []*mcp.ClientSession{}
ctx, cancel := context.WithCancel(context.Background())
// Get the list of all the tools that the Agent will be esposed to
for _, server := range remote.Servers {
log.Debug().Msgf("[MCP remote server] Configuration : %+v", server)
// Create HTTP client with custom roundtripper for bearer token injection
client := &http.Client{
httpClient := &http.Client{
Timeout: 360 * time.Second,
Transport: newBearerTokenRoundTripper(server.Token, http.DefaultTransport),
}
tools, err := mcpToolsFromTransport(ctx,
&mcp.StreamableClientTransport{Endpoint: server.URL, HTTPClient: client},
)
transport := &mcp.StreamableClientTransport{Endpoint: server.URL, HTTPClient: httpClient}
mcpSession, err := client.Connect(ctx, transport, nil)
if err != nil {
return nil, err
log.Error().Err(err).Msgf("Failed to connect to MCP server %s", server.URL)
continue
}
allTools = append(allTools, tools...)
log.Debug().Msgf("[MCP remote server] Connected to MCP server %s", server.URL)
cache.cache[name] = append(cache.cache[name], mcpSession)
}
for _, server := range stdio.Servers {
@@ -48,18 +71,24 @@ func ToolsFromMCPConfig(ctx context.Context, remote config.MCPGenericConfig[conf
for key, value := range server.Env {
command.Env = append(command.Env, key+"="+value)
}
tools, err := mcpToolsFromTransport(ctx,
&mcp.CommandTransport{
Command: command},
)
transport := &mcp.CommandTransport{Command: command}
mcpSession, err := client.Connect(ctx, transport, nil)
if err != nil {
return nil, err
log.Error().Err(err).Msgf("Failed to start MCP server %s", command)
continue
}
allTools = append(allTools, tools...)
log.Debug().Msgf("[MCP stdio server] Connected to MCP server %s", command)
cache.cache[name] = append(cache.cache[name], mcpSession)
}
return allTools, nil
signals.RegisterGracefulTerminationHandler(func() {
for _, session := range allSessions {
session.Close()
}
cancel()
})
return allSessions, nil
}
// bearerTokenRoundTripper is a custom roundtripper that injects a bearer token
@@ -87,146 +116,3 @@ func newBearerTokenRoundTripper(token string, base http.RoundTripper) http.Round
base: base,
}
}
type MCPTool struct {
name, description string
inputSchema ToolInputSchema
session *mcp.ClientSession
ctx context.Context
props map[string]jsonschema.Definition
}
func (t *MCPTool) Run(args map[string]any) (string, error) {
// Call a tool on the server.
params := &mcp.CallToolParams{
Name: t.name,
Arguments: args,
}
res, err := t.session.CallTool(t.ctx, params)
if err != nil {
log.Error().Msgf("CallTool failed: %v", err)
return "", err
}
if res.IsError {
log.Error().Msgf("tool failed")
return "", errors.New("tool failed")
}
result := ""
for _, c := range res.Content {
result += c.(*mcp.TextContent).Text
}
return result, nil
}
func (t *MCPTool) Tool() openai.Tool {
return openai.Tool{
Type: openai.ToolTypeFunction,
Function: &openai.FunctionDefinition{
Name: t.name,
Description: t.description,
Parameters: jsonschema.Definition{
Type: jsonschema.Object,
Properties: t.props,
Required: t.inputSchema.Required,
},
},
}
}
func (t *MCPTool) Close() {
t.session.Close()
}
type ToolInputSchema struct {
Type string `json:"type"`
Properties map[string]interface{} `json:"properties,omitempty"`
Required []string `json:"required,omitempty"`
}
// probe the MCP remote and generate tools that are compliant with cogito
// TODO: Maybe move this to cogito?
func mcpToolsFromTransport(ctx context.Context, transport mcp.Transport) ([]*MCPTool, error) {
allTools := []*MCPTool{}
// Create a new client, with no features.
client := mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil)
session, err := client.Connect(ctx, transport, nil)
if err != nil {
log.Error().Msgf("Error connecting to MCP server: %v", err)
return nil, err
}
tools, err := session.ListTools(ctx, nil)
if err != nil {
log.Error().Msgf("Error listing tools: %v", err)
return nil, err
}
for _, tool := range tools.Tools {
dat, err := json.Marshal(tool.InputSchema)
if err != nil {
log.Error().Msgf("Error marshalling input schema: %v", err)
continue
}
// XXX: This is a wild guess, to verify (data types might be incompatible)
var inputSchema ToolInputSchema
err = json.Unmarshal(dat, &inputSchema)
if err != nil {
log.Error().Msgf("Error unmarshalling input schema: %v", err)
continue
}
props := map[string]jsonschema.Definition{}
dat, err = json.Marshal(inputSchema.Properties)
if err != nil {
log.Error().Msgf("Error marshalling input schema: %v", err)
continue
}
err = json.Unmarshal(dat, &props)
if err != nil {
log.Error().Msgf("Error unmarshalling input schema properties: %v", err)
continue
}
allTools = append(allTools, &MCPTool{
name: tool.Name,
description: tool.Description,
session: session,
ctx: ctx,
props: props,
inputSchema: inputSchema,
})
}
// We make sure we run Close on signal
handleSignal(allTools)
return allTools, nil
}
func handleSignal(tools []*MCPTool) {
// Create a channel to receive OS signals
sigChan := make(chan os.Signal, 1)
// Register for interrupt and terminate signals
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
// Handle signals in a separate goroutine
go func() {
sig := <-sigChan
log.Printf("Received signal %v, shutting down gracefully...", sig)
for _, t := range tools {
t.Close()
}
// Exit the application
os.Exit(0)
}()
}

39
core/http/endpoints/openai/mcp.go
View File

@@ -5,11 +5,10 @@ import (
"errors"
"fmt"
"strings"
"sync"
"time"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/http/endpoints/mcp"
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
"github.com/mudler/LocalAI/core/http/middleware"
"github.com/gofiber/fiber/v2"
@@ -27,10 +26,6 @@ import (
// @Success 200 {object} schema.OpenAIResponse "Response"
// @Router /mcp/v1/completions [post]
func MCPCompletionEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig) func(c *fiber.Ctx) error {
toolsCache := map[string][]*mcp.MCPTool{}
mu := sync.Mutex{}
// We do not support streaming mode (Yet?)
return func(c *fiber.Ctx) error {
created := int(time.Now().Unix())
@@ -54,37 +49,17 @@ func MCPCompletionEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader,
return fmt.Errorf("no MCP servers configured")
}
allTools := []*mcp.MCPTool{}
// Get MCP config from model config
remote, stdio := config.MCP.MCPConfigFromYAML()
// Check if we have tools in cache, or we have to have an initial connection
mu.Lock()
tools, exists := toolsCache[config.Name]
if exists {
allTools = append(allTools, tools...)
} else {
tools, err := mcp.ToolsFromMCPConfig(ctx, remote, stdio)
if err != nil {
mu.Unlock()
return err
}
toolsCache[config.Name] = tools
allTools = append(allTools, tools...)
}
mu.Unlock()
cogitoTools := []cogito.Tool{}
for _, tool := range allTools {
cogitoTools = append(cogitoTools, tool)
// defer tool.Close()
sessions, err := mcpTools.SessionsFromMCPConfig(config.Name, remote, stdio)
if err != nil {
return err
}
if len(cogitoTools) == 0 {
return fmt.Errorf("no tools found in the specified MCP servers")
if len(sessions) == 0 {
return fmt.Errorf("no working MCP servers found")
}
fragment := cogito.NewEmptyFragment()
@@ -109,7 +84,7 @@ func MCPCompletionEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader,
log.Debug().Msgf("[model agent] [model: %s] Status: %s", config.Name, s)
}),
cogito.WithContext(ctx),
cogito.WithTools(cogitoTools...),
cogito.WithMCPs(sessions...),
cogito.WithIterations(3), // default to 3 iterations
cogito.WithMaxAttempts(3), // default to 3 attempts
}