package cogito_test import ( "fmt" "strings" . "github.com/mudler/cogito" "github.com/mudler/cogito/tests/mock" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/sashabaranov/go-openai" "github.com/sashabaranov/go-openai/jsonschema" ) var _ = Describe("ExecuteTools", func() { var mockLLM *mock.MockOpenAIClient var originalFragment Fragment BeforeEach(func() { mockLLM = mock.NewMockOpenAIClient() originalFragment = NewEmptyFragment(). AddMessage(UserMessageRole, "What is photosynthesis?"). AddMessage(AssistantMessageRole, "Photosynthesis is the process by which plants convert sunlight into energy.") }) Context("ToolDefinition", func() { It("should create a valid ToolDefinition", func() { mockToolDef := mock.NewMockTool("search", "Search for information") mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) toolDefinition := ToolDefinition[map[string]any]{ ToolRunner: mockToolDefT.ToolRunner, Name: "search", Description: "Search for information", InputArguments: &struct { Query string `json:"query"` }{}, } tool := toolDefinition.Tool() Expect(tool.Function.Name).To(Equal("search")) Expect(tool.Function.Description).To(Equal("Search for information")) Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ Type: jsonschema.Object, AdditionalProperties: false, Properties: map[string]jsonschema.Definition{ "query": { Type: jsonschema.String, Enum: nil, }, }, Required: []string{"query"}, Defs: map[string]jsonschema.Definition{}, })) }) It("should create a valid ToolDefinition with enums and description", func() { mockToolDef := mock.NewMockTool("search", "Search for information") mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) toolDefinition := ToolDefinition[map[string]any]{ ToolRunner: mockToolDefT.ToolRunner, Name: "search", Description: "Search for information", InputArguments: &struct { Query string `json:"query" enum:"foo,bar" description:"The query to search for"` }{}, } tool := toolDefinition.Tool() Expect(tool.Function.Name).To(Equal("search")) Expect(tool.Function.Description).To(Equal("Search for information")) Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ Type: jsonschema.Object, AdditionalProperties: false, Properties: map[string]jsonschema.Definition{ "query": { Type: jsonschema.String, Enum: []string{"foo", "bar"}, Description: "The query to search for", }, }, Required: []string{"query"}, Defs: map[string]jsonschema.Definition{}, })) }) It("should create a valid ToolDefinition which arg is not required", func() { mockToolDef := mock.NewMockTool("search", "Search for information") mockToolDefT := mockToolDef.(*ToolDefinition[map[string]any]) toolDefinition := ToolDefinition[map[string]any]{ ToolRunner: mockToolDefT.ToolRunner, Name: "search", Description: "Search for information", InputArguments: &struct { Query string `json:"query" required:"false"` }{}, } tool := toolDefinition.Tool() Expect(tool.Function.Name).To(Equal("search")) Expect(tool.Function.Description).To(Equal("Search for information")) Expect(tool.Function.Parameters).To(Equal(jsonschema.Definition{ Type: jsonschema.Object, AdditionalProperties: false, Properties: map[string]jsonschema.Definition{ "query": { Type: jsonschema.String, }, }, Required: nil, Defs: map[string]jsonschema.Definition{}, })) }) }) Context("ExecuteTools with tools", func() { It("should execute tools when provided", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection and execution mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // After tool execution, ToolReEvaluator (toolSelection) picks next tool mockLLM.AddCreateChatCompletionFunction("search", `{"query": "grass"}`) // Second tool selection and execution // (The "grass" tool call above will be picked as nextAction) mock.SetRunResult(mockTool, "Grass is a plant that grows on the ground.") // After tool execution, ToolReEvaluator (toolSelection) picks next tool mockLLM.AddCreateChatCompletionFunction("search", `{"query": "baz"}`) // Third tool selection and execution // (The "baz" tool call above will be picked as nextAction) mock.SetRunResult(mockTool, "Baz is a plant that grows on the ground.") // After ToolReEvaluator returns no tool, Ask() is called to get final response mockLLM.SetAskResponse("Here is the final response with all the information gathered.") result, err := ExecuteTools(mockLLM, originalFragment, WithIterations(3), WithTools(mockTool)) Expect(err).ToNot(HaveOccurred()) // ExecuteTools now calls Ask() at the end to get a final response // when ToolReEvaluator returns no more tools Expect(len(mockLLM.FragmentHistory)).To(Equal(1), fmt.Sprintf("Fragment history: %v", mockLLM.FragmentHistory)) Expect(result).ToNot(BeNil()) Expect(len(result.Status.ToolsCalled)).To(Equal(3)) Expect(len(result.Status.ToolResults)).To(Equal(3)) Expect(result.Status.ToolResults[0].Executed).To(BeTrue()) Expect(result.Status.ToolResults[0].Name).To(Equal("search")) Expect(result.Status.ToolResults[0].Result).To(Equal("Chlorophyll is a green pigment found in plants.")) Expect(result.Status.ToolResults[1].Executed).To(BeTrue()) Expect(result.Status.ToolResults[1].Name).To(Equal("search")) Expect(result.Status.ToolResults[1].Result).To(Equal("Grass is a plant that grows on the ground.")) Expect(result.Status.ToolResults[2].Executed).To(BeTrue()) Expect(result.Status.ToolResults[2].Name).To(Equal("search")) Expect(result.Status.ToolResults[2].Result).To(Equal("Baz is a plant that grows on the ground.")) }) It("should execute tools when provided with guidelines", func() { mockTool := mock.NewMockTool("search", "Search for information") mockWeatherTool := mock.NewMockTool("get_weather", "Get the weather") // First iteration // 1. Guidelines selection: mockLLM.SetAskResponse("Only the first guideline is relevant.") mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [1]}`) // 2. Tool selection (direct): mockLLM.AddCreateChatCompletionFunction("search", `{"query": "chlorophyll"}`) mock.SetRunResult(mockTool, "Chlorophyll is a green pigment found in plants.") // Second iteration // 1. Guidelines selection: mockLLM.SetAskResponse("Only the first guideline is relevant.") mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [1]}`) // 2. Tool selection: mockLLM.AddCreateChatCompletionFunction("search", `{"query": "grass"}`) mock.SetRunResult(mockTool, "Grass is a plant that grows on the ground.") // Third iteration // 1. Guidelines selection: mockLLM.SetAskResponse("Only the second guideline is relevant.") mockLLM.AddCreateChatCompletionFunction("json", `{"guidelines": [2]}`) // 2. Tool selection: mockLLM.AddCreateChatCompletionFunction("get_weather", `{"query": "baz"}`) mock.SetRunResult(mockWeatherTool, "Baz is a plant that grows on the ground.") // When max iterations is reached, Ask() is called to get final response mockLLM.SetAskResponse("All tasks completed.") result, err := ExecuteTools(mockLLM, originalFragment, WithIterations(3), WithTools(mockTool, mockWeatherTool), EnableStrictGuidelines, WithGuidelines( Guideline{ Condition: "User asks about informations", Action: "Use the search tool to find information.", Tools: Tools{mockTool}, }, Guideline{ Condition: "User asks for the weather in a city", Action: "Use the weather tool to find the weather in the city.", Tools: Tools{mockWeatherTool}, }, )) Expect(err).ToNot(HaveOccurred()) // Check fragments history to see if we behaved as expected // Guidelines selection: 3 iterations × 1 Ask() + Final response when max iterations reached: 1 Ask() = 4 Ask() calls Expect(len(mockLLM.FragmentHistory)).To(Equal(4), fmt.Sprintf("Fragment history: %v", mockLLM.FragmentHistory)) // Iteration 1: [0] Guidelines Expect(mockLLM.FragmentHistory[0].String()).To(ContainSubstring("You are an AI assistant that needs to understand if any of the guidelines should be applied")) // Iteration 2: [1] Guidelines Expect(mockLLM.FragmentHistory[1].String()).To(ContainSubstring("You are an AI assistant that needs to understand if any of the guidelines should be applied")) // Iteration 3: [2] Guidelines Expect(mockLLM.FragmentHistory[2].String()).To(ContainSubstring("You are an AI assistant that needs to understand if any of the guidelines should be applied")) // [3] Final Ask when max iterations reached - check that it contains the conversation Expect(mockLLM.FragmentHistory[3].String()).To(ContainSubstring(`get_weather({"query":"baz"})`)) Expect(result).ToNot(BeNil()) Expect(len(result.Status.ToolsCalled)).To(Equal(3)) Expect(len(result.Status.ToolResults)).To(Equal(3)) Expect(result.Status.ToolResults[0].Executed).To(BeTrue()) Expect(result.Status.ToolResults[0].Name).To(Equal("search")) Expect(result.Status.ToolResults[0].Result).To(Equal("Chlorophyll is a green pigment found in plants.")) Expect(result.Status.ToolResults[1].Executed).To(BeTrue()) Expect(result.Status.ToolResults[1].Name).To(Equal("search")) Expect(result.Status.ToolResults[1].Result).To(Equal("Grass is a plant that grows on the ground.")) Expect(result.Status.ToolResults[2].Executed).To(BeTrue()) Expect(result.Status.ToolResults[2].Name).To(Equal("get_weather")) Expect(result.Status.ToolResults[2].Result).To(Equal("Baz is a plant that grows on the ground.")) }) It("should execute autoplan basic functionality", func() { mockTool := mock.NewMockTool("search", "Search for information") // Mock planning decision - decide that planning is needed mockLLM.SetAskResponse("Yes, this task requires planning to be completed effectively.") mockLLM.AddCreateChatCompletionFunction("json", `{"extract_boolean": true}`) // Mock goal extraction mockLLM.SetAskResponse("The goal is to research information about photosynthesis.") mockLLM.AddCreateChatCompletionFunction("json", `{"goal": "Research information about photosynthesis"}`) // Mock plan creation (first step of plan extraction) mockLLM.SetAskResponse("Here is a plan with subtasks: 1. Search for basic information about photosynthesis") // Mock subtask extraction (second step of plan extraction) - this uses CreateChatCompletion mockLLM.AddCreateChatCompletionFunction("json", `{"subtasks": ["Search for basic information about photosynthesis"]}`) // Mock first subtask execution - search mockLLM.AddCreateChatCompletionFunction("search", `{"query": "photosynthesis basics"}`) mockLLM.SetAskResponse("Photosynthesis is the process by which plants convert sunlight into energy.") mock.SetRunResult(mockTool, "Photosynthesis is the process by which plants convert sunlight into energy.") mockLLM.SetAskResponse("Photosynthesis is the process by which plants convert sunlight into energy.") // Mock goal achievement check for first subtask mockLLM.AddCreateChatCompletionFunction("json", `{"extract_boolean": true}`) result, err := ExecuteTools(mockLLM, originalFragment, EnableAutoPlan, WithTools(mockTool)) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) // Verify that planning was executed by checking fragment history // PlanDecision + GoalExtraction + PlanCreation + GoalCheck = 4 Ask() calls // ToolReEvaluator uses toolSelection (CreateChatCompletion), not Ask() Expect(len(mockLLM.FragmentHistory)).To(BeNumerically("==", 5), fmt.Sprintf("Fragment history: %v", mockLLM.FragmentHistory)) // Check that planning decision was made Expect(mockLLM.FragmentHistory[0].String()).To( And( ContainSubstring("You are an AI assistant that decides if planning and executing subtasks in sequence is needed from a conversation"), ContainSubstring("What is photosynthesis"), )) // Check that goal extraction was called Expect(mockLLM.FragmentHistory[1].String()).To( ContainSubstring("Analyze the following text and the context to identify the goal")) // Check that plan creation was called Expect(mockLLM.FragmentHistory[2].String()).To( ContainSubstring("You are an AI assistant that breaks down a goal into a series of actionable steps")) Expect(len(result.Messages)).To(Equal(4), fmt.Sprintf("Messages: %+v", result.Messages)) Expect(result.Messages[len(result.Messages)-1].Content).To( And( ContainSubstring("hotosynthesis is the process by which plants convert sunlight into energy"), ), fmt.Sprintf("Result: %+v", result), ) Expect(len(result.Status.Plans)).To(Equal(1)) // Verify tools were called correctly Expect(len(result.Status.ToolsCalled)).To(Equal(1)) Expect(len(result.Status.ToolResults)).To(Equal(1)) Expect(result.Status.ToolResults[0].Executed).To(BeTrue()) Expect(result.Status.ToolResults[0].Name).To(Equal("search")) Expect(result.Status.ToolResults[0].Result).To(Equal("Photosynthesis is the process by which plants convert sunlight into energy.")) }) It("should not execute autoplan when planning is not needed", func() { mockTool := mock.NewMockTool("search", "Search for information") // Mock planning decision - decide that planning is NOT needed mockLLM.SetAskResponse("No, this task does not require planning.") mockLLM.AddCreateChatCompletionFunction("json", `{"extract_boolean": false}`) // Mock regular tool execution (since planning is not needed, it falls back to normal tool execution) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "photosynthesis"}`) mockLLM.SetAskResponse("Photosynthesis is the process by which plants convert sunlight into energy.") mock.SetRunResult(mockTool, "Photosynthesis is the process by which plants convert sunlight into energy.") // After tool execution, ToolReEvaluator (toolSelection) returns no tool (text response) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, EnableAutoPlan, WithTools(mockTool)) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) // Verify that planning decision was made but no plan was executed // PlanDecision = 1 Ask() call // ToolReEvaluator uses toolSelection (CreateChatCompletion), not Ask() Expect(len(mockLLM.FragmentHistory)).To(Equal(2), fmt.Sprintf("Fragment history: %v", mockLLM.FragmentHistory)) // Check that planning decision was made Expect(mockLLM.FragmentHistory[0].String()).To( And( ContainSubstring("You are an AI assistant that decides if planning and executing subtasks in sequence is needed from a conversation"), ContainSubstring("What is photosynthesis"), )) // Check that tools were called (regular tool execution, not planning) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) Expect(len(result.Status.ToolResults)).To(Equal(1)) Expect(result.Status.ToolResults[0].Executed).To(BeTrue()) Expect(result.Status.ToolResults[0].Name).To(Equal("search")) Expect(result.Status.ToolResults[0].Result).To(Equal("Photosynthesis is the process by which plants convert sunlight into energy.")) }) }) Context("Tool Call Callbacks", func() { It("should call the callback with ToolChoice and SessionState", func() { mockTool := mock.NewMockTool("search", "Search for information") var receivedTool *ToolChoice var receivedState *SessionState // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Test result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { receivedTool = tool receivedState = state return ToolCallDecision{Approved: true} })) Expect(err).ToNot(HaveOccurred()) Expect(receivedTool).ToNot(BeNil()) Expect(receivedTool.Name).To(Equal("search")) Expect(receivedState).ToNot(BeNil()) Expect(receivedState.ToolChoice).To(Equal(receivedTool)) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) }) It("should interrupt execution when Approved is false", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { return ToolCallDecision{Approved: false} })) Expect(err).To(HaveOccurred()) Expect(err).To(Equal(ErrToolCallCallbackInterrupted)) Expect(len(result.Status.ToolsCalled)).To(Equal(0)) }) It("should skip tool call when Skip is true", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) // After skipping, ToolReEvaluator returns no tool (this happens after the skip) mockLLM.SetAskResponse("LLM result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { return ToolCallDecision{Approved: true, Skip: true} })) // When skipping with DisableToolReEvaluator, we might get ErrNoToolSelected // because no tools were actually executed. This is expected behavior. if err != nil { Expect(err).To(Equal(ErrNoToolSelected)) } // Tool should not be executed Expect(len(result.Status.ToolsCalled)).To(Equal(0)) // But should be in the conversation Expect(len(result.Messages)).To(BeNumerically(">", 0)) // Check that skip message was added foundSkipMessage := false for _, msg := range result.Messages { if msg.Role == "tool" && msg.Content == "Tool call skipped by user" { foundSkipMessage = true break } } Expect(foundSkipMessage).To(BeTrue()) }) It("should use directly modified tool choice when Modified is set", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection (will be modified) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) mock.SetRunResult(mockTool, "Modified result") mockLLM.SetAskResponse("LLM result") var executedArgs map[string]any result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { // Directly modify the tool arguments modified := *tool modified.Arguments = map[string]any{ "query": "modified_query", } return ToolCallDecision{ Approved: true, Modified: &modified, } })) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) Expect(result.LastMessage().Content).To(Equal("LLM result")) // Check that the modified arguments were used executedArgs = result.Status.ToolResults[0].ToolArguments.Arguments Expect(executedArgs["query"]).To(Equal("modified_query")) }) It("should handle adjustment feedback and re-evaluate tool call", func() { mockTool := mock.NewMockTool("search", "Search for information") callbackCount := 0 // First tool selection (original) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) // Adjustment: LLM re-evaluates with feedback mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted"}`) mock.SetRunResult(mockTool, "Adjusted result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithMaxAdjustmentAttempts(3), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { callbackCount++ if callbackCount == 1 { // First call: provide adjustment feedback return ToolCallDecision{ Approved: true, Adjustment: "Please use a more specific query", } } // Second call: approve the adjusted tool return ToolCallDecision{Approved: true} })) Expect(err).ToNot(HaveOccurred()) Expect(callbackCount).To(Equal(2)) // Called twice: original + adjusted Expect(len(result.Status.ToolsCalled)).To(Equal(1)) // Check that adjusted arguments were used Expect(result.Status.ToolResults[0].ToolArguments.Arguments["query"]).To(Equal("adjusted")) }) It("should respect max adjustment attempts limit", func() { mockTool := mock.NewMockTool("search", "Search for information") callbackCount := 0 // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) // Adjustment attempts (will hit max) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted1"}`) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted2"}`) mock.SetRunResult(mockTool, "Final result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithMaxAdjustmentAttempts(2), // Limit to 2 attempts WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { callbackCount++ // Always provide adjustment to test max limit if callbackCount < 3 { return ToolCallDecision{ Approved: true, Adjustment: "Keep adjusting", } } return ToolCallDecision{Approved: true} })) Expect(err).ToNot(HaveOccurred()) // Should have hit max attempts (2 adjustments + 1 final = 3 calls max) Expect(callbackCount).To(BeNumerically("<=", 3)) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) }) It("should handle skip during adjustment loop", func() { mockTool := mock.NewMockTool("search", "Search for information") callbackCount := 0 // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) // Adjustment attempt mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted"}`) mockLLM.SetAskResponse("LLM result") result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { callbackCount++ if callbackCount == 1 { // First call: provide adjustment return ToolCallDecision{ Approved: true, Adjustment: "Please adjust", } } // Second call: skip return ToolCallDecision{ Approved: true, Skip: true, } })) // When skipping with DisableToolReEvaluator, we might get ErrNoToolSelected if err != nil { Expect(err).To(Equal(ErrNoToolSelected)) } Expect(callbackCount).To(Equal(2)) // Tool should not be executed Expect(len(result.Status.ToolsCalled)).To(Equal(0)) }) It("should handle direct modification during adjustment loop", func() { mockTool := mock.NewMockTool("search", "Search for information") mock.SetRunResult(mockTool, "Directly modified result") // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) // Adjustment attempt (will be modified directly, so this won't be used) mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted"}`) mockLLM.SetAskResponse("LLM result") // After modification, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { // First call: provide adjustment if tool.Arguments["query"] == "original" { return ToolCallDecision{ Approved: true, Adjustment: "Please adjust", } } // During adjustment: directly modify modified := *tool modified.Arguments = map[string]any{ "query": "directly_modified", } return ToolCallDecision{ Approved: true, Modified: &modified, } })) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) // Check that directly modified arguments were used Expect(result.Status.ToolResults[0].ToolArguments.Arguments["query"]).To(Equal("directly_modified")) }) }) Context("SessionState and Resume", func() { It("should create SessionState with ToolChoice and Fragment", func() { mockTool := mock.NewMockTool("search", "Search for information") var savedState *SessionState // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Test result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) _, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { savedState = state return ToolCallDecision{Approved: true} })) Expect(err).ToNot(HaveOccurred()) Expect(savedState).ToNot(BeNil()) Expect(savedState.ToolChoice).ToNot(BeNil()) Expect(savedState.ToolChoice.Name).To(Equal("search")) Expect(savedState.Fragment).ToNot(BeNil()) }) It("should resume execution from SessionState", func() { mockTool := mock.NewMockTool("search", "Search for information") var savedState *SessionState // First execution - interrupt after saving state mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mockLLM.SetAskResponse("LLM result") _, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { savedState = state return ToolCallDecision{Approved: false} // Interrupt })) Expect(err).To(HaveOccurred()) Expect(savedState).ToNot(BeNil()) // Resume execution mock.SetRunResult(mockTool, "Resumed result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) resumedFragment, err := savedState.Resume(mockLLM, WithTools(mockTool)) Expect(err).ToNot(HaveOccurred()) Expect(len(resumedFragment.Status.ToolsCalled)).To(Equal(1)) Expect(resumedFragment.Status.ToolResults[0].Result).To(Equal("Resumed result")) }) }) Context("WithStartWithAction", func() { It("should start execution with a pre-selected tool", func() { mockTool := mock.NewMockTool("search", "Search for information") mock.SetRunResult(mockTool, "Pre-selected result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) initialTool := &ToolChoice{ Name: "search", Arguments: map[string]any{ "query": "pre_selected_query", }, } result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithStartWithAction(initialTool)) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) Expect(result.Status.ToolResults[0].ToolArguments.Arguments["query"]).To(Equal("pre_selected_query")) }) It("should start execution with multiple pre-selected tools", func() { mockSearchTool := mock.NewMockTool("search", "Search for information") mockWeatherTool := mock.NewMockTool("get_weather", "Get weather information") mock.SetRunResult(mockSearchTool, "Search result") mock.SetRunResult(mockWeatherTool, "Weather result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) initialTools := []*ToolChoice{ { Name: "search", Arguments: map[string]any{ "query": "test query", }, }, { Name: "get_weather", Arguments: map[string]any{ "city": "San Francisco", }, }, } result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockSearchTool, mockWeatherTool), WithStartWithAction(initialTools...)) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(2)) Expect(result.Status.ToolResults[0].ToolArguments.Arguments["query"]).To(Equal("test query")) Expect(result.Status.ToolResults[1].ToolArguments.Arguments["city"]).To(Equal("San Francisco")) }) }) Context("Multiple Tool Selection", func() { It("should handle multiple tool selections sequentially", func() { mockSearchTool := mock.NewMockTool("search", "Search for information") mockWeatherTool := mock.NewMockTool("get_weather", "Get weather information") mock.SetRunResult(mockSearchTool, "Search result") mock.SetRunResult(mockWeatherTool, "Weather result") mockLLM.SetAskResponse("LLM result") // LLM selects multiple tools in a single response mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), ToolCalls: []openai.ToolCall{ { ID: "call_1", Type: openai.ToolTypeFunction, Function: openai.FunctionCall{ Name: "search", Arguments: `{"query": "test"}`, }, }, { ID: "call_2", Type: openai.ToolTypeFunction, Function: openai.FunctionCall{ Name: "get_weather", Arguments: `{"city": "SF"}`, }, }, }, }, }, }, }) // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockSearchTool, mockWeatherTool)) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(2)) Expect(result.Status.ToolResults[0].Name).To(Equal("search")) Expect(result.Status.ToolResults[1].Name).To(Equal("get_weather")) }) }) Context("Parallel Tool Execution", func() { It("should execute multiple tools in parallel when enabled", func() { mockSearchTool := mock.NewMockTool("search", "Search for information") mockWeatherTool := mock.NewMockTool("get_weather", "Get weather information") mock.SetRunResult(mockSearchTool, "Search result") mock.SetRunResult(mockWeatherTool, "Weather result") mockLLM.SetAskResponse("LLM result") // LLM selects multiple tools using the parallel intention tool // First, reasoning step - now uses the reasoning tool mockLLM.AddCreateChatCompletionFunction("reasoning", `{"reasoning": "I need to search and get weather information."}`) // Then, tool selection with multiple tools mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), ToolCalls: []openai.ToolCall{ { ID: "call_1", Type: openai.ToolTypeFunction, Function: openai.FunctionCall{ Name: "pick_tools", Arguments: `{"tools": ["search", "get_weather"], "reasoning": "Both tools are needed"}`, }, }, }, }, }, }, }) // Parameter generation for search (with forced reasoning, needs parameter reasoning first) // 1. Parameter reasoning for search - now uses the reasoning tool mockLLM.AddCreateChatCompletionFunction("reasoning", `{"reasoning": "The search tool needs a query parameter to search for information."}`) // 2. Parameter generation for search mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) // Parameter generation for weather (with forced reasoning, needs parameter reasoning first) // 3. Parameter reasoning for weather - now uses the reasoning tool mockLLM.AddCreateChatCompletionFunction("reasoning", `{"reasoning": "The weather tool needs a city parameter to get weather information."}`) // 4. Parameter generation for weather mockLLM.AddCreateChatCompletionFunction("get_weather", `{"city": "SF"}`) // After tool execution, ToolReEvaluator uses forced reasoning, so it needs: // 1. Reasoning step response - now uses the reasoning tool mockLLM.AddCreateChatCompletionFunction("reasoning", `{"reasoning": "The tools have been executed successfully. No more tools are needed."}`) // 2. Intention tool response (return sink state "reply" to indicate no tools needed) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), ToolCalls: []openai.ToolCall{ { ID: "call_reval", Type: openai.ToolTypeFunction, Function: openai.FunctionCall{ Name: "pick_tools", Arguments: `{"tools": ["reply"], "reasoning": "No more tools needed"}`, }, }, }, }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockSearchTool, mockWeatherTool), EnableParallelToolExecution, WithForceReasoning(), WithForceReasoningTool()) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Status.ToolsCalled)).To(Equal(2)) }) }) Context("WithMaxAdjustmentAttempts", func() { It("should use default max adjustment attempts when not specified", func() { mockTool := mock.NewMockTool("search", "Search for information") callbackCount := 0 // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "original"}`) // Adjustment attempts mockLLM.AddCreateChatCompletionFunction("search", `{"query": "adjusted"}`) mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") // After tool execution, ToolReEvaluator returns no tool mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), // Don't specify WithMaxAdjustmentAttempts - should use default (5) WithToolCallBack(func(tool *ToolChoice, state *SessionState) ToolCallDecision { callbackCount++ if callbackCount == 1 { return ToolCallDecision{ Approved: true, Adjustment: "Adjust", } } return ToolCallDecision{Approved: true} })) Expect(err).ToNot(HaveOccurred()) Expect(callbackCount).To(Equal(2)) Expect(len(result.Status.ToolsCalled)).To(Equal(1)) }) }) }) var _ = Describe("ExecuteTools with Compaction", func() { var mockLLM *mock.MockOpenAIClient var originalFragment Fragment BeforeEach(func() { mockLLM = mock.NewMockOpenAIClient() originalFragment = NewEmptyFragment(). AddMessage(UserMessageRole, "Task 1"). AddMessage(AssistantMessageRole, "Done 1") }) Context("WithCompactionThreshold", func() { It("should not compact when threshold is disabled (0)", func() { mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mockTool := mock.NewMockTool("search", "Search for information") mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") mockLLM.SetUsage(100, 100, 1000) mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Usage: openai.Usage{ TotalTokens: 1000, PromptTokens: 100, CompletionTokens: 100, }, Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) result, err := ExecuteTools(mockLLM, originalFragment, WithTools(mockTool), WithCompactionThreshold(0), ) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Messages)).ToNot(Equal(len(originalFragment.Messages)), fmt.Sprintf("result: %+v", result)) Expect(result.Status.LastUsage.TotalTokens).To(BeNumerically(">", 0)) Expect(len(result.Messages)).To(Equal(6)) }) It("should not compact when tokens below threshold", func() { mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mockTool := mock.NewMockTool("search", "Search for information") mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) // Create fragment with low token count smallFragment := NewEmptyFragment(). AddMessage(UserMessageRole, "Hi"). AddMessage(AssistantMessageRole, "Hello") result, err := ExecuteTools(mockLLM, smallFragment, WithTools(mockTool), WithCompactionThreshold(100000), WithCompactionKeepMessages(2)) Expect(err).ToNot(HaveOccurred()) // Should not be compacted - still has original messages Expect(len(result.Messages)).To(BeNumerically(">", 2)) }) It("should compact when token threshold is exceeded", func() { mockTool := mock.NewMockTool("search", "Search for information") // First tool selection mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") // After tool execution, no more tools needed mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) // Create a large fragment with high token count largeFragment := NewEmptyFragment(). AddMessage(UserMessageRole, "Task 1"). AddMessage(AssistantMessageRole, "Answer to task 1"). AddMessage(ToolMessageRole, "Result 1"). AddMessage(UserMessageRole, "Task 2"). AddMessage(AssistantMessageRole, "Answer to task 2"). AddMessage(ToolMessageRole, "Result 2") // Set the usage to exceed threshold mockLLM.SetUsage(100, 100, 5000) // Mock the compaction summary response summaryFragment := NewEmptyFragment(). AddMessage(AssistantMessageRole, "Summary of conversation history.") mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) result, err := ExecuteTools(mockLLM, largeFragment, WithTools(mockTool), WithCompactionThreshold(1000), WithCompactionKeepMessages(1)) Expect(err).ToNot(HaveOccurred()) Expect(len(result.Messages)).To(BeNumerically(">", 0)) Expect(result.Messages[0].Role).To(Equal("system"), fmt.Sprintf("result: %+v", result)) Expect(result.Messages[0].Content).To(ContainSubstring("compacted"), fmt.Sprintf("result: %+v", result)) Expect(len(result.Messages)).To(BeNumerically("<", len(largeFragment.Messages))) }) It("should preserve parent fragment after compaction", func() { mockTool := mock.NewMockTool("search", "Search for information") mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) // Create a fragment with a parent parentFragment := NewEmptyFragment().AddMessage(UserMessageRole, "Parent task") largeFragment := NewEmptyFragment(). AddMessage(UserMessageRole, "Task 1"). AddMessage(AssistantMessageRole, strings.Repeat("response ", 5000)) largeFragment.ParentFragment = &parentFragment // Set usage to exceed threshold mockLLM.SetUsage(100, 100, 5000) // Mock the compaction summary response (may be used in-loop and again before final Ask) summaryFragment := NewEmptyFragment(). AddMessage(AssistantMessageRole, "Summary of conversation.") mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) result, err := ExecuteTools(mockLLM, largeFragment, WithTools(mockTool), WithCompactionThreshold(1000), WithCompactionKeepMessages(1)) Expect(err).ToNot(HaveOccurred()) Expect(result.ParentFragment).ToNot(BeNil()) Expect(result.ParentFragment.Messages[0].Role).To(Equal(UserMessageRole.String())) }) It("should preserve status after compaction", func() { mockTool := mock.NewMockTool("search", "Search for information") mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) // Create fragment with status largeFragment := NewEmptyFragment(). AddMessage(UserMessageRole, "Task 1"). AddMessage(AssistantMessageRole, strings.Repeat("response ", 5000)) largeFragment.Status = &Status{ Iterations: 5, ReasoningLog: []string{"reasoning1", "reasoning2"}, } // Set usage to exceed threshold mockLLM.SetUsage(100, 100, 5000) // Mock the compaction summary response (may be used in-loop and again before final Ask) summaryFragment := NewEmptyFragment(). AddMessage(AssistantMessageRole, "Summary of conversation.") mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) result, err := ExecuteTools(mockLLM, largeFragment, WithTools(mockTool), WithCompactionThreshold(1000), WithCompactionKeepMessages(1)) Expect(err).ToNot(HaveOccurred()) Expect(result.Status).ToNot(BeNil()) // Original had Iterations: 5; one tool loop iteration was run, so 6 Expect(result.Status.Iterations).To(Equal(6)) Expect(result.Status.ReasoningLog).To(Equal([]string{"reasoning1", "reasoning2"})) }) It("should use rough token estimate when LastUsage is not set", func() { mockTool := mock.NewMockTool("search", "Search for information") mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`) mock.SetRunResult(mockTool, "Result") mockLLM.SetAskResponse("LLM result") mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{ Choices: []openai.ChatCompletionChoice{ { Message: openai.ChatCompletionMessage{ Role: AssistantMessageRole.String(), Content: "No more tools needed.", }, }, }, }) // Large fragment without LastUsage set largeFragment := NewEmptyFragment(). AddMessage(UserMessageRole, "Task 1"). AddMessage(AssistantMessageRole, strings.Repeat("response with lots of content ", 500)). AddMessage(ToolMessageRole, "Result 1") // Mock the compaction summary response (may be used in-loop and again before final Ask) summaryFragment := NewEmptyFragment(). AddMessage(AssistantMessageRole, "Summary.") mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) mockLLM.AskResponses = append([]Fragment{summaryFragment}, mockLLM.AskResponses...) result, err := ExecuteTools(mockLLM, largeFragment, WithTools(mockTool), WithCompactionThreshold(1000), WithCompactionKeepMessages(1)) Expect(err).ToNot(HaveOccurred()) // Should be compacted based on rough estimate if len(result.Messages) > 0 { Expect(result.Messages[0].Role).To(Equal("system")) } }) }) })