Code Sample — SmolAgents — for planning and data analysis

Recently I have started using and enjoying the SmolAgents library to put together agentic solutions. An example below shows how to use SmolAgents for a sample Data Analysis task.

"""
# !pip install smolagents[litellm]
"""

import os
from smolagents import CodeAgent, LiteLLMModel

from smolagents import tool

current_dir = os.path.dirname(os.path.abspath(__file__))

output_path = os.path.join(current_dir, "output")

openai_api_key = os.environ.get("OPENAI_API_KEY")
if openai_api_key is None:
    raise ValueError("Please set the OPENAI_API_KEY environment variable.")

@tool
def get_user_inputs(reason_for_inputs: str) -> str:
    """
    Function to get user inputs.
    Args:
        reason_for_inputs (str): Reason for requesting user inputs.

    Returns:
        str: User input for the data analysis request.
    """
    user_request = input(f" {reason_for_inputs}: ")
    return user_request

authorized_imports = [
    "os",
    "pandas",
    "numpy",
    "random",
    "string",
    "math",
    "numpy.random",
    "sklearn",
    "sklearn.cluster",
    "sklearn.metrics",
    "scipy",
    "scipy.*",
    "sklearn.*",
    "matplotlib",
    "matplotlib.pyplot",
    "seaborn",
    "plotly"
]

authorized_imports_str = ", ".join(authorized_imports)

model = LiteLLMModel(model_id="gpt-4o-mini", api_key=openai_api_key)

coder_sub_agent = CodeAgent(
    name="CoderAgent",
    description=f"This agent solves the user's query by generating code. It can use libraries like the ones built into Python, and can also use the following libraries: {authorized_imports_str}. ",
    model=model,
    tools=[],
    additional_authorized_imports=authorized_imports,
    max_steps=4,
)

manager_agent_prompt = """
This agent solves the user's query by adaptive planning which is dependent on the complexity of user's task. 
Perform extensive planning only if needed. 
If the task is simple, try to solve it in one or few steps. 

Also, this agent does not write code, it only manages the subagents, and generates an execution plan and tracks the progress of the subagents until the task is completed. 
It can solve problems with the help of subagents like CoderAgent to write and execute. 

On your plan indicate which steps can be executed in parallel and when possible, try to execute those steps in parallel without compromising on the plan quality. 
If you see an error in plan execution, try to fix it via a different approach.  

If the user has provided some datasets and context, then only use those datasets and context to solve the problem.
If you are not sure about the user's request, ask clarifying questions to get more information.
"""

manager_agent = CodeAgent(
    name="ManagerAgent",
    description=manager_agent_prompt,
    model=model,
    tools=[get_user_inputs],
    managed_agents=[coder_sub_agent],
    planning_interval=1,
    max_steps=50,
    additional_authorized_imports=authorized_imports,
)

"""Accept users request via input and run the agent until user enters quit or exit"""

while True:
    user_request = input("Enter your data analysis request (or type 'exit' to quit): ")

    if user_request.strip().lower() == "exit" or user_request.strip().lower() == "quit":
        print("Exiting the program.")
        break

    # Run the manager agent with the user's request
    manager_agent.run(user_request, reset=False)
    manager_agent.save(output_path)

Explanation of key points

• Above implements a sample ( but complete ) data analysis agentic workflow.
• We integrate with LiteLLM to support a diverse suite of models from a variety of providers.
• Overall lines of code is less than 100 since about 50 lines reference to prompts as well as dependencies.
• The multi agent system can plan out solutions to user’s request.
• We integrated with a simple tool to accept user’s request during code execution, as well as to suggest the next recommended action. We can generalize this to suggest multiple actions.
• We separated out planning / task management from code execution, and leverage a sub-agent for the latter.
• We give 50 step allowance to the top level agent, to be able to perform an elaborate task via planning.
• By setting planning interval as 1, we reflect on our plan at each step of the way, hopefully with better results than planning less frequently.
• We restrict our universe of libraries to the ones commonly used for data analysis tasks.
• When we run the agent we set, reset=False, so that we can continue the conversation over multiple tasks, thus maintaining conversation history.
• Last but not the least, we specify an output folder where we can easily save necessary artifacts from an agent run, including code, data analysis artifacts etc.
• Overall, we have a fairly powerful system at hand now, with minimal lines of code.