cookbooks/notebooks/agent/code_act_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prebuilt CodeAct Agent w/ LlamaIndex\n",
    "\n",
    "LlamaIndex offers a prebuilt CodeAct Agent that can be used to write and execute code, inspired by the original [CodeAct paper](https://arxiv.org/abs/2402.01030).\n",
    "\n",
    "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.\n",
    "\n",
    "Some advantages of using the CodeAct Agent:\n",
    "\n",
    "- No need to exhaustively list out all the possible functions that the agent might need\n",
    "- The agent can develop complex workflows around your existing functions\n",
    "- Can integrate directly with existing API's\n",
    "\n",
    "Let's walk through a simple example of how to use the CodeAct Agent.\n",
    "\n",
    "**NOTE:** This example includes code that will execute arbitrary code. This is dangerous, and proper sandboxing should be used in production environments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's configure the LLM we want to use, and provide some functions that we can use in our code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U llama-index-core llama-index-llms-ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "# Configure the LLM\n",
    "llm = OpenAI(model=\"gpt-4o-mini\", api_key=\"sk-...\")\n",
    "\n",
    "\n",
    "# Define a few helper functions\n",
    "def add(a: int, b: int) -> int:\n",
    "    \"\"\"Add two numbers together\"\"\"\n",
    "    return a + b\n",
    "\n",
    "\n",
    "def subtract(a: int, b: int) -> int:\n",
    "    \"\"\"Subtract two numbers\"\"\"\n",
    "    return a - b\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiply two numbers\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def divide(a: int, b: int) -> float:\n",
    "    \"\"\"Divide two numbers\"\"\"\n",
    "    return a / b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Code Executor \n",
    "\n",
    "The `CodeActAgent` will require a specific `code_execute_fn` to execute the code generated by the agent.\n",
    "\n",
    "Below, we define a simple `code_execute_fn` that will execute the code in-process and maintain execution state.\n",
    "\n",
    "**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.\n",
    "\n",
    "With this executor, we can pass in a dictionary of local and global variables to use in the execution context.\n",
    "\n",
    "- `locals`: Local variables to use in the execution context, this includes our functions that we want the LLM to code around\n",
    "- `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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, Dict, Tuple\n",
    "import io\n",
    "import contextlib\n",
    "import ast\n",
    "import traceback\n",
    "\n",
    "\n",
    "class SimpleCodeExecutor:\n",
    "    \"\"\"\n",
    "    A simple code executor that runs Python code with state persistence.\n",
    "\n",
    "    This executor maintains a global and local state between executions,\n",
    "    allowing for variables to persist across multiple code runs.\n",
    "\n",
    "    NOTE: not safe for production use! Use with caution.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, locals: Dict[str, Any], globals: Dict[str, Any]):\n",
    "        \"\"\"\n",
    "        Initialize the code executor.\n",
    "\n",
    "        Args:\n",
    "            locals: Local variables to use in the execution context\n",
    "            globals: Global variables to use in the execution context\n",
    "        \"\"\"\n",
    "        # State that persists between executions\n",
    "        self.globals = globals\n",
    "        self.locals = locals\n",
    "\n",
    "    def execute(self, code: str) -> Tuple[bool, str, Any]:\n",
    "        \"\"\"\n",
    "        Execute Python code and capture output and return values.\n",
    "\n",
    "        Args:\n",
    "            code: Python code to execute\n",
    "\n",
    "        Returns:\n",
    "            Dict with keys `success`, `output`, and `return_value`\n",
    "        \"\"\"\n",
    "        # Capture stdout and stderr\n",
    "        stdout = io.StringIO()\n",
    "        stderr = io.StringIO()\n",
    "\n",
    "        output = \"\"\n",
    "        return_value = None\n",
    "        try:\n",
    "            # Execute with captured output\n",
    "            with contextlib.redirect_stdout(\n",
    "                stdout\n",
    "            ), contextlib.redirect_stderr(stderr):\n",
    "                # Try to detect if there's a return value (last expression)\n",
    "                try:\n",
    "                    tree = ast.parse(code)\n",
    "                    last_node = tree.body[-1] if tree.body else None\n",
    "\n",
    "                    # If the last statement is an expression, capture its value\n",
    "                    if isinstance(last_node, ast.Expr):\n",
    "                        # Split code to add a return value assignment\n",
    "                        last_line = code.rstrip().split(\"\\n\")[-1]\n",
    "                        exec_code = (\n",
    "                            code[: -len(last_line)]\n",
    "                            + \"\\n__result__ = \"\n",
    "                            + last_line\n",
    "                        )\n",
    "\n",
    "                        # Execute modified code\n",
    "                        exec(exec_code, self.globals, self.locals)\n",
    "                        return_value = self.locals.get(\"__result__\")\n",
    "                    else:\n",
    "                        # Normal execution\n",
    "                        exec(code, self.globals, self.locals)\n",
    "                except:\n",
    "                    # If parsing fails, just execute the code as is\n",
    "                    exec(code, self.globals, self.locals)\n",
    "\n",
    "            # Get output\n",
    "            output = stdout.getvalue()\n",
    "            if stderr.getvalue():\n",
    "                output += \"\\n\" + stderr.getvalue()\n",
    "\n",
    "        except Exception as e:\n",
    "            # Capture exception information\n",
    "            output = f\"Error: {type(e).__name__}: {str(e)}\\n\"\n",
    "            output += traceback.format_exc()\n",
    "\n",
    "        if return_value is not None:\n",
    "            output += \"\\n\\n\" + str(return_value)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "code_executor = SimpleCodeExecutor(\n",
    "    # give access to our functions defined above\n",
    "    locals={\n",
    "        \"add\": add,\n",
    "        \"subtract\": subtract,\n",
    "        \"multiply\": multiply,\n",
    "        \"divide\": divide,\n",
    "    },\n",
    "    globals={\n",
    "        # give access to all builtins\n",
    "        \"__builtins__\": __builtins__,\n",
    "        # give access to numpy\n",
    "        \"np\": __import__(\"numpy\"),\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the CodeAct Agent\n",
    "\n",
    "Now that we have our code executor, we can setup the CodeAct Agent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import CodeActAgent\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "agent = CodeActAgent(\n",
    "    code_execute_fn=code_executor.execute,\n",
    "    llm=llm,\n",
    "    tools=[add, subtract, multiply, divide],\n",
    ")\n",
    "\n",
    "# context to hold the agent's session/state/chat history\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the Agent\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import (\n",
    "    ToolCall,\n",
    "    ToolCallResult,\n",
    "    AgentStream,\n",
    ")\n",
    "\n",
    "\n",
    "async def run_agent_verbose(agent, ctx, query):\n",
    "    handler = agent.run(query, ctx=ctx)\n",
    "    print(f\"User:  {query}\")\n",
    "    async for event in handler.stream_events():\n",
    "        if isinstance(event, ToolCallResult):\n",
    "            print(\n",
    "                f\"\\n-----------\\nCode execution result:\\n{event.tool_output}\"\n",
    "            )\n",
    "        elif isinstance(event, ToolCall):\n",
    "            print(f\"\\n-----------\\nParsed code:\\n{event.tool_kwargs['code']}\")\n",
    "        elif isinstance(event, AgentStream):\n",
    "            print(f\"{event.delta}\", end=\"\", flush=True)\n",
    "\n",
    "    return await handler"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, the agent uses some built-in functions to calculate the sum of all numbers from 1 to 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Calculate the sum of all numbers from 1 to 10\n",
      "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.\n",
      "\n",
      "<execute>\n",
      "# Calculate the sum of numbers from 1 to 10\n",
      "total_sum = sum(range(1, 11))\n",
      "print(total_sum)\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "# Calculate the sum of numbers from 1 to 10\n",
      "total_sum = sum(range(1, 11))\n",
      "print(total_sum)\n",
      "\n",
      "-----------\n",
      "Code execution result:\n",
      "55\n",
      "\n",
      "The sum of all numbers from 1 to 10 is 55."
     ]
    }
   ],
   "source": [
    "response = await run_agent_verbose(\n",
    "    agent, ctx, \"Calculate the sum of all numbers from 1 to 10\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we get the agent to use the tools that we passed in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Add 5 and 3, then multiply the result by 2\n",
      "I will perform the addition of 5 and 3, and then multiply the result by 2.\n",
      "\n",
      "<execute>\n",
      "# Perform the calculation\n",
      "addition_result = add(5, 3)\n",
      "final_result = multiply(addition_result, 2)\n",
      "print(final_result)\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "# Perform the calculation\n",
      "addition_result = add(5, 3)\n",
      "final_result = multiply(addition_result, 2)\n",
      "print(final_result)\n",
      "\n",
      "-----------\n",
      "Code execution result:\n",
      "16\n",
      "\n",
      "The result of adding 5 and 3, then multiplying by 2, is 16."
     ]
    }
   ],
   "source": [
    "response = await run_agent_verbose(\n",
    "    agent, ctx, \"Add 5 and 3, then multiply the result by 2\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can even get the agent to define new functions for us!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Calculate the sum of the first 10 fibonacci numbers\n",
      "I will calculate the sum of the first 10 Fibonacci numbers.\n",
      "\n",
      "<execute>\n",
      "def fibonacci(n):\n",
      "    fib_sequence = [0, 1]\n",
      "    for i in range(2, n):\n",
      "        fib_sequence.append(fib_sequence[-1] + fib_sequence[-2])\n",
      "    return fib_sequence\n",
      "\n",
      "# Calculate the sum of the first 10 Fibonacci numbers\n",
      "first_10_fib = fibonacci(10)\n",
      "fibonacci_sum = sum(first_10_fib)\n",
      "print(fibonacci_sum)\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "def fibonacci(n):\n",
      "    fib_sequence = [0, 1]\n",
      "    for i in range(2, n):\n",
      "        fib_sequence.append(fib_sequence[-1] + fib_sequence[-2])\n",
      "    return fib_sequence\n",
      "\n",
      "# Calculate the sum of the first 10 Fibonacci numbers\n",
      "first_10_fib = fibonacci(10)\n",
      "fibonacci_sum = sum(first_10_fib)\n",
      "print(fibonacci_sum)\n",
      "\n",
      "-----------\n",
      "Code execution result:\n",
      "88\n",
      "\n",
      "The sum of the first 10 Fibonacci numbers is 88."
     ]
    }
   ],
   "source": [
    "response = await run_agent_verbose(\n",
    "    agent, ctx, \"Calculate the sum of the first 10 fibonacci numbers\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then reuse those new functions in a new task!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Calculate the sum of the first 20 fibonacci numbers\n",
      "I will calculate the sum of the first 20 Fibonacci numbers.\n",
      "\n",
      "<execute>\n",
      "# Calculate the sum of the first 20 Fibonacci numbers\n",
      "first_20_fib = fibonacci(20)\n",
      "fibonacci_sum_20 = sum(first_20_fib)\n",
      "print(fibonacci_sum_20)\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "# Calculate the sum of the first 20 Fibonacci numbers\n",
      "first_20_fib = fibonacci(20)\n",
      "fibonacci_sum_20 = sum(first_20_fib)\n",
      "print(fibonacci_sum_20)\n",
      "\n",
      "-----------\n",
      "Code execution result:\n",
      "10945\n",
      "\n",
      "The sum of the first 20 Fibonacci numbers is 10,945."
     ]
    }
   ],
   "source": [
    "response = await run_agent_verbose(\n",
    "    agent, ctx, \"Calculate the sum of the first 20 fibonacci numbers\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-caVs7DDe-py3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}