AI Agents Are All You Need ? Hands-on guide to creating a better code-solving AI Agent with AutoGen

7 min readMay 18, 2024

Introduction

Hey there! So Lately, there’s been a lot of buzz around AI agents. Intrigued by the hype, I did some reading and stumbled upon AutoGen, a great library for building AI agents. But what exactly are AI agents, and why do they matter? Let’s dive into it.

LLMs have been around for a few years now, and they are rapidly evolving towards AI Agents and Agentic workflows. Don’t get me wrong, LLMs are great, but they are still not efficient at automating the stuff on their own. LLMs combined with other tools are a really efficient way to utilize the general intelligence LLMs possess, by consuming massive amounts of language data. AI agents are software entities that can perform tasks autonomously, interacting with their environment to make decisions. They are revolutionizing how we approach problems, from customer service chatbots to complex data analysis. Today, I’m excited to show you how to build a simple AI agent using AutoGen.

Ready? Let’s go!

Why AI Agents Matter

AI agents bring several benefits to the table:

Automation: They handle repetitive tasks without human intervention.
Efficiency: They process data and perform tasks faster than humans.
Scalability: They can operate 24/7, providing services without downtime.
Intelligence: They learn and improve from interactions, getting better over time.
Given these advantages, AI agents are being adopted across various industries, from healthcare to finance.

Building a Simple AI Agent with AutoGen
Let’s get hands-on and build a basic AI agent using AutoGen. For this demonstration, we’ll solve the “Distribute Coins in a binary tree” problem from LeetCode. This problem is a classic in coding interviews, where you find two numbers in an array that add up to a target.

Colab Link — https://colab.research.google.com/drive/1dbP9XP-_yyAWEPmFqdIzvcn5XU6rlnVk?usp=sharing

Step 1: Install AutoGen

First, we need to install the AutoGen library. Run the following command in your environment:

!pip install pyautogen -q --progress-bar off

Step 2: Set Up Your Environment

I’m using Google Colab for this demonstration. Here’s how to set up your environment and load your OpenAI API key securely:

import os
import autogen
from google.colab import userdata

# Load your OpenAI API key
userdata.get('OPENAI_API_KEY')

llm_config = {
    "config_list": [{"model": "gpt-3.5-turbo", "api_key": userdata.get('OPENAI_API_KEY')}],
    "cache_seed": 0,
    "temperature": 0,
}

Step 3: Define the Problem

We’ll use the a LC Medium “Distribute Coins in a binary tree” problem from LeetCode.

LEETCODE_QUESTION = """
Title : Distribute Coins in a binary tree

You are given the root of a binary tree with n nodes where each node in the tree has node.val coins. There are n coins in total throughout the whole tree.

In one move, we may choose two adjacent nodes and move one coin from one node to another. A move may be from parent to child, or from child to parent.

Return the minimum number of moves required to make every node have exactly one coin.

 

Example 1:


Input: root = [3,0,0]
Output: 2
Explanation: From the root of the tree, we move one coin to its left child, and one coin to its right child.
Example 2:


Input: root = [0,3,0]
Output: 3
Explanation: From the left child of the root, we move two coins to the root [taking two moves]. Then, we move one coin from the root of the tree to the right child.

Follow-up: Can you come up with an algorithm that has a better time complexity?
"""

Step 4: Set Up the Agents

AutoGen uses two types of agents: the AssistantAgent and the UserProxyAgent. The AssistantAgent suggests code solutions, while the UserProxyAgent executes the code and checks if it works correctly.

SYSTEM_MESSAGE = """You are a helpful AI assistant. Solve tasks using your coding and language skills..."""

assistant = autogen.AssistantAgent(
    name="assistant",
    llm_config=llm_config,
    system_message=SYSTEM_MESSAGE
)

user_proxy = autogen.UserProxyAgent(
    name="user_proxy",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=4,
    is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"),
    code_execution_config={
        "work_dir": "coding",
        "use_docker": False,
    },
)

The human_input_mode is set to "NEVER" because we don’t plan to provide any input ourselves, and max_consecutive_auto_reply limits the conversation turns.

Step 5: Initiate the Conversation

We’ll start the conversation by sending a message from our UserProxyAgent to our Assistant.

with autogen.Cache.disk(cache_seed=7) as cache:
  chat_res = user_proxy.initiate_chat(
      assistant,
      message="""Solve the following LeetCode problem and also comment on its time and space complexity:\n\n""" + LEETCODE_QUESTION
  )

Output

Here’s a snippet of the conversation that follows:

User Proxy: Solve the following LeetCode problem and also comment on its time and space complexity…

Assistant: To solve this problem efficiently, we can use a dictionary to store the elements we have seen so far along with their indices. Here’s the Python code…

user_proxy (to assistant):

Solve the following LeetCode problem and also comment on its time and space complexity:


Title : Distribute Coins in a binary tree

You are given the root of a binary tree with n nodes where each node in the tree has node.val coins. There are n coins in total throughout the whole tree.

In one move, we may choose two adjacent nodes and move one coin from one node to another. A move may be from parent to child, or from child to parent.

Return the minimum number of moves required to make every node have exactly one coin.

 

Example 1:


Input: root = [3,0,0]
Output: 2
Explanation: From the root of the tree, we move one coin to its left child, and one coin to its right child.
Example 2:


Input: root = [0,3,0]
Output: 3
Explanation: From the left child of the root, we move two coins to the root [taking two moves]. Then, we move one coin from the root of the tree to the right child.

Follow-up: Can you come up with an algorithm that has a better time complexity?


--------------------------------------------------------------------------------
assistant (to user_proxy):

To solve this problem, we can perform a depth-first search (DFS) traversal of the binary tree and calculate the excess coins at each node. The excess coins at a node can be calculated as the sum of coins at the node and the excess coins from its left and right children, minus 1 (since the node itself should have 1 coin).

Here is the Python code to solve the problem:

```python
class Solution:
    def distributeCoins(self, root: TreeNode) -> int:
        self.moves = 0
        
        def dfs(node):
            if not node:
                return 0
            
            left_excess = dfs(node.left)
            right_excess = dfs(node.right)
            
            self.moves += abs(left_excess) + abs(right_excess)
            
            return node.val + left_excess + right_excess - 1
        
        dfs(root)
        
        return self.moves
```

Time complexity analysis:
- The DFS traversal visits each node once, so the time complexity is O(n), where n is the number of nodes in the binary tree.

Space complexity analysis:
- The space complexity is O(h), where h is the height of the binary tree. This is due to the recursive calls in the DFS traversal. In the worst case, the space complexity can be O(n) for a skewed tree.

To improve the time complexity, we can modify the algorithm to avoid redundant calculations and achieve a linear time complexity. This can be done by returning both the excess coins and the number of moves from the DFS function. This way, we can eliminate the need to traverse the tree twice.

--------------------------------------------------------------------------------

>>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...
user_proxy (to assistant):

exitcode: 1 (execution failed)
Code output: 
Traceback (most recent call last):
  File "", line 1, in <module>
    class Solution:
  File "", line 2, in Solution
    def distributeCoins(self, root: TreeNode) -> int:
NameError: name 'TreeNode' is not defined

Notice how the UserProxyAgent got an error. And it calls back the other agent for its solution.

--------------------------------------------------------------------------------
assistant (to user_proxy):

It seems like the code snippet I provided is missing the definition of the `TreeNode` class. Here is an updated version of the code that includes the definition of the `TreeNode` class:

```python
# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def distributeCoins(self, root: TreeNode) -> int:
        self.moves = 0
        
        def dfs(node):
            if not node:
                return 0
            
            left_excess = dfs(node.left)
            right_excess = dfs(node.right)
            
            self.moves += abs(left_excess) + abs(right_excess)
            
            return node.val + left_excess + right_excess - 1
        
        dfs(root)
        
        return self.moves
```

You can now use this updated code to solve the LeetCode problem "Distribute Coins in a binary tree".

--------------------------------------------------------------------------------

>>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...
user_proxy (to assistant):

exitcode: 0 (execution succeeded)
Code output: 


--------------------------------------------------------------------------------
assistant (to user_proxy):

Great! I'm glad to hear that the code executed successfully. If you have any more questions or need further assistance, feel free to ask!

--------------------------------------------------------------------------------
user_proxy (to assistant):



--------------------------------------------------------------------------------
assistant (to user_proxy):

If you have any more questions or need help with anything else, feel free to ask. I'm here to help!

--------------------------------------------------------------------------------
user_proxy (to assistant):



--------------------------------------------------------------------------------
assistant (to user_proxy):

If you have any questions in the future or need assistance, don't hesitate to ask. Have a great day!

--------------------------------------------------------------------------------

Conclusion

By using AutoGen’s conversable agents, we:

Automatically verified that the Python code suggested by the LLM works.
Created a framework where the LLM can respond to syntax or logical errors by reading the console output.

Thanks for reading! If you found this article useful, please leave a comment or a clap.
Follow me to stay updated on my latest articles!