Is there a way to return plan of a CodeAgent?

[yjiahao]

Is your feature request related to a problem? Please describe.
Hi, I'm currently trying to create an agentic workflow with smolagents.

The workflow I want to create is this:

1. CodeAgent (call it agent from now on) takes in some query from user
2. agent formulates the plan to achieve the task, then sends the plan back to the user for approval
3. user approves/edits the plan and gives it back to the agent
4. agent executes this edited plan

Describe the solution you'd like
I think it would be beneficial if the user could directly access the agent's initial plan and correct the agent right there and then before it continues to retrieve the wrong information.

Is this not possible with the current options.
Not that I know of.

Describe alternatives you've considered
I am able to intercept the agent's steps and get the initial plan as shown below:

def get_plan(query, sql_agent, engine):
    # Reset agent memory to start fresh
    sql_agent.memory.steps = []
    
    # Get the planning step
    for step in sql_agent.run(query, stream=True, additional_args={'engine': engine}):
        if isinstance(step, PlanningStep):
            plan = step.plan
            # Store the original step for later reference
            return plan, step
    
    return "No plan was generated.", None

# Second function - Execute with the approved/amended plan
def execute_plan(query, amended_plan, original_step, sql_agent, engine):
    # Clear previous planning steps
    sql_agent.memory.steps = [s for s in sql_agent.memory.steps 
                             if not isinstance(s, PlanningStep)]
    
    # Create new planning step with amended plan
    new_planning_step = PlanningStep(
        plan=amended_plan,
        model_input_messages=original_step.model_input_messages,
        model_output_message=original_step.model_output_message,
        timing=original_step.timing  # Add the timing parameter from original step
    )
    
    # Add to memory
    sql_agent.memory.steps.append(new_planning_step)

    original_memory = sql_agent.memory.steps
    
    # Execute with the plan
    response = sql_agent.run(query, additional_args={'engine': engine})

    after_2nd_run_mem = sql_agent.memory.steps
    return response, original_memory, after_2nd_run_mem

# Example query
query = "Count how many distinct survey names we have in each database table"

# Step 1: Get the plan
print("Getting plan...")
plan, original_step = get_plan(query, sql_agent, engine)

print("=" * 80)
print("PROPOSED PLAN:")
print("=" * 80)
print(plan)
print("=" * 80)

amended_plan = input("Edit the plan if needed (or press Enter to keep as is):\n")

# Use original plan if no edits were made
if not amended_plan.strip():
    amended_plan = plan
    print("Using original plan")
else:
    print("Using amended plan")

print("\nExecuting plan...")
result, og_mem, final_mem = execute_plan(query, amended_plan, original_step, sql_agent, engine)

print("=" * 80)
print("RESULT:")
print("=" * 80)
print(result)

However, I think the agent's memory gets cleared with each .run() call, so even after appending the new plan, it forgets it after the second run.

[albertvillanova]

Hi @yjiahao.

Thanks for raising this interesting use case! 🤗

What about using step_callbacks to trigger agent.interrupt once the plan is ready? That would allow you to pause the execution, present the plan to the user for approval or editing, and then resume by calling run again with reset=False to preserve the agent's memory.

[albertvillanova]

Hi @yjiahao,

Please note that currently the step_callbacks are not called after the planning step. I proposed a PR to fix this.

In the meantime, you could run the agent step by step passing stream=True:

for step in agent.run(task, stream=True):
    if isinstance(step, Planningstep):
        # ask the user to confirm or edit the step.plan, then continue

[spisupat]

almost raised a duplicate issue/fix 😅

I would find this feature useful since it allows for human-in-the-loop intervention of the planning step

[albertvillanova]

I added an example demonstrating how to customize agent plan here: https://github.com/huggingface/smolagents/tree/main/examples/plan_customization

• Add complete Plan Customization example #1506
• Improve Plan Customization example with Human-in-the-Loop #1513