diff --git a/core/agent/agent.go b/core/agent/agent.go index e492becb..a15d798e 100644 --- a/core/agent/agent.go +++ b/core/agent/agent.go @@ -297,7 +297,7 @@ func (a *Agent) Stop() { a.Lock() defer a.Unlock() xlog.Debug("Stopping agent", "agent", a.Character.Name) - a.closeMCPSTDIOServers() + a.closeMCPServers() a.context.Cancel() } diff --git a/core/agent/mcp.go b/core/agent/mcp.go index 62d85b37..7c505bff 100644 --- a/core/agent/mcp.go +++ b/core/agent/mcp.go @@ -2,6 +2,7 @@ package agent import ( "context" + "encoding/base64" "encoding/json" "errors" "os" @@ -61,11 +62,19 @@ func (m *mcpAction) Run(ctx context.Context, sharedState *types.AgentSharedState } result := "" + imageBase64Result := "" for _, c := range resp.Content { - result += c.(*mcp.TextContent).Text + switch content := c.(type) { + case *mcp.TextContent: + result += content.Text + case *mcp.ImageContent: + imageBase64Result = "data:" + content.MIMEType + ";base64," + base64.StdEncoding.EncodeToString(content.Data) + default: + log.Error().Msgf("[Unknown content type received: %T]", content) + } } - return types.ActionResult{Result: result}, nil + return types.ActionResult{Result: result, ImageBase64Result: imageBase64Result}, nil } func (m *mcpAction) Definition() types.ActionDefinition { @@ -170,7 +179,7 @@ func (a *Agent) initMCPActions() error { generatedActions := types.Actions{} client := mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil) - // Connect to a server over stdin/stdout. + a.closeMCPServers() // Make sure we stop all previous servers if any is active // MCP HTTP Servers for _, mcpServer := range a.options.mcpServers { @@ -200,8 +209,6 @@ func (a *Agent) initMCPActions() error { // MCP STDIO Servers - a.closeMCPSTDIOServers() // Make sure we stop all previous servers if any is active - if a.options.mcpPrepareScript != "" { xlog.Debug("Preparing MCP", "script", a.options.mcpPrepareScript) @@ -240,7 +247,7 @@ func (a *Agent) initMCPActions() error { return err } -func (a *Agent) closeMCPSTDIOServers() { +func (a *Agent) closeMCPServers() { for _, s := range a.mcpSessions { s.Close() }