docs/examples/agent/code_act_agent.ipynb
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:
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.
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
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 aroundglobals: Global variables to use in the execution context, this includes the builtins and other imported modules we want to use in the execution contextfrom 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"),
},
)
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)
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"
)
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"
)
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"
)
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"
)