mirror of
https://github.com/mudler/cogito.git
synced 2026-07-24 19:05:32 -04:00
37b69f1aec
* feat: add automatic conversation compaction based on token threshold This commit adds automatic conversation compaction to prevent context overflow during long-running tool execution sessions. Key changes: - Added LLMUsage struct to track token usage from LLM responses - Modified LLM interface to return token usage alongside Fragment - Added WithCompactionThreshold option to set token count threshold - Added WithCompactionKeepMessages option to configure recent messages to keep - Added compaction logic in ExecuteTools after LLM calls - Added helper functions: compactFragment, checkAndCompact, estimateTokens - Added PromptConversationCompaction for generating conversation summaries - Updated OpenAI and LocalAI clients to return token usage - Updated mock client for testing When compactionThreshold is set (> 0), the conversation will be automatically compacted when estimated token count exceeds the threshold. The compaction generates a summary of the conversation history using an LLM call while preserving recent messages. Signed-off-by: Autonomous Coding Agent <agent@autonomous> * fix: use actual usage tokens from LLM response for compaction - Store LastUsage in Status struct from LLM responses - checkAndCompact now uses actual TotalTokens from LLM response - Removed estimateTokens function (no longer needed) - Fallback estimate only used on first iteration when no usage data available * fix: capture usage tokens after sink state LLM call for compaction The sink state handling was not capturing usage tokens from the LLM response, which meant the compaction check would use the rough estimate instead of actual usage tokens. This change ensures LastUsage is stored after the llm.Ask call in the hasSinkState block, allowing proper token-based compaction. * fix: move compaction check to beginning of tool loop - Removed compaction check after max iterations (not needed) - Removed compaction check after sink state (not needed) - Added compaction check at beginning of tool loop (after totalIterations++) - Uses actual usage tokens from LLM response * fix: update Ask to return usage tokens from LocalAIClient * fix: set LastUsage in Ask function return fragment This addresses reviewer feedback that Ask() should automatically update the Fragment's LastUsage, not have callers do it. The OpenAIClient and LocalAIClient Ask functions now set Status.LastUsage before returning. * refactor: Ask() updates Fragment.Status.LastUsage directly Instead of returning LLMUsage from Ask(), the LLM clients now update the Fragment's Status.LastUsage directly. This simplifies the interface and ensures usage is always tracked in the fragment. Changes: - LLM.Ask() now returns (Fragment, error) instead of (Fragment, LLMUsage, error) - Clients (openai_client.go, localai_client.go) set LastUsage on the returned fragment - Mock client also updated to set usage in Status - All callers updated to use new 2-value return signature This addresses reviewer feedback on PR #41. * Apply suggestion from @mudler * Apply suggestion from @mudler * Apply suggestion from @mudler * Apply suggestion from @mudler * Apply suggestions from code review * Apply suggestion from @mudler * test: add mocked tests for compaction functionality - Add DefaultPrompts() function to prompt package for tests - Export CompactFragment and CheckAndCompact functions for testing - Add comprehensive unit tests for compaction logic using mocks - Remove duplicate Ginkgo compaction tests that have import issues * test: add compaction tests to tools_test.go suite Add Ginkgo tests for compaction functionality within the existing tools_test.go suite. Tests cover: - No compaction when threshold is disabled (0) - No compaction when tokens below threshold - Compaction when token threshold is exceeded - Parent fragment preservation after compaction - Status preservation after compaction - Rough token estimate usage when LastUsage is not set This addresses the reviewer's request to keep tests consistent with other ginkgo tests in tools_test.go. * chore: run go fmt and add compaction docs to README * Make CompactFragment and CheckAndCompact private Per reviewer request: - Changed CompactFragment to compactFragment (private) - Changed CheckAndCompact to checkAndCompact (private) - Removed tools_compaction_test.go (tests should be in tools_test.go) The compaction functionality is still available internally via ExecuteTools with WithCompactionThreshold option. * chore: remove exported functions from README, keep them private * chore: verify all changes applied - build passes * chore: verify build and vet pass --------- Signed-off-by: Autonomous Coding Agent <agent@autonomous> Co-authored-by: Ettore Di Giacinto <mudler@users.noreply.github.com>
1237 lines
47 KiB
Go
1237 lines
47 KiB
Go
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(5))
|
||
})
|
||
|
||
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"))
|
||
}
|
||
})
|
||
})
|
||
})
|