feat(tools): allow to automatically plan

fragment.go

@@ -219,18 +219,18 @@ func (f Fragment) LastMessage() *openai.ChatCompletionMessage {
 	return &f.Messages[len(f.Messages)-1]
 }
 
-func (f Fragment) LastAssistantMessages() []openai.ChatCompletionMessage {
+func (f Fragment) LastAssistantAndToolMessages() []openai.ChatCompletionMessage {
 
 	lastMessages := []openai.ChatCompletionMessage{}
 	found := false
 	for i := len(f.Messages) - 1; i >= 0; i-- {
 
-		if f.Messages[i].Role == "assistant" {
+		if f.Messages[i].Role == "assistant" || f.Messages[i].Role == "tool" {
 			found = true
 			lastMessages = append([]openai.ChatCompletionMessage{f.Messages[i]}, lastMessages...)
 		}
 
-		if found && f.Messages[i].Role != "assistant" {
+		if found && (f.Messages[i].Role != "assistant" && f.Messages[i].Role != "tool") {
 			break
 		}
 	}

fragment_e2e_test.go

@@ -79,11 +79,11 @@ var _ = Describe("Fragment test", func() {
 			fragment = fragment.AddMessage("assistant", "Byee!")
 
 			Expect(len(fragment.Messages)).To(Equal(4))
-			Expect(len(fragment.LastAssistantMessages())).To(Equal(2))
-			Expect(fragment.LastAssistantMessages()[0].Content).To(Equal("Hi!"))
-			Expect(fragment.LastAssistantMessages()[1].Content).To(Equal("Byee!"))
+			Expect(len(fragment.LastAssistantAndToolMessages())).To(Equal(2))
+			Expect(fragment.LastAssistantAndToolMessages()[0].Content).To(Equal("Hi!"))
+			Expect(fragment.LastAssistantAndToolMessages()[1].Content).To(Equal("Byee!"))
 			conv := NewEmptyFragment()
-			conv.Messages = append(conv.Messages, fragment.LastAssistantMessages()...)
+			conv.Messages = append(conv.Messages, fragment.LastAssistantAndToolMessages()...)
 
 			Expect(conv.Messages[0].Content).To(Equal("Hi!"))
 			Expect(conv.Messages[1].Content).To(Equal("Byee!"))

fragment_test.go

@@ -29,11 +29,11 @@ var _ = Describe("Fragment test", func() {
 			fragment = fragment.AddMessage("assistant", "Byee!")
 
 			Expect(len(fragment.Messages)).To(Equal(4))
-			Expect(len(fragment.LastAssistantMessages())).To(Equal(2))
-			Expect(fragment.LastAssistantMessages()[0].Content).To(Equal("Hi!"))
-			Expect(fragment.LastAssistantMessages()[1].Content).To(Equal("Byee!"))
+			Expect(len(fragment.LastAssistantAndToolMessages())).To(Equal(2))
+			Expect(fragment.LastAssistantAndToolMessages()[0].Content).To(Equal("Hi!"))
+			Expect(fragment.LastAssistantAndToolMessages()[1].Content).To(Equal("Byee!"))
 			conv := NewEmptyFragment()
-			conv.Messages = append(conv.Messages, fragment.LastAssistantMessages()...)
+			conv.Messages = append(conv.Messages, fragment.LastAssistantAndToolMessages()...)
 
 			Expect(conv.Messages[0].Content).To(Equal("Hi!"))
 			Expect(conv.Messages[1].Content).To(Equal("Byee!"))

guidelines.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/mudler/cogito/prompt"
 	"github.com/mudler/cogito/structures"
+	"github.com/sashabaranov/go-openai"
 )
 
 type Guidelines []Guideline
@@ -99,23 +100,31 @@ func GetRelevantGuidelines(llm LLM, guidelines Guidelines, fragment Fragment, op
 	return g, nil
 }
 
-func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines, error) {
+func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines, []openai.ChatCompletionMessage, error) {
 
 	o := defaultOptions()
 	o.Apply(opts...)
 
 	tools := slices.Clone(o.tools)
 
 	guidelines := o.guidelines
+	prompts := []openai.ChatCompletionMessage{}
 
 	for _, session := range o.mcpSessions {
 		mcpTools, err := mcpToolsFromTransport(o.context, session)
 		if err != nil {
-			return Tools{}, Guidelines{}, fmt.Errorf("failed to get MCP tools: %w", err)
+			return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get MCP tools: %w", err)
 		}
 		for _, tool := range mcpTools {
 			tools = append(tools, tool)
 		}
+		toolPrompts, err := mcpPromptsFromTransport(o.context, session, o.mcpArgs)
+		if err != nil {
+			return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get MCP prompts: %w", err)
+		}
+		if o.mcpPrompts {
+			prompts = append(prompts, toolPrompts...)
+		}
 	}
 
 	if len(o.guidelines) > 0 {
@@ -125,12 +134,12 @@ func usableTools(llm LLM, fragment Fragment, opts ...Option) (Tools, Guidelines,
 		var err error
 		guidelines, err = GetRelevantGuidelines(llm, o.guidelines, fragment, opts...)
 		if err != nil {
-			return Tools{}, Guidelines{}, fmt.Errorf("failed to get relevant guidelines: %w", err)
+			return Tools{}, Guidelines{}, nil, fmt.Errorf("failed to get relevant guidelines: %w", err)
 		}
 		for _, guideline := range guidelines {
 			tools = append(tools, guideline.Tools...)
 		}
 	}
 
-	return tools, guidelines, nil
+	return tools, guidelines, prompts, nil
 }

mcp.go

@@ -69,6 +69,33 @@ type toolInputSchema struct {
 	Required   []string               `json:"required,omitempty"`
 }
 
+func mcpPromptsFromTransport(ctx context.Context, session *mcp.ClientSession, arguments map[string]string) ([]openai.ChatCompletionMessage, error) {
+	prompts, err := session.ListPrompts(ctx, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	promptsList := []openai.ChatCompletionMessage{}
+
+	for _, prompt := range prompts.Prompts {
+		p, err := session.GetPrompt(ctx, &mcp.GetPromptParams{Name: prompt.Name, Arguments: arguments})
+		if err != nil {
+			return nil, err
+		}
+		for _, message := range p.Messages {
+			switch message.Content.(type) {
+			case *mcp.TextContent:
+				promptsList = append(promptsList, openai.ChatCompletionMessage{
+					Role:    string(message.Role),
+					Content: message.Content.(*mcp.TextContent).Text,
+				})
+			}
+		}
+	}
+
+	return promptsList, nil
+}
+
 // probe the MCP remote and generate tools that are compliant with cogito
 func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession) ([]*mcpTool, error) {
 	allTools := []*mcpTool{}

options.go

@@ -16,6 +16,8 @@ type Options struct {
 	deepContext            bool
 	toolReasoner           bool
 	toolReEvaluator        bool
+	autoPlan               bool
+	planReEvaluator        bool
 	statusCallback         func(string)
 	gaps                   []string
 	context                context.Context
@@ -27,6 +29,8 @@ type Options struct {
 	strictGuidelines       bool
 	mcpSessions            []*mcp.ClientSession
 	guidelines             Guidelines
+	mcpPrompts             bool
+	mcpArgs                map[string]string
 }
 
 type Option func(*Options)
@@ -75,6 +79,21 @@ var (
 	EnableStrictGuidelines Option = func(o *Options) {
 		o.strictGuidelines = true
 	}
+
+	// EnableAutoPlan enables cogito to automatically use planning if needed
+	EnableAutoPlan Option = func(o *Options) {
+		o.autoPlan = true
+	}
+
+	// EnableAutoPlanReEvaluator enables cogito to automatically re-evaluate the need to use planning
+	EnableAutoPlanReEvaluator Option = func(o *Options) {
+		o.planReEvaluator = true
+	}
+
+	// EnableMCPPrompts enables the use of MCP prompts
+	EnableMCPPrompts Option = func(o *Options) {
+		o.mcpPrompts = true
+	}
 )
 
 // WithIterations allows to set the number of refinement iterations
@@ -166,3 +185,10 @@ func WithMCPs(sessions ...*mcp.ClientSession) func(o *Options) {
 		o.mcpSessions = append(o.mcpSessions, sessions...)
 	}
 }
+
+// WithMCPArgs sets the arguments for the MCP prompts
+func WithMCPArgs(args map[string]string) func(o *Options) {
+	return func(o *Options) {
+		o.mcpArgs = args
+	}
+}

plan.go

@@ -167,7 +167,7 @@ func ExecutePlan(llm LLM, conv Fragment, plan *structures.Plan, goal *structures
 			return Fragment{}, nil
 		}
 
-		conv.Messages = append(conv.Messages, subtaskConvResult.LastAssistantMessages()...)
+		conv.Messages = append(conv.Messages, subtaskConvResult.LastAssistantAndToolMessages()...)
 		conv.Status.Iterations = conv.Status.Iterations + 1
 		conv.Status.ToolsCalled = append(conv.Status.ToolsCalled, subtaskConvResult.Status.ToolsCalled...)
 		conv.Status.ToolResults = append(conv.Status.ToolResults, subtaskConvResult.Status.ToolResults...)

plan_test.go

@@ -145,12 +145,17 @@ var _ = Describe("Plannings with tools", func() {
 				ContainSubstring("What is photosynthesis?"),
 			)
 
-			Expect(len(result.Messages)).To(Equal(3))
+			Expect(len(result.Status.ToolResults)).To(Equal(2))
+			Expect(len(result.Status.ToolsCalled)).To(Equal(2))
+			Expect(len(result.Messages)).To(Equal(5))
+
 			Expect(result.Messages[0].Content).To(Equal("What is photosynthesis?"))
 			Expect(result.Messages[1].ToolCalls[0].Function.Arguments).To(Equal("{\"query\": \"chlorophyll\"}"))
 			Expect(result.Messages[1].ToolCalls[0].Function.Name).To(Equal("search"))
-			Expect(result.Messages[2].ToolCalls[0].Function.Arguments).To(Equal("{\"query\": \"photosynthesis\"}"))
-			Expect(result.Messages[2].ToolCalls[0].Function.Name).To(Equal("search"))
+			Expect(result.Messages[2].Content).To(Equal("Chlorophyll is a green pigment found in plants."))
+			Expect(result.Messages[3].ToolCalls[0].Function.Arguments).To(Equal("{\"query\": \"photosynthesis\"}"))
+			Expect(result.Messages[3].ToolCalls[0].Function.Name).To(Equal("search"))
+			Expect(result.Messages[4].Content).To(Equal("Photosynthesis is the process by which plants convert sunlight into energy."))
 		})
 	})
 })

prompt/prompt.go

@@ -16,6 +16,7 @@ const (
 	PromptPlanExecutionType        PromptType = iota
 	PromptGuidelinesType           PromptType = iota
 	PromptGuidelinesExtractionType PromptType = iota
+	PromptPlanDecisionType         PromptType = iota
 )
 
 var (
@@ -33,6 +34,7 @@ var (
 		PromptPlanExecutionType:        PromptPlanExecution,
 		PromptGuidelinesType:           PromptGuidelines,
 		PromptGuidelinesExtractionType: PromptGuidelinesExtraction,
+		PromptPlanDecisionType:         DecideIfPlanningIsNeeded,
 	}
 
 	PromptGuidelinesExtraction = NewPrompt("What guidelines should be applied? return only the numbers of the guidelines by using the json tool with a list of integers corresponding to the guidelines.")
@@ -281,4 +283,25 @@ Context:
 You will use the "json" tool with the option "extract_boolean" set to either yes or no.
 Reply with the appropriate boolean extraction tool with yes or no, based on the context. 
 If the context speaks about, let's say doing something, you will replay with yes, or a no otherwise.`)
+
+	DecideIfPlanningIsNeeded = NewPrompt(`You are an AI assistant that decides if planning and executing subtasks in sequence is needed from a conversation.
+
+Conversation: 
+{{.Context}}
+
+{{if ne .AdditionalContext ""}}
+AdditionalContext:
+{{.AdditionalContext}}
+{{end}}
+
+Available tools:
+{{ range $index, $tool := .Tools }}
+- Tool name: "{{$tool.Name}}" 
+  Tool description: {{$tool.Description}}
+  Tool arguments: {{$tool.Parameters | toJson}}
+{{ end }}
+
+Based on the conversation, context, and available tools, decide if planning and executing subtasks in sequence is needed.
+Keep in mind that Planning will later involve in breaking down the problem into a set of subtasks that require running tools in sequence and evaluating their results.
+If you think planning is needed, reply with yes, otherwise reply with no.`)
 )