diff --git a/adk/prebuilt/plan_execute.go b/adk/prebuilt/plan_execute.go index 044b67b1..c6515465 100644 --- a/adk/prebuilt/plan_execute.go +++ b/adk/prebuilt/plan_execute.go @@ -18,6 +18,7 @@ package prebuilt import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -33,9 +34,22 @@ import ( "github.com/cloudwego/eino/schema" ) -// Plan represents a structured output format with a list of steps to be executed. -// This struct is used for JSON serialization/deserialization of the plan output -// generated by the model. +// Plan represents an execution plan with a sequence of actionable steps. +// It supports JSON serialization and provides access to the first step. +type Plan interface { + // FirstStep returns the first step to be executed in the plan. + FirstStep() string + + // MarshalJSON marshals the Plan into JSON. + json.Marshaler + // UnmarshalJSON unmarshals the JSON into the Plan. + json.Unmarshaler +} + +// NewPlan creates a new Plan instance. +type NewPlan func(ctx context.Context) Plan + +// defaultPlan the default implementation of Plan. // // JSON Schema: // @@ -52,12 +66,29 @@ import ( // }, // "required": ["steps"] // } -type Plan struct { +type defaultPlan struct { // Steps contains the ordered list of actions to be taken. // Each step should be clear, actionable, and arranged in a logical sequence. Steps []string `json:"steps"` } +func (p *defaultPlan) FirstStep() string { + if len(p.Steps) == 0 { + return "" + } + return p.Steps[0] +} + +func (p *defaultPlan) MarshalJSON() ([]byte, error) { + type planTyp defaultPlan + return sonic.Marshal((*planTyp)(p)) +} + +func (p *defaultPlan) UnmarshalJSON(bytes []byte) error { + type planTyp defaultPlan + return sonic.Unmarshal(bytes, (*planTyp)(p)) +} + // Response represents the final response to the user. // This struct is used for JSON serialization/deserialization of the final response // generated by the model. @@ -101,20 +132,57 @@ var ( ), } - ExecutorUserPrompt = prompt.FromMessages(schema.FString, schema.UserMessage( - "Given the following plan:\n{plan}\nYour task is to execute the first step, which is: {task}.")) + // PlannerPrompt is the prompt template for the planner. + // It provides context and guidance to the planner on how to generate the Plan. + PlannerPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage(`You are an expert planning agent. Given an objective, create a comprehensive step-by-step plan to achieve the objective. - ReplannerUserPrompt = prompt.FromMessages(schema.FString, schema.UserMessage( - `You are going to review the progress toward an objective. Analyze the current state and determine the optimal next action. +## YOUR TASK +Analyze the objective and generate a strategic plan that breaks down the goal into manageable, executable steps. -## OBJECTIVE -{input} +## PLANNING REQUIREMENTS +Each step in your plan must be: +- **Specific and actionable**: Clear instructions that can be executed without ambiguity +- **Self-contained**: Include all necessary context, parameters, and requirements +- **Independently executable**: Can be performed by an agent without dependencies on other steps +- **Logically sequenced**: Arranged in optimal order for efficient execution +- **Objective-focused**: Directly contribute to achieving the main goal -## ORIGINAL PLAN -{plan} +## PLANNING GUIDELINES +- Eliminate redundant or unnecessary steps +- Include relevant constraints, parameters, and success criteria for each step +- Ensure the final step produces a complete answer or deliverable +- Anticipate potential challenges and include mitigation strategies +- Structure steps to build upon each other logically +- Provide sufficient detail for successful execution +## QUALITY CRITERIA +- Plan completeness: Does it address all aspects of the objective? +- Step clarity: Can each step be understood and executed independently? +- Logical flow: Do steps follow a sensible progression? +- Efficiency: Is this the most direct path to the objective? +- Adaptability: Can the plan handle unexpected results or changes?`), + schema.MessagesPlaceholder("input", false), + ) + + // ExecutorPrompt is the prompt template for the executor. + // It provides context and guidance to the executor on how to execute the Task. + ExecutorPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage(`You are a diligent and meticulous executor agent. Follow the given plan and execute your tasks carefully and thoroughly.`), + schema.UserMessage(`## OBJECTIVE +{input} +## Given the following plan: +{plan} ## COMPLETED STEPS & RESULTS -{executed_results} +{executed_steps} +## Your task is to execute the first step, which is: +{step}`)) + + // ReplannerPrompt is the prompt template for the replanner. + // It provides context and guidance to the replanner on how to regenerate the Plan. + ReplannerPrompt = prompt.FromMessages(schema.FString, + schema.SystemMessage( + `You are going to review the progress toward an objective. Analyze the current state and determine the optimal next action. ## YOUR TASK Based on the progress above, you MUST choose exactly ONE action: @@ -149,50 +217,30 @@ Each step in your plan must be: - Has the original objective been completely satisfied? - Are there any remaining requirements or sub-goals? - Do the results suggest a need for strategy adjustment? -- What specific actions are still required?`)) +- What specific actions are still required?`), + schema.UserMessage(`## OBJECTIVE +{input} + +## ORIGINAL PLAN +{plan} + +## COMPLETED STEPS & RESULTS +{executed_steps}`), + ) ) const ( + // PlanExecuteUserInputSessionKey is the session key for the user input. PlanExecuteUserInputSessionKey = "UserInput" - // PlannerInstruction is the system instruction for the planner. - // It provides context and guidance to the planner on how to generate the Plan. - PlannerInstruction = `You are an expert planning agent. Given an objective, create a comprehensive step-by-step plan to achieve the objective. - -## YOUR TASK -Analyze the objective and generate a strategic plan that breaks down the goal into manageable, executable steps. - -## PLANNING REQUIREMENTS -Each step in your plan must be: -- **Specific and actionable**: Clear instructions that can be executed without ambiguity -- **Self-contained**: Include all necessary context, parameters, and requirements -- **Independently executable**: Can be performed by an agent without dependencies on other steps -- **Logically sequenced**: Arranged in optimal order for efficient execution -- **Objective-focused**: Directly contribute to achieving the main goal - -## PLANNING GUIDELINES -- Eliminate redundant or unnecessary steps -- Include relevant constraints, parameters, and success criteria for each step -- Ensure the final step produces a complete answer or deliverable -- Anticipate potential challenges and include mitigation strategies -- Structure steps to build upon each other logically -- Provide sufficient detail for successful execution - -## QUALITY CRITERIA -- Plan completeness: Does it address all aspects of the objective? -- Step clarity: Can each step be understood and executed independently? -- Logical flow: Do steps follow a sensible progression? -- Efficiency: Is this the most direct path to the objective? -- Adaptability: Can the plan handle unexpected results or changes?` - // PlanSessionKey is the session key for the plan. PlanSessionKey = "Plan" - // ExecuteResultSessionKey is the session key for the execute result. - ExecuteResultSessionKey = "ExecuteResult" + // ExecutedStepSessionKey is the session key for the execute result. + ExecutedStepSessionKey = "ExecutedStep" - // ExecuteResultsSessionKey is the session key for the execute results. - ExecuteResultsSessionKey = "ExecuteResults" + // ExecutedStepsSessionKey is the session key for the execute results. + ExecutedStepsSessionKey = "ExecutedSteps" ) // PlannerConfig provides configuration options for creating a planner agent. @@ -214,16 +262,43 @@ type PlannerConfig struct { // If not provided, PlanToolInfo will be used as the default. ToolInfo *schema.ToolInfo - // Instruction is the system instruction for the planner. - // It provides context and guidance to the planner on how to generate the Plan. - // If not provided, PlannerInstruction will be used as the default. - Instruction string + // GenInputFn is a function that generates the input messages for the planner. + // If not provided, defaultGenPlannerInputFn will be used as the default. + GenInputFn GenPlannerInputFn + + // NewPlan creates a new Plan instance for JSON. + // The returned Plan will be used to unmarshal the model-generated JSON output. + // If not provided, defaultNewPlan will be used as the default. + NewPlan NewPlan +} + +// PlannerInput is the input information for the planner. +type PlannerInput struct { + Input []adk.Message +} + +// GenPlannerInputFn is a function that generates the input messages for the planner. +type GenPlannerInputFn func(ctx context.Context, in *PlannerInput) ([]adk.Message, error) + +func defaultNewPlan(ctx context.Context) Plan { + return &defaultPlan{} +} + +func defaultGenPlannerInputFn(ctx context.Context, in *PlannerInput) ([]adk.Message, error) { + msgs, err := PlannerPrompt.Format(ctx, map[string]any{ + "input": in.Input, + }) + if err != nil { + return nil, err + } + return msgs, nil } type planner struct { - toolCall bool - chatModel model.BaseChatModel - sysMsg *schema.Message + toolCall bool + chatModel model.BaseChatModel + genInputFn GenPlannerInputFn + newPlan NewPlan } func (p *planner) Name(_ context.Context) string { @@ -260,9 +335,13 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, generator.Close() }() - msgs := make([]*schema.Message, 0, len(input.Messages)+1) - msgs = append(msgs, p.sysMsg) - msgs = append(msgs, input.Messages...) + msgs, err := p.genInputFn(ctx, &PlannerInput{ + Input: input.Messages, + }) + if err != nil { + generator.Send(&adk.AgentEvent{Err: err}) + return + } var modelCallOptions []model.Option if p.toolCall { modelCallOptions = append(modelCallOptions, model.WithToolChoice(schema.ToolChoiceForced)) @@ -270,9 +349,9 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, var msg adk.Message if input.EnableStreaming { - s, err := p.chatModel.Stream(ctx, msgs, modelCallOptions...) - if err != nil { - generator.Send(&adk.AgentEvent{Err: err}) + s, err_ := p.chatModel.Stream(ctx, msgs, modelCallOptions...) + if err_ != nil { + generator.Send(&adk.AgentEvent{Err: err_}) return } @@ -287,9 +366,9 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, event := adk.EventFromMessage(nil, sOutput, schema.Assistant, "") generator.Send(event) - msg, err = schema.ConcatMessageStream(ss[1]) - if err != nil { - generator.Send(&adk.AgentEvent{Err: err}) + msg, err_ = schema.ConcatMessageStream(ss[1]) + if err_ != nil { + generator.Send(&adk.AgentEvent{Err: err_}) return } @@ -298,10 +377,10 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, return } } else { - var err error - msg, err = p.chatModel.Generate(ctx, msgs, modelCallOptions...) - if err != nil { - generator.Send(&adk.AgentEvent{Err: err}) + var err_ error + msg, err_ = p.chatModel.Generate(ctx, msgs, modelCallOptions...) + if err_ != nil { + generator.Send(&adk.AgentEvent{Err: err_}) return } @@ -320,21 +399,21 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, generator.Send(event) } - var plan Plan var planJSON string if p.toolCall { planJSON = msg.ToolCalls[0].Function.Arguments } else { planJSON = msg.Content } - err := sonic.UnmarshalString(planJSON, &plan) + plan := p.newPlan(ctx) + err = plan.UnmarshalJSON([]byte(planJSON)) if err != nil { err = fmt.Errorf("unmarshal plan error: %w", err) generator.Send(&adk.AgentEvent{Err: err}) return } - adk.SetSessionValue(ctx, PlanSessionKey, &plan) + adk.SetSessionValue(ctx, PlanSessionKey, plan) }() return iterator @@ -348,12 +427,6 @@ func (p *planner) Run(ctx context.Context, input *adk.AgentInput, // If ToolCallingChatModel is provided, it will be configured with ToolInfo (or PlanToolInfo by default) // to generate structured Plan output. func NewPlanner(_ context.Context, cfg *PlannerConfig) (adk.Agent, error) { - instruction := cfg.Instruction - if instruction == "" { - instruction = PlannerInstruction - } - sysMsg := schema.SystemMessage(instruction) - var chatModel model.BaseChatModel var toolCall bool if cfg.ChatModelWithFormattedOutput != nil { @@ -372,73 +445,103 @@ func NewPlanner(_ context.Context, cfg *PlannerConfig) (adk.Agent, error) { } } + inputFn := cfg.GenInputFn + if inputFn == nil { + inputFn = defaultGenPlannerInputFn + } + + planParser := cfg.NewPlan + if planParser == nil { + planParser = defaultNewPlan + } + return &planner{ - toolCall: toolCall, - chatModel: chatModel, - sysMsg: sysMsg, + toolCall: toolCall, + chatModel: chatModel, + genInputFn: inputFn, + newPlan: planParser, }, nil } +// PlanExecuteInput is the input information for the executor and the planner. +type PlanExecuteInput struct { + Input []adk.Message + Plan Plan + ExecutedSteps []ExecutedStep +} + +// GenPlanExecuteInputFn is a function that generates the input messages for the executor. +type GenPlanExecuteInputFn func(ctx context.Context, in *PlanExecuteInput) ([]adk.Message, error) + +// ExecutorConfig provides configuration options for creating a executor agent. type ExecutorConfig struct { - Instruction string - Model model.ToolCallingChatModel + // Model is the chat model used by the executor. + Model model.ToolCallingChatModel + + // ToolsConfig is the tools configuration used by the executor. ToolsConfig adk.ToolsConfig - MaxStep int -} -type ExecuteResult struct { - Task string - Result string -} + // MaxStep is the maximum number of steps allowed for the executor. + MaxStep int -func formatPlan(p *Plan) string { - var formattedPlan strings.Builder - for i, step := range p.Steps { - formattedPlan.WriteString(fmt.Sprintf("%d. %s\n", i+1, step)) - } + // GenInputFn is the function that generates the input messages for the Executor. + GenInputFn GenPlanExecuteInputFn +} - return formattedPlan.String() +type ExecutedStep struct { + Step string + Result string } // NewExecutor creates a new executor agent. func NewExecutor(ctx context.Context, cfg *ExecutorConfig) (adk.Agent, error) { - genInput := func(ctx context.Context, instruction string, _ *adk.AgentInput) ([]adk.Message, error) { - msgs := make([]adk.Message, 0, 2) - if instruction != "" { - sp := schema.SystemMessage(instruction) - msgs = append(msgs, sp) - } + genInputFn := cfg.GenInputFn + if genInputFn == nil { + genInputFn = defaultGenExecutorInputFn + } + genInput := func(ctx context.Context, instruction string, _ *adk.AgentInput) ([]adk.Message, error) { plan, ok := adk.GetSessionValue(ctx, PlanSessionKey) if !ok { panic("impossible: plan not found") } - plan_ := plan.(*Plan) - task := plan_.Steps[0] + plan_ := plan.(Plan) - userMsgs, err := ExecutorUserPrompt.Format(ctx, map[string]any{ - "plan": formatPlan(plan_), - "task": task, - }) + userInput, ok := adk.GetSessionValue(ctx, PlanExecuteUserInputSessionKey) + if !ok { + panic("impossible: user input not found") + } + userInput_ := userInput.([]adk.Message) + + var executedSteps_ []ExecutedStep + executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey) + if ok { + executedSteps_ = executedStep.([]ExecutedStep) + } + + in := &PlanExecuteInput{ + Input: userInput_, + Plan: plan_, + ExecutedSteps: executedSteps_, + } + + msgs, err := genInputFn(ctx, in) if err != nil { return nil, err } - msgs = append(msgs, userMsgs...) - return msgs, nil } agent, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{ Name: "Executor", Description: "an executor agent", - Instruction: cfg.Instruction, Model: cfg.Model, ToolsConfig: cfg.ToolsConfig, GenModelInput: genInput, MaxStep: cfg.MaxStep, - OutputKey: ExecuteResultSessionKey, + OutputKey: ExecutedStepSessionKey, }) if err != nil { return nil, err @@ -447,16 +550,36 @@ func NewExecutor(ctx context.Context, cfg *ExecutorConfig) (adk.Agent, error) { return agent, nil } +func defaultGenExecutorInputFn(ctx context.Context, in *PlanExecuteInput) ([]adk.Message, error) { + + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + + userMsgs, err := ExecutorPrompt.Format(ctx, map[string]any{ + "input": formatInput(in.Input), + "plan": string(planContent), + "executed_steps": formatExecutedSteps(in.ExecutedSteps), + "step": in.Plan.FirstStep(), + }) + if err != nil { + return nil, err + } + + return userMsgs, nil +} + type replanner struct { chatModel model.ToolCallingChatModel - instruction string planTool *schema.ToolInfo respondTool *schema.ToolInfo + + genInputFn GenPlanExecuteInputFn + newPlan NewPlan } type ReplannerConfig struct { - // Instruction is the system instruction for the replanner. - Instruction string // ChatModel is the model that supports tool calling capabilities. // It will be configured with PlanTool and RespondTool to generate updated plans or responses. @@ -469,8 +592,18 @@ type ReplannerConfig struct { // RespondTool defines the schema for the response tool that can be used with ToolCallingChatModel. // If not provided, the default RespondToolInfo will be used. RespondTool *schema.ToolInfo + + // GenInputFn is the function that generates the input messages for the Replanner. + // if not provided, buildDefaultReplannerInputFn will be used. + GenInputFn GenPlanExecuteInputFn + + // NewPlan creates a new Plan instance. + // The returned Plan will be used to unmarshal the model-generated JSON output from PlanTool. + // If not provided, defaultNewPlan will be used as the default. + NewPlan NewPlan } +// formatInput formats the input messages into a string. func formatInput(input []adk.Message) string { var sb strings.Builder for _, msg := range input { @@ -481,10 +614,10 @@ func formatInput(input []adk.Message) string { return sb.String() } -func formatExecutedResults(results []ExecuteResult) string { +func formatExecutedSteps(results []ExecutedStep) string { var sb strings.Builder for _, result := range results { - sb.WriteString(fmt.Sprintf("Task: %s\nResult: %s\n\n", result.Task, result.Result)) + sb.WriteString(fmt.Sprintf("Step: %s\nResult: %s\n\n", result.Step, result.Result)) } return sb.String() @@ -499,37 +632,31 @@ func (r *replanner) Description(_ context.Context) string { } func (r *replanner) genInput(ctx context.Context) ([]adk.Message, error) { - msgs := make([]adk.Message, 0, 2) - - if r.instruction != "" { - sp := schema.SystemMessage(r.instruction) - msgs = append(msgs, sp) - } - executeResult, ok := adk.GetSessionValue(ctx, ExecuteResultSessionKey) + executedStep, ok := adk.GetSessionValue(ctx, ExecutedStepSessionKey) if !ok { panic("impossible: execute result not found") } - executeResult_ := executeResult.(string) + executedStep_ := executedStep.(string) plan, ok := adk.GetSessionValue(ctx, PlanSessionKey) if !ok { panic("impossible: plan not found") } - plan_ := plan.(*Plan) - task := plan_.Steps[0] + plan_ := plan.(Plan) + step := plan_.FirstStep() - var executeResults_ []ExecuteResult - executeResults, ok := adk.GetSessionValue(ctx, ExecuteResultsSessionKey) + var executedSteps_ []ExecutedStep + executedSteps, ok := adk.GetSessionValue(ctx, ExecutedStepsSessionKey) if ok { - executeResults_ = executeResults.([]ExecuteResult) + executedSteps_ = executedSteps.([]ExecutedStep) } - executeResults_ = append(executeResults_, ExecuteResult{ - Task: task, - Result: executeResult_, + executedSteps_ = append(executedSteps_, ExecutedStep{ + Step: step, + Result: executedStep_, }) - adk.SetSessionValue(ctx, ExecuteResultsSessionKey, executeResults_) + adk.SetSessionValue(ctx, ExecutedStepsSessionKey, executedSteps_) userInput, ok := adk.GetSessionValue(ctx, PlanExecuteUserInputSessionKey) if !ok { @@ -537,19 +664,20 @@ func (r *replanner) genInput(ctx context.Context) ([]adk.Message, error) { } userInput_ := userInput.([]adk.Message) - userMsgs, err := ReplannerUserPrompt.Format(ctx, map[string]any{ - "plan": formatPlan(plan_), - "input": formatInput(userInput_), - "executed_results": formatExecutedResults(executeResults_), - "plan_tool": r.planTool.Name, - "respond_tool": r.respondTool.Name, - }) + in := &PlanExecuteInput{ + Input: userInput_, + Plan: plan_, + ExecutedSteps: executedSteps_, + } + genInputFn := r.genInputFn + if genInputFn == nil { + genInputFn = buildDefaultReplannerInputFn(in, r.planTool.Name, r.respondTool.Name) + } + msgs, err := genInputFn(ctx, in) if err != nil { return nil, err } - msgs = append(msgs, userMsgs...) - return msgs, nil } @@ -661,20 +789,41 @@ func (r *replanner) Run(ctx context.Context, input *adk.AgentInput, _ ...adk.Age return } - var plan Plan - err = sonic.UnmarshalString(planMsg.ToolCalls[0].Function.Arguments, &plan) + plan_ := r.newPlan(ctx) + err = plan_.UnmarshalJSON([]byte(planMsg.ToolCalls[0].Function.Arguments)) if err != nil { err = fmt.Errorf("unmarshal plan error: %w", err) generator.Send(&adk.AgentEvent{Err: err}) return } - adk.SetSessionValue(ctx, PlanSessionKey, &plan) + adk.SetSessionValue(ctx, PlanSessionKey, plan_) }() return iterator } +func buildDefaultReplannerInputFn(in *PlanExecuteInput, planToolName, respondToolName string) GenPlanExecuteInputFn { + return func(ctx context.Context, in *PlanExecuteInput) ([]adk.Message, error) { + planContent, err := in.Plan.MarshalJSON() + if err != nil { + return nil, err + } + msgs, err := ReplannerPrompt.Format(ctx, map[string]any{ + "plan": string(planContent), + "input": formatInput(in.Input), + "executed_steps": formatExecutedSteps(in.ExecutedSteps), + "plan_tool": planToolName, + "respond_tool": respondToolName, + }) + if err != nil { + return nil, err + } + + return msgs, nil + } +} + func NewReplanner(_ context.Context, cfg *ReplannerConfig) (adk.Agent, error) { planTool := cfg.PlanTool if planTool == nil { @@ -691,25 +840,38 @@ func NewReplanner(_ context.Context, cfg *ReplannerConfig) (adk.Agent, error) { return nil, err } + planParser := cfg.NewPlan + if planParser == nil { + planParser = defaultNewPlan + } + return &replanner{ chatModel: chatModel, - instruction: cfg.Instruction, planTool: planTool, respondTool: respondTool, + genInputFn: cfg.GenInputFn, + newPlan: planParser, }, nil } +// PlanExecuteConfig provides configuration options for creating a plan execute agent. type PlanExecuteConfig struct { - Planner adk.Agent - Executor adk.Agent - Replanner adk.Agent + Planner adk.Agent + Executor adk.Agent + Replanner adk.Agent + MaxIterations int } +// NewPlanExecuteAgent creates a new plan execute agent with the given configuration. func NewPlanExecuteAgent(ctx context.Context, cfg *PlanExecuteConfig) (adk.Agent, error) { + maxIterations := cfg.MaxIterations + if maxIterations <= 0 { + maxIterations = 10 + } loop, err := adk.NewLoopAgent(ctx, &adk.LoopAgentConfig{ Name: "execute_replan", SubAgents: []adk.Agent{cfg.Executor, cfg.Replanner}, - MaxIterations: 10, + MaxIterations: maxIterations, }) if err != nil { return nil, err diff --git a/adk/prebuilt/plan_execute_test.go b/adk/prebuilt/plan_execute_test.go index 23862c6d..120fbd47 100644 --- a/adk/prebuilt/plan_execute_test.go +++ b/adk/prebuilt/plan_execute_test.go @@ -45,7 +45,6 @@ func TestNewPlannerWithFormattedOutput(t *testing.T) { // Create the PlannerConfig conf := &PlannerConfig{ ChatModelWithFormattedOutput: mockChatModel, - Instruction: "Custom instruction", } // Create the planner @@ -128,13 +127,14 @@ func TestPlannerRunWithFormattedOutput(t *testing.T) { event, ok = iterator.Next() assert.False(t, ok) - var plan Plan - err = sonic.UnmarshalString(msg.Content, &plan) + plan := defaultNewPlan(ctx) + err = plan.UnmarshalJSON([]byte(msg.Content)) assert.NoError(t, err) - assert.Equal(t, 3, len(plan.Steps)) - assert.Equal(t, "Step 1", plan.Steps[0]) - assert.Equal(t, "Step 2", plan.Steps[1]) - assert.Equal(t, "Step 3", plan.Steps[2]) + plan_ := plan.(*defaultPlan) + assert.Equal(t, 3, len(plan_.Steps)) + assert.Equal(t, "Step 1", plan_.Steps[0]) + assert.Equal(t, "Step 2", plan_.Steps[1]) + assert.Equal(t, "Step 3", plan_.Steps[2]) } // TestPlannerRunWithToolCalling tests the Run method of a planner created with ToolCallingChatModel @@ -194,13 +194,15 @@ func TestPlannerRunWithToolCalling(t *testing.T) { _, ok = iterator.Next() assert.False(t, ok) - var plan Plan - err = sonic.UnmarshalString(msg.Content, &plan) + plan := defaultNewPlan(ctx) + err = plan.UnmarshalJSON([]byte(msg.Content)) assert.NoError(t, err) - assert.Equal(t, 3, len(plan.Steps)) - assert.Equal(t, "Step 1", plan.Steps[0]) - assert.Equal(t, "Step 2", plan.Steps[1]) - assert.Equal(t, "Step 3", plan.Steps[2]) + plan_ := plan.(*defaultPlan) + assert.NoError(t, err) + assert.Equal(t, 3, len(plan_.Steps)) + assert.Equal(t, "Step 1", plan_.Steps[0]) + assert.Equal(t, "Step 2", plan_.Steps[1]) + assert.Equal(t, "Step 3", plan_.Steps[2]) } // TestNewExecutor tests the NewExecutor function @@ -216,9 +218,8 @@ func TestNewExecutor(t *testing.T) { // Create the ExecutorConfig conf := &ExecutorConfig{ - Instruction: "Custom instruction", - Model: mockToolCallingModel, - MaxStep: 5, + Model: mockToolCallingModel, + MaxStep: 5, } // Create the executor @@ -243,7 +244,7 @@ func TestExecutorRun(t *testing.T) { mockToolCallingModel := mockModel.NewMockToolCallingChatModel(ctrl) // Store a plan in the session - plan := &Plan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} adk.SetSessionValue(ctx, PlanSessionKey, plan) // Set up expectations for the mock model @@ -263,16 +264,16 @@ func TestExecutorRun(t *testing.T) { // Create the ExecutorConfig conf := &ExecutorConfig{ - Instruction: "Custom instruction", - Model: mockToolCallingModel, - MaxStep: 5, + Model: mockToolCallingModel, + MaxStep: 5, } // Create the executor executor, err := NewExecutor(ctx, conf) assert.NoError(t, err) executor, err = AgentWithSessionKVs(ctx, executor, map[string]any{ - PlanSessionKey: plan, + PlanSessionKey: plan, + PlanExecuteUserInputSessionKey: []adk.Message{schema.UserMessage("no input")}, }) assert.NoError(t, err) @@ -320,7 +321,6 @@ func TestNewReplanner(t *testing.T) { // Create the ReplannerConfig conf := &ReplannerConfig{ - Instruction: "Custom instruction", ChatModel: mockToolCallingModel, PlanTool: planTool, RespondTool: respondTool, @@ -378,7 +378,6 @@ func TestReplannerRunWithPlan(t *testing.T) { // Create the ReplannerConfig conf := &ReplannerConfig{ - Instruction: "Custom instruction", ChatModel: mockToolCallingModel, PlanTool: planTool, RespondTool: respondTool, @@ -389,10 +388,10 @@ func TestReplannerRunWithPlan(t *testing.T) { assert.NoError(t, err) // Store necessary values in the session - plan := &Plan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} rp, err = AgentWithSessionKVs(ctx, rp, map[string]any{ PlanSessionKey: plan, - ExecuteResultSessionKey: "Execution result", + ExecutedStepSessionKey: "Execution result", PlanExecuteUserInputSessionKey: []adk.Message{schema.UserMessage("User input")}, }) assert.NoError(t, err) @@ -417,19 +416,19 @@ func TestReplannerRunWithPlan(t *testing.T) { // Verify the updated plan was stored in the session planValue, ok := kvs[PlanSessionKey] assert.True(t, ok) - updatedPlan, ok := planValue.(*Plan) + updatedPlan, ok := planValue.(*defaultPlan) assert.True(t, ok) assert.Equal(t, 2, len(updatedPlan.Steps)) assert.Equal(t, "Updated Step 1", updatedPlan.Steps[0]) assert.Equal(t, "Updated Step 2", updatedPlan.Steps[1]) // Verify the execute results were updated - executeResultsValue, ok := kvs[ExecuteResultsSessionKey] + executeResultsValue, ok := kvs[ExecutedStepsSessionKey] assert.True(t, ok) - executeResults, ok := executeResultsValue.([]ExecuteResult) + executeResults, ok := executeResultsValue.([]ExecutedStep) assert.True(t, ok) assert.Equal(t, 1, len(executeResults)) - assert.Equal(t, "Step 1", executeResults[0].Task) + assert.Equal(t, "Step 1", executeResults[0].Step) assert.Equal(t, "Execution result", executeResults[0].Result) _, ok = iterator.Next() @@ -478,7 +477,6 @@ func TestReplannerRunWithRespond(t *testing.T) { // Create the ReplannerConfig conf := &ReplannerConfig{ - Instruction: "Custom instruction", ChatModel: mockToolCallingModel, PlanTool: planTool, RespondTool: respondTool, @@ -489,10 +487,10 @@ func TestReplannerRunWithRespond(t *testing.T) { assert.NoError(t, err) // Store necessary values in the session - plan := &Plan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + plan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} rp, err = AgentWithSessionKVs(ctx, rp, map[string]any{ PlanSessionKey: plan, - ExecuteResultSessionKey: "Execution result", + ExecutedStepSessionKey: "Execution result", PlanExecuteUserInputSessionKey: []adk.Message{schema.UserMessage("User input")}, }) assert.NoError(t, err) @@ -578,9 +576,9 @@ func TestPlanExecuteAgentWithReplan(t *testing.T) { mockReplanner.EXPECT().Description(gomock.Any()).Return("a replanner agent").AnyTimes() // Create a plan - originalPlan := &Plan{Steps: []string{"Step 1", "Step 2", "Step 3"}} + originalPlan := &defaultPlan{Steps: []string{"Step 1", "Step 2", "Step 3"}} // Create an updated plan with fewer steps (after replanning) - updatedPlan := &Plan{Steps: []string{"Updated Step 2", "Updated Step 3"}} + updatedPlan := &defaultPlan{Steps: []string{"Updated Step 2", "Updated Step 3"}} // Create execute result originalExecuteResult := "Execution result for Step 1" updatedExecuteResult := "Execution result for Updated Step 2" @@ -616,15 +614,15 @@ func TestPlanExecuteAgentWithReplan(t *testing.T) { iterator, generator := adk.NewAsyncIteratorPair[*adk.AgentEvent]() plan, _ := adk.GetSessionValue(ctx, PlanSessionKey) - currentPlan := plan.(*Plan) + currentPlan := plan.(*defaultPlan) var msg adk.Message // Check if this is the first replanning (original plan has 3 steps) if len(currentPlan.Steps) == 3 { msg = schema.AssistantMessage(originalExecuteResult, nil) - adk.SetSessionValue(ctx, ExecuteResultSessionKey, originalExecuteResult) + adk.SetSessionValue(ctx, ExecutedStepSessionKey, originalExecuteResult) } else { msg = schema.AssistantMessage(updatedExecuteResult, nil) - adk.SetSessionValue(ctx, ExecuteResultSessionKey, updatedExecuteResult) + adk.SetSessionValue(ctx, ExecutedStepSessionKey, updatedExecuteResult) } event := adk.EventFromMessage(msg, nil, schema.Assistant, "") generator.Send(event) @@ -642,7 +640,7 @@ func TestPlanExecuteAgentWithReplan(t *testing.T) { // First call: Update the plan // Get the current plan from the session plan, _ := adk.GetSessionValue(ctx, PlanSessionKey) - currentPlan := plan.(*Plan) + currentPlan := plan.(*defaultPlan) // Check if this is the first replanning (original plan has 3 steps) if len(currentPlan.Steps) == 3 { @@ -654,8 +652,8 @@ func TestPlanExecuteAgentWithReplan(t *testing.T) { // Set the updated plan & execute result in the session adk.SetSessionValue(ctx, PlanSessionKey, updatedPlan) - adk.SetSessionValue(ctx, ExecuteResultsSessionKey, []ExecuteResult{{ - Task: currentPlan.Steps[0], + adk.SetSessionValue(ctx, ExecutedStepsSessionKey, []ExecutedStep{{ + Step: currentPlan.Steps[0], Result: originalExecuteResult, }}) } else {