Prebuilt CodeAct Agent w/ LlamaIndex

LlamaIndex offers a prebuilt CodeAct Agent that can be used to write and execute code, inspired by the original CodeAct paper.

With this agent, you provide an agent with a set of functions, and the agent will write code that uses those functions to help complete the task you give it.

Some advantages of using the CodeAct Agent:

• No need to exhaustively list out all the possible functions that the agent might need
• The agent can develop complex workflows around your existing functions
• Can integrate directly with existing API’s

Let’s walk through a simple example of how to use the CodeAct Agent.

NOTE: This example includes code that will execute arbitrary code. This is dangerous, and proper sandboxing should be used in production environments.

Setup

First, let’s configure the LLM we want to use, and provide some functions that we can use in our code.

%pip install -U llama-index-core llama-index-llms-ollama

from llama_index.llms.openai import OpenAI

# Configure the LLM
llm = OpenAI(model="gpt-4o-mini", api_key="sk-...")


# Define a few helper functions
def add(a: int, b: int) -> int:
    """Add two numbers together"""
    return a + b


def subtract(a: int, b: int) -> int:
    """Subtract two numbers"""
    return a - b


def multiply(a: int, b: int) -> int:
    """Multiply two numbers"""
    return a * b


def divide(a: int, b: int) -> float:
    """Divide two numbers"""
    return a / b

Create a Code Executor

The CodeActAgent will require a specific code_execute_fn to execute the code generated by the agent.

Below, we define a simple code_execute_fn that will execute the code in-process and maintain execution state.

NOTE: In a production environment, you should use a more robust method of executing code. This is just for demonstration purposes, and executing code in-process is dangerous. Consider using docker or external services to execute code.

With this executor, we can pass in a dictionary of local and global variables to use in the execution context.

• locals: Local variables to use in the execution context, this includes our functions that we want the LLM to code around
• globals: Global variables to use in the execution context, this includes the builtins and other imported modules we want to use in the execution context

from typing import Any, Dict, Tuple
import io
import contextlib
import ast
import traceback


class SimpleCodeExecutor:
    """
    A simple code executor that runs Python code with state persistence.

    This executor maintains a global and local state between executions,
    allowing for variables to persist across multiple code runs.

    NOTE: not safe for production use! Use with caution.
    """

    def __init__(self, locals: Dict[str, Any], globals: Dict[str, Any]):
        """
        Initialize the code executor.

        Args:
            locals: Local variables to use in the execution context
            globals: Global variables to use in the execution context
        """
        # State that persists between executions
        self.globals = globals
        self.locals = locals

    def execute(self, code: str) -> Tuple[bool, str, Any]:
        """
        Execute Python code and capture output and return values.

        Args:
            code: Python code to execute

        Returns:
            Dict with keys `success`, `output`, and `return_value`
        """
        # Capture stdout and stderr
        stdout = io.StringIO()
        stderr = io.StringIO()

        output = ""
        return_value = None
        try:
            # Execute with captured output
            with contextlib.redirect_stdout(
                stdout
            ), contextlib.redirect_stderr(stderr):
                # Try to detect if there's a return value (last expression)
                try:
                    tree = ast.parse(code)
                    last_node = tree.body[-1] if tree.body else None

                    # If the last statement is an expression, capture its value
                    if isinstance(last_node, ast.Expr):
                        # Split code to add a return value assignment
                        last_line = code.rstrip().split("\n")[-1]
                        exec_code = (
                            code[: -len(last_line)]
                            + "\n__result__ = "
                            + last_line
                        )

                        # Execute modified code
                        exec(exec_code, self.globals, self.locals)
                        return_value = self.locals.get("__result__")
                    else:
                        # Normal execution
                        exec(code, self.globals, self.locals)
                except:
                    # If parsing fails, just execute the code as is
                    exec(code, self.globals, self.locals)

            # Get output
            output = stdout.getvalue()
            if stderr.getvalue():
                output += "\n" + stderr.getvalue()

        except Exception as e:
            # Capture exception information
            output = f"Error: {type(e).__name__}: {str(e)}\n"
            output += traceback.format_exc()

        if return_value is not None:
            output += "\n\n" + str(return_value)

        return output

code_executor = SimpleCodeExecutor(
    # give access to our functions defined above
    locals={
        "add": add,
        "subtract": subtract,
        "multiply": multiply,
        "divide": divide,
    },
    globals={
        # give access to all builtins
        "__builtins__": __builtins__,
        # give access to numpy
        "np": __import__("numpy"),
    },
)

Setup the CodeAct Agent

Now that we have our code executor, we can setup the CodeAct Agent.

from llama_index.core.agent.workflow import CodeActAgent
from llama_index.core.workflow import Context

agent = CodeActAgent(
    code_execute_fn=code_executor.execute,
    llm=llm,
    tools=[add, subtract, multiply, divide],
)

# context to hold the agent's session/state/chat history
ctx = Context(agent)

Use the Agent

Now that we have our agent, we can use it to complete tasks! Since we gave it some math functions, we will prompt it for tasks that require calculations.

from llama_index.core.agent.workflow import (
    ToolCall,
    ToolCallResult,
    AgentStream,
)


async def run_agent_verbose(agent, ctx, query):
    handler = agent.run(query, ctx=ctx)
    print(f"User:  {query}")
    async for event in handler.stream_events():
        if isinstance(event, ToolCallResult):
            print(
                f"\n-----------\nCode execution result:\n{event.tool_output}"
            )
        elif isinstance(event, ToolCall):
            print(f"\n-----------\nParsed code:\n{event.tool_kwargs['code']}")
        elif isinstance(event, AgentStream):
            print(f"{event.delta}", end="", flush=True)

    return await handler

Here, the agent uses some built-in functions to calculate the sum of all numbers from 1 to 10.

response = await run_agent_verbose(
    agent, ctx, "Calculate the sum of all numbers from 1 to 10"
)

User:  Calculate the sum of all numbers from 1 to 10
The sum of all numbers from 1 to 10 can be calculated using the formula for the sum of an arithmetic series. However, I will compute it directly for you.

<execute>
# Calculate the sum of numbers from 1 to 10
total_sum = sum(range(1, 11))
print(total_sum)
</execute>
-----------
Parsed code:
# Calculate the sum of numbers from 1 to 10
total_sum = sum(range(1, 11))
print(total_sum)

-----------
Code execution result:
55

The sum of all numbers from 1 to 10 is 55.

Next, we get the agent to use the tools that we passed in.

response = await run_agent_verbose(
    agent, ctx, "Add 5 and 3, then multiply the result by 2"
)

User:  Add 5 and 3, then multiply the result by 2
I will perform the addition of 5 and 3, and then multiply the result by 2.

<execute>
# Perform the calculation
addition_result = add(5, 3)
final_result = multiply(addition_result, 2)
print(final_result)
</execute>
-----------
Parsed code:
# Perform the calculation
addition_result = add(5, 3)
final_result = multiply(addition_result, 2)
print(final_result)

-----------
Code execution result:
16

The result of adding 5 and 3, then multiplying by 2, is 16.

We can even get the agent to define new functions for us!

response = await run_agent_verbose(
    agent, ctx, "Calculate the sum of the first 10 fibonacci numbers"
)

User:  Calculate the sum of the first 10 fibonacci numbers
I will calculate the sum of the first 10 Fibonacci numbers.

<execute>
def fibonacci(n):
    fib_sequence = [0, 1]
    for i in range(2, n):
        fib_sequence.append(fib_sequence[-1] + fib_sequence[-2])
    return fib_sequence

# Calculate the sum of the first 10 Fibonacci numbers
first_10_fib = fibonacci(10)
fibonacci_sum = sum(first_10_fib)
print(fibonacci_sum)
</execute>
-----------
Parsed code:
def fibonacci(n):
    fib_sequence = [0, 1]
    for i in range(2, n):
        fib_sequence.append(fib_sequence[-1] + fib_sequence[-2])
    return fib_sequence

# Calculate the sum of the first 10 Fibonacci numbers
first_10_fib = fibonacci(10)
fibonacci_sum = sum(first_10_fib)
print(fibonacci_sum)

-----------
Code execution result:
88

The sum of the first 10 Fibonacci numbers is 88.

And then reuse those new functions in a new task!

response = await run_agent_verbose(
    agent, ctx, "Calculate the sum of the first 20 fibonacci numbers"
)

User:  Calculate the sum of the first 20 fibonacci numbers
I will calculate the sum of the first 20 Fibonacci numbers.

<execute>
# Calculate the sum of the first 20 Fibonacci numbers
first_20_fib = fibonacci(20)
fibonacci_sum_20 = sum(first_20_fib)
print(fibonacci_sum_20)
</execute>
-----------
Parsed code:
# Calculate the sum of the first 20 Fibonacci numbers
first_20_fib = fibonacci(20)
fibonacci_sum_20 = sum(first_20_fib)
print(fibonacci_sum_20)

-----------
Code execution result:
10945

The sum of the first 20 Fibonacci numbers is 10,945.