cookbooks/notebooks/agent/from_scratch_code_act_agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a CodeAct Agent From Scratch\n",
    "\n",
    "While LlamaIndex provides a pre-built [CodeActAgent](https://docs.llamaindex.ai/en/stable/examples/agent/code_act_agent/), we can also create our own from scratch.\n",
    "\n",
    "This way, we can fully understand and customize the agent's behaviour beyond what is provided by the pre-built agent.\n",
    "\n",
    "In this notebook, we will\n",
    "1. Create a workflow for generating and parsing code\n",
    "2. Implement basic code execution\n",
    "3. Add memory and state to the agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up Functions for our Agent\n",
    "\n",
    "If we want our agent to execute our code, we need to deine the code for it to execute!\n",
    "\n",
    "For now, let's use a few basic math functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": [
    "## Creating a Code Executor \n",
    "\n",
    "In order to execute code, we need to create a code executor.\n",
    "\n",
    "Here, we will use a simple in-process code executor that maintains it's own state.\n",
    "\n",
    "**NOTE:** This is a simple example, and does not include proper sandboxing. In a production environment, you should use tools like docker or proper code sandboxing environments."
   ]
  },
  {
   "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": [
    "## Defining the CodeAct Agent\n",
    "\n",
    "Now, we can using LlamaIndex Workflows to define the workflow for our agent.\n",
    "\n",
    "The basic flow is:\n",
    "- take in our prompt + chat history\n",
    "- parse out the code to execute (if any)\n",
    "- execute the code\n",
    "- provide the output of the code execution back to the agent\n",
    "- repeat until the agent is satisfied with the answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we can create the events in the workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.core.workflow import Event\n",
    "\n",
    "\n",
    "class InputEvent(Event):\n",
    "    input: list[ChatMessage]\n",
    "\n",
    "\n",
    "class StreamEvent(Event):\n",
    "    delta: str\n",
    "\n",
    "\n",
    "class CodeExecutionEvent(Event):\n",
    "    code: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can define the workflow that orchestrates using these events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "import re\n",
    "from typing import Any, Callable, List\n",
    "\n",
    "from llama_index.core.llms import ChatMessage, LLM\n",
    "from llama_index.core.memory import ChatMemoryBuffer\n",
    "from llama_index.core.tools.types import BaseTool\n",
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    Workflow,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    step,\n",
    ")\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "\n",
    "CODEACT_SYSTEM_PROMPT = \"\"\"\n",
    "You are a helpful assistant that can execute code.\n",
    "\n",
    "Given the chat history, you can write code within <execute>...</execute> tags to help the user with their question.\n",
    "\n",
    "In your code, you can reference any previously used variables or functions.\n",
    "\n",
    "The user has also provided you with some predefined functions:\n",
    "{fn_str}\n",
    "\n",
    "To execute code, write the code between <execute>...</execute> tags.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class CodeActAgent(Workflow):\n",
    "    def __init__(\n",
    "        self,\n",
    "        fns: List[Callable],\n",
    "        code_execute_fn: Callable,\n",
    "        llm: LLM | None = None,\n",
    "        **workflow_kwargs: Any,\n",
    "    ) -> None:\n",
    "        super().__init__(**workflow_kwargs)\n",
    "        self.fns = fns or []\n",
    "        self.code_execute_fn = code_execute_fn\n",
    "        self.llm = llm or OpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "        # parse the functions into truncated function strings\n",
    "        self.fn_str = \"\\n\\n\".join(\n",
    "            f'def {fn.__name__}{str(inspect.signature(fn))}:\\n    \"\"\" {fn.__doc__} \"\"\"\\n    ...'\n",
    "            for fn in self.fns\n",
    "        )\n",
    "        self.system_message = ChatMessage(\n",
    "            role=\"system\",\n",
    "            content=CODEACT_SYSTEM_PROMPT.format(fn_str=self.fn_str),\n",
    "        )\n",
    "\n",
    "    def _parse_code(self, response: str) -> str | None:\n",
    "        # find the code between <execute>...</execute> tags\n",
    "        matches = re.findall(r\"<execute>(.*?)</execute>\", response, re.DOTALL)\n",
    "        if matches:\n",
    "            return \"\\n\\n\".join(matches)\n",
    "\n",
    "        return None\n",
    "\n",
    "    @step\n",
    "    async def prepare_chat_history(\n",
    "        self, ctx: Context, ev: StartEvent\n",
    "    ) -> InputEvent:\n",
    "        # check if memory is setup\n",
    "        memory = await ctx.store.get(\"memory\", default=None)\n",
    "        if not memory:\n",
    "            memory = ChatMemoryBuffer.from_defaults(llm=self.llm)\n",
    "\n",
    "        # get user input\n",
    "        user_input = ev.get(\"user_input\")\n",
    "        if user_input is None:\n",
    "            raise ValueError(\"user_input kwarg is required\")\n",
    "        user_msg = ChatMessage(role=\"user\", content=user_input)\n",
    "        memory.put(user_msg)\n",
    "\n",
    "        # get chat history\n",
    "        chat_history = memory.get()\n",
    "\n",
    "        # update context\n",
    "        await ctx.store.set(\"memory\", memory)\n",
    "\n",
    "        # add the system message to the chat history and return\n",
    "        return InputEvent(input=[self.system_message, *chat_history])\n",
    "\n",
    "    @step\n",
    "    async def handle_llm_input(\n",
    "        self, ctx: Context, ev: InputEvent\n",
    "    ) -> CodeExecutionEvent | StopEvent:\n",
    "        chat_history = ev.input\n",
    "\n",
    "        # stream the response\n",
    "        response_stream = await self.llm.astream_chat(chat_history)\n",
    "        async for response in response_stream:\n",
    "            ctx.write_event_to_stream(StreamEvent(delta=response.delta or \"\"))\n",
    "\n",
    "        # save the final response, which should have all content\n",
    "        memory = await ctx.store.get(\"memory\")\n",
    "        memory.put(response.message)\n",
    "        await ctx.store.set(\"memory\", memory)\n",
    "\n",
    "        # get the code to execute\n",
    "        code = self._parse_code(response.message.content)\n",
    "\n",
    "        if not code:\n",
    "            return StopEvent(result=response)\n",
    "        else:\n",
    "            return CodeExecutionEvent(code=code)\n",
    "\n",
    "    @step\n",
    "    async def handle_code_execution(\n",
    "        self, ctx: Context, ev: CodeExecutionEvent\n",
    "    ) -> InputEvent:\n",
    "        # execute the code\n",
    "        ctx.write_event_to_stream(ev)\n",
    "        output = self.code_execute_fn(ev.code)\n",
    "\n",
    "        # update the memory\n",
    "        memory = await ctx.store.get(\"memory\")\n",
    "        memory.put(ChatMessage(role=\"assistant\", content=output))\n",
    "        await ctx.store.set(\"memory\", memory)\n",
    "\n",
    "        # get the latest chat history and loop back to the start\n",
    "        chat_history = memory.get()\n",
    "        return InputEvent(input=[self.system_message, *chat_history])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the CodeAct Agent\n",
    "\n",
    "Now, we can test out the CodeAct Agent!\n",
    "\n",
    "We'll create a simple agent and slowly build up the complexity with requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import Context\n",
    "\n",
    "agent = CodeActAgent(\n",
    "    fns=[add, subtract, multiply, divide],\n",
    "    code_execute_fn=code_executor.execute,\n",
    "    llm=OpenAI(model=\"gpt-4o-mini\", api_key=\"sk-...\"),\n",
    ")\n",
    "\n",
    "# context to hold the agent's state / memory\n",
    "ctx = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def run_agent_verbose(agent: CodeActAgent, ctx: Context, query: str):\n",
    "    handler = agent.run(user_input=query, ctx=ctx)\n",
    "    print(f\"User:  {query}\")\n",
    "    async for event in handler.stream_events():\n",
    "        if isinstance(event, StreamEvent):\n",
    "            print(f\"{event.delta}\", end=\"\", flush=True)\n",
    "        elif isinstance(event, CodeExecutionEvent):\n",
    "            print(f\"\\n-----------\\nParsed code:\\n{event.code}\\n\")\n",
    "\n",
    "    return await handler"
   ]
  },
  {
   "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",
      "To calculate the sum of all numbers from 1 to 10, we can use the `add` function in a loop. Here's how we can do it:\n",
      "\n",
      "<execute>\n",
      "total_sum = 0\n",
      "for number in range(1, 11):\n",
      "    total_sum = add(total_sum, number)\n",
      "total_sum\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "\n",
      "total_sum = 0\n",
      "for number in range(1, 11):\n",
      "    total_sum = add(total_sum, number)\n",
      "total_sum\n",
      "\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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Add 5 and 3, then multiply the result by 2\n",
      "To perform the calculation, we will first add 5 and 3 using the `add` function, and then multiply the result by 2 using the `multiply` function. Here's how we can do it:\n",
      "\n",
      "<execute>\n",
      "result_addition = add(5, 3)\n",
      "final_result = multiply(result_addition, 2)\n",
      "final_result\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "\n",
      "result_addition = add(5, 3)\n",
      "final_result = multiply(result_addition, 2)\n",
      "final_result\n",
      "\n",
      "\n",
      "The final 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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Calculate the sum of the first 10 fibonacci numbers0\n",
      "To calculate the sum of the first 10 Fibonacci numbers, we first need to generate the Fibonacci sequence up to the 10th number and then sum those numbers. The Fibonacci sequence starts with 0 and 1, and each subsequent number is the sum of the two preceding ones.\n",
      "\n",
      "Here's how we can do it:\n",
      "\n",
      "<execute>\n",
      "def fibonacci(n: int) -> int:\n",
      "    \"\"\" Return the nth Fibonacci number \"\"\"\n",
      "    if n == 0:\n",
      "        return 0\n",
      "    elif n == 1:\n",
      "        return 1\n",
      "    else:\n",
      "        a, b = 0, 1\n",
      "        for _ in range(2, n + 1):\n",
      "            a, b = b, a + b\n",
      "        return b\n",
      "\n",
      "# Calculate the sum of the first 10 Fibonacci numbers\n",
      "fibonacci_sum = 0\n",
      "for i in range(10):\n",
      "    fibonacci_sum = add(fibonacci_sum, fibonacci(i))\n",
      "\n",
      "fibonacci_sum\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "\n",
      "def fibonacci(n: int) -> int:\n",
      "    \"\"\" Return the nth Fibonacci number \"\"\"\n",
      "    if n == 0:\n",
      "        return 0\n",
      "    elif n == 1:\n",
      "        return 1\n",
      "    else:\n",
      "        a, b = 0, 1\n",
      "        for _ in range(2, n + 1):\n",
      "            a, b = b, a + b\n",
      "        return b\n",
      "\n",
      "# Calculate the sum of the first 10 Fibonacci numbers\n",
      "fibonacci_sum = 0\n",
      "for i in range(10):\n",
      "    fibonacci_sum = add(fibonacci_sum, fibonacci(i))\n",
      "\n",
      "fibonacci_sum\n",
      "\n",
      "\n",
      "The sum of the first 10 Fibonacci numbers is 55."
     ]
    }
   ],
   "source": [
    "response = await run_agent_verbose(\n",
    "    agent, ctx, \"Calculate the sum of the first 10 fibonacci numbers0\"\n",
    ")"
   ]
  },
  {
   "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",
      "To calculate the sum of the first 20 Fibonacci numbers, we can use the same approach as before, but this time we will iterate up to 20. Here's how we can do it:\n",
      "\n",
      "<execute>\n",
      "# Calculate the sum of the first 20 Fibonacci numbers\n",
      "fibonacci_sum_20 = 0\n",
      "for i in range(20):\n",
      "    fibonacci_sum_20 = add(fibonacci_sum_20, fibonacci(i))\n",
      "\n",
      "fibonacci_sum_20\n",
      "</execute>\n",
      "-----------\n",
      "Parsed code:\n",
      "\n",
      "# Calculate the sum of the first 20 Fibonacci numbers\n",
      "fibonacci_sum_20 = 0\n",
      "for i in range(20):\n",
      "    fibonacci_sum_20 = add(fibonacci_sum_20, fibonacci(i))\n",
      "\n",
      "fibonacci_sum_20\n",
      "\n",
      "\n",
      "The sum of the first 20 Fibonacci numbers is 6765."
     ]
    }
   ],
   "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
}