Tutorial 2 - Shared Variables and Global Context.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2fa1e388-b8f4-4fa3-be5a-ba7c6caae038",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Tutorial 2 - Shared Variables and Global Context\n",
    "\n",
    "- Video: https://www.youtube.com/watch?v=-oo5oi7KjhQ\n",
    "\n",
    "- This is a very important feature of AgentJo, as it helps with learning through tasks and dynamic adaptation of the Agent with changing environments\n",
    "- `Shared Variables` help to keep track of persistent states for the agent\n",
    "- Global Context can provide some of the `Shared Variables` as a global context to the Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "745333b4",
   "metadata": {},
   "source": [
    "# Setup Guide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9843e04a",
   "metadata": {},
   "source": [
    "## Step 1: Install AgentJo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4a16f454",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install agentjo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a14492d",
   "metadata": {},
   "source": [
    "## Step 2: Import required functions and setup relevant API keys for your LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a7228b2e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set up API key and do the necessary imports\n",
    "from agentjo import *\n",
    "import os\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e41f3d8",
   "metadata": {},
   "source": [
    "## Step 3: Define your own LLM\n",
    "- Take in a `system_prompt`, `user_prompt`, and outputs llm response string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8283f0aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def llm(system_prompt: str, user_prompt: str) -> str:\n",
    "    ''' Here, we use OpenAI for illustration, you can change it to your own LLM '''\n",
    "    # ensure your LLM imports are all within this function\n",
    "    from openai import OpenAI\n",
    "    \n",
    "    # define your own LLM here\n",
    "    client = OpenAI()\n",
    "    response = client.chat.completions.create(\n",
    "        model='gpt-4o-mini',\n",
    "        temperature = 0,\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": user_prompt}\n",
    "        ]\n",
    "    )\n",
    "    return response.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "171092c2-b30e-4ea7-9abf-32405b053b91",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Neutral'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Verify that llm function is working\n",
    "llm(system_prompt = 'You are a classifier to classify the sentiment of a sentence', \n",
    "    user_prompt = 'It is a hot and sunny day')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b87cb375-0d69-4cf2-9b46-4404bfadcc04",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 1. Shared Variables\n",
    "\n",
    "*\"Because text is not enough\"* - Anonymous\n",
    "\n",
    "- `shared_variables` is a dictionary, that is initialised in Agent (default empty dictionary), and can be referenced by any function of the agent (including Inner Agents and their functions)\n",
    "- This can be useful for non-text modalitiies (e.g. audio, pdfs, image) and lengthy text modalities, which we do not want to output into `subtasks_completed` directly\n",
    "- To use, simply define an External Function with `shared_variables` as the first input variable, from which you can access and modify `shared_variables` directly\n",
    "- The agent will also be able to be self-referenced in the External Function via `shared_variables['agent']`, so you can change the agent's internal parameters via `shared_variables`\n",
    "- If the function has no output because the output is stored in `shared_variables`, the default return value will be `{'Status': 'Completed'}`\n",
    "\n",
    "### Example External Function using `shared_variables`\n",
    "```python\n",
    "# Use shared_variables as input to your external function to access and modify the shared variables\n",
    "def generate_quotes(shared_variables, number_of_quotes: int, category: str):\n",
    "    ''' Generates number_of_quotes quotes about category '''\n",
    "    # Retrieve from shared variables\n",
    "    my_quote_list = shared_variables['Quote List']\n",
    "    \n",
    "    # Generate the quotes\n",
    "    res = strict_json(system_prompt = f'''Generate {number_of_quotes} sentences about {category}. \n",
    "Do them in the format \"<Quote> - <Person>\", e.g. \"The way to get started is to quit talking and begin doing. - Walt Disney\"\n",
    "Ensure your quotes contain only ' within the quote, and are enclosed by \" ''',\n",
    "                      user_prompt = '',\n",
    "                      output_format = {'Quote List': f'list of {number_of_quotes} quotes, type: List[str]'},\n",
    "                      llm = llm)\n",
    "    \n",
    "    my_quote_list.extend([f'Category: {category}. '+ x for x in res['Quote List']])\n",
    "    \n",
    "    # Store back to shared variables\n",
    "    shared_variables['Quote List'] = my_quote_list\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8cbb3cc7-3532-4894-888d-542991ee3333",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Use shared_variables as input to your external function to access and modify the shared variables\n",
    "def generate_quotes(shared_variables, number_of_quotes: int, category: str):\n",
    "    ''' Generates number_of_quotes quotes about category '''\n",
    "    # Retrieve from shared variables\n",
    "    my_quote_list = shared_variables['Quote List']\n",
    "    \n",
    "    # Generate the quotes\n",
    "    res = strict_json(system_prompt = f'''Generate {number_of_quotes} sentences about {category}. \n",
    "Do them in the format \"<Quote> - <Person>\", e.g. \"The way to get started is to quit talking and begin doing. - Walt Disney\"\n",
    "Ensure your quotes contain only ' within the quote, and are enclosed by \" ''',\n",
    "                      user_prompt = '',\n",
    "                      output_format = {'Quote List': f'list of {number_of_quotes} quotes, type: List[str]'},\n",
    "                      llm = llm)\n",
    "    \n",
    "    my_quote_list.extend([f'Category: {category}. '+ x for x in res['Quote List']])\n",
    "    \n",
    "    # Store back to shared variables\n",
    "    shared_variables['Quote List'] = my_quote_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fca7eb4b-bc07-4587-9d7c-5eec0fd448c9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Define the quote generator agent and the shared_variables - Note the naming convention of s_ at the start of the names for shared variables\n",
    "my_agent = Agent('Quote Generator', 'Generates Quotes according to category',\n",
    "                 default_to_llm = False, # do not provide llm as a default function to Agent to prevent hallucinations\n",
    "                 shared_variables = {'Quote List': []},\n",
    "                 llm = llm).assign_functions([generate_quotes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5a4a96d3-2580-4989-8e3b-f240d1f12190",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet for the assigned task of generating three quotes about life.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the assigned task, I need to generate three quotes specifically about the category of life.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Generate three quotes about life using the equipped function.\u001b[0m\n",
      "Calling function generate_quotes with parameters {'number_of_quotes': 3, 'category': 'life'}\n",
      "> {'Status': 'Completed'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The subtask to generate three quotes about life has been completed successfully.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since the quotes have been generated, the next step is to pass the final output to the user.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "output = my_agent.run('Generate three quotes about life')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "057c7ecb-3c59-4273-80d0-1ac1590b5f1c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quote List: [\"Category: life. 'Life is what happens when you're busy making other plans.' - John Lennon\", \"Category: life. 'In the end, it's not the years in your life that count. It's the life in your years.' - Abraham Lincoln\", \"Category: life. 'Life is either a daring adventure or nothing at all.' - Helen Keller\"]\n"
     ]
    }
   ],
   "source": [
    "# visualise quote list\n",
    "print('Quote List:', my_agent.shared_variables['Quote List'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f28a9ef3-408b-4b63-b725-bbbd5ac10909",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet for the assigned task of generating three quotes about happiness.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the assigned task, I need to generate three quotes specifically focused on the theme of happiness.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Use the generate_quotes function to create three quotes about happiness.\u001b[0m\n",
      "Calling function generate_quotes with parameters {'number_of_quotes': 3, 'category': 'happiness'}\n",
      "> {'Status': 'Completed'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The subtask to generate three quotes about happiness has been completed successfully.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since the quotes have been generated, the next step is to pass the final output to the user.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "my_agent.reset() # always reset agent if the task is something new to prevent misinterpretation\n",
    "output = my_agent.run('Generate three quotes about happiness')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d9e9bc5a-4d6a-41bd-aaac-b7638594ac0a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Quote List: [\"Category: life. 'Life is what happens when you're busy making other plans.' - John Lennon\", \"Category: life. 'In the end, it's not the years in your life that count. It's the life in your years.' - Abraham Lincoln\", \"Category: life. 'Life is either a daring adventure or nothing at all.' - Helen Keller\", \"Category: happiness. 'Happiness is not something ready made. It comes from your own actions.' - Dalai Lama\", \"Category: happiness. 'The purpose of our lives is to be happy.' - Dalai Lama\", \"Category: happiness. 'For every minute you are angry you lose sixty seconds of happiness.' - Ralph Waldo Emerson\"]\n"
     ]
    }
   ],
   "source": [
    "# visualise quote list\n",
    "print('Quote List:', my_agent.shared_variables['Quote List'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa83ba4c-81bf-457b-b1f9-21254aa9f6f4",
   "metadata": {},
   "source": [
    "# 2. Global Context\n",
    "\n",
    "- `Global Context` is a very powerful feature in AgentJo, as it allows the Agent to be updated with the latest environmental state before every decision it makes\n",
    "- It also allows for learnings in `shared_variables` to be carried across tasks, making the Agent teachable and learn through experiences\n",
    "- A recommended practice is to always store the learnings of the Agent during the External Function call, and reset the Agent after each task, so that `subtasks_completed` will be as short as possible to avoid confusion to the Agent\n",
    "\n",
    "- There are two ways to use `Global Context`, and both can be used concurrently:\n",
    "    - 1. `global_context`\n",
    "        - If all you need in the global context is `shared_variables` without any modification to it, then you can use `global_context`\n",
    "        - `global_context` is a string with `<shared_variables_name>` enclosed with `<>`. These <> will be replaced with the actual variable in `shared_variables`\n",
    "    - 2. `get_global_context` \n",
    "        - `get_global_context` is a function that takes in the agent's internal parameters (self) and outputs a string to the LLM to append to the prompts of any LLM-based calls internally, e.g. `get_next_subtask`, `use_llm`, `reply_to_user`\n",
    "    - You have full flexibility to access anything the agent knows and process the `shared_variables` as required and configure a global prompt to the agent\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd4bf6c-08bd-4bfe-95f3-7fd96bda67a4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Example for `global_context` : Inventory Manager\n",
    "- We can use `Global Context` to keep track of inventory state\n",
    "- We simply get the functions `add_item_to_inventory` and `remove_item_from_inventory` to modify the `shared_variable` named `Inventory`\n",
    "- Note we can also put rule-based checks like checking if item is in inventory before removing inside the function\n",
    "- Even after task reset, the Agent still knows the inventory because of `Global Context`\n",
    "\n",
    "```python\n",
    "def add_item_to_inventory(shared_variables, item: str) -> str:\n",
    "    ''' Adds item to inventory, and returns outcome of action '''\n",
    "    shared_variables['Inventory'].append(item)\n",
    "    return f'{item} successfully added to Inventory'\n",
    "    \n",
    "def remove_item_from_inventory(shared_variables, item: str) -> str:\n",
    "    ''' Removes item from inventory and returns outcome of action '''\n",
    "    if item in shared_variables['Inventory']:\n",
    "        shared_variables['Inventory'].remove(item)\n",
    "        return f'{item} successfully removed from Inventory'\n",
    "    else:\n",
    "        return f'{item} not found in Inventory, unable to remove'\n",
    "    \n",
    "agent = Agent('Inventory Manager', \n",
    "              'Adds and removes items in Inventory. Only able to remove items if present in Inventory',\n",
    "              shared_variables = {'Inventory': []},\n",
    "              global_context = 'Inventory: <Inventory>', # Add in Global Context here with shared_variables Inventory\n",
    "              llm = llm).assign_functions([add_item_to_inventory, remove_item_from_inventory])\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9c7c9bb3-288e-4034-896f-690be5984872",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_item_to_inventory(shared_variables, item: str) -> str:\n",
    "    ''' Adds item to inventory, and returns outcome of action '''\n",
    "    shared_variables['Inventory'].append(item)\n",
    "    return f'{item} successfully added to Inventory'\n",
    "    \n",
    "def remove_item_from_inventory(shared_variables, item: str) -> str:\n",
    "    ''' Removes item from inventory and returns outcome of action '''\n",
    "    if item in shared_variables['Inventory']:\n",
    "        shared_variables['Inventory'].remove(item)\n",
    "        return f'{item} successfully removed from Inventory'\n",
    "    else:\n",
    "        return f'{item} not found in Inventory, unable to remove'\n",
    "    \n",
    "agent = Agent('Inventory Manager', \n",
    "              'Adds and removes items in Inventory. Only able to remove items if present in Inventory',\n",
    "              shared_variables = {'Inventory': []},\n",
    "              global_context = 'Inventory: <Inventory>', # Add in Global Context here with shared_variables Inventory\n",
    "              llm = llm).assign_functions([add_item_to_inventory, remove_item_from_inventory])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "95303279-099f-488d-b0d9-4e5e5937b6c8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No items have been added to the inventory yet, and the inventory is currently empty.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the assigned task of adding apples and oranges, I need to add each item to the inventory one at a time.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Add apples to the inventory.\u001b[0m\n",
      "Calling function add_item_to_inventory with parameters {'item': 'apples'}\n",
      "> {'output_1': 'apples successfully added to Inventory'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The task to add apples has been completed successfully, but oranges have not yet been added to the inventory.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the remainder of the assigned task, I need to add oranges to the inventory since apples are already present.\u001b[0m\n",
      "An exception occurred: list index out of range\n",
      "Current invalid json format: { '###item###':}\n",
      "\u001b[1m\u001b[34mSubtask identified: Add oranges to the inventory.\u001b[0m\n",
      "Calling function add_item_to_inventory with parameters {'item': 'oranges'}\n",
      "> {'output_1': 'oranges successfully added to Inventory'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: Both apples and oranges have been successfully added to the inventory.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since the task is to add apples and oranges, and both have already been added, the task is complete.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "output = agent.run('Add apples and oranges')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9c2f1a8e-99b6-4611-b78c-c9ca85034f48",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['apples', 'oranges']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visualise the inventory - we should have both apples and oranges\n",
    "agent.shared_variables['Inventory']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2f61f1cc-fd22-4b82-a891-39eae8dc3003",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: The task is to remove apples from the inventory, which currently contains apples and oranges. No subtasks have been completed yet.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since apples are present in the inventory, I can proceed to remove them using the appropriate function.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Remove apples from the inventory by using the remove_item_from_inventory function.\u001b[0m\n",
      "Calling function remove_item_from_inventory with parameters {'item': 'apples'}\n",
      "> {'output_1': 'apples successfully removed from Inventory'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The task to remove apples has been attempted, but apples are not present in the inventory, which currently only contains oranges.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since apples are not in the inventory, I cannot remove them. The task cannot be completed as there are no apples to remove.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# even after agent reset, we will still know the Inventory because of global context\n",
    "agent.reset()\n",
    "output = agent.run('Remove only apples')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5ea3c3b0-2dd6-4260-a682-260ec3687c24",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['oranges']"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visualise the inventory - we should only have oranges\n",
    "agent.shared_variables['Inventory']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "79f15971-d755-4fbf-99cd-ca6e42caadfb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No items have been removed from the inventory yet, and the only item present is oranges.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since pears are not in the inventory, I cannot remove them. The task cannot be completed as there are no pears to remove.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# even after agent reset, we will still know the Inventory because of global context\n",
    "# Here we will know that pears are not part of inventory and cannot be removed\n",
    "agent.reset()\n",
    "output = agent.run('Remove only pears')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "751e16cb-910c-4c0c-99ed-b609ab9d7e51",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['oranges']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visualise the inventory - we should only have oranges\n",
    "agent.shared_variables['Inventory']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "13e727ff-458e-43ad-8faa-2a76cefbc798",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'Observation': 'No items have been removed from the inventory yet, and the only item present is oranges.',\n",
       "  'Thoughts': 'Since pears are not in the inventory, I cannot remove them. The task cannot be completed as there are no pears to remove.',\n",
       "  'Current Subtask': 'End the task since the requested item to remove is not present in the inventory.',\n",
       "  'Equipped Function Name': 'end_task',\n",
       "  'Equipped Function Inputs': {}}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.thoughts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f4db221-8b94-49f5-bd61-b1bcd1266396",
   "metadata": {},
   "source": [
    "# Advanced Example for `get_global_context`: Maze Navigator\n",
    "- We can use the global context to let agent know present state in a 2D grid, and obstacle positions that the agent has seen\n",
    "- We use the `get_global_context` function as we want to configure more details about how we use the `shared_variables`\n",
    "- Task: Given current position and end position, navigate to end position using actions: Up, Down, Left, Right, Stay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "11ccbceb-9f47-4225-aa08-f36400810f3c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# These are the utility functions\n",
    "def generate_grid(size, start_pos, exit_pos, obstacles):\n",
    "    '''Generates a grid with obstacles'''\n",
    "    grid = [[' ' for _ in range(size)] for _ in range(size)]\n",
    "    grid[start_pos[0]][start_pos[1]] = 'S'  # Start\n",
    "    grid[exit_pos[0]][exit_pos[1]] = 'E'  # Exit\n",
    "    \n",
    "    # Mark a basic path (optional, for simplicity in ensuring a path)\n",
    "    # This part could be removed or replaced with a more sophisticated path marking\n",
    "    path = set()\n",
    "    for i in range(min(start_pos[0], exit_pos[0]), max(start_pos[0], exit_pos[0]) + 1):\n",
    "        path.add((i, start_pos[1]))\n",
    "    for j in range(min(start_pos[1], exit_pos[1]), max(start_pos[1], exit_pos[1]) + 1):\n",
    "        path.add((exit_pos[0], j))\n",
    "    \n",
    "    # Randomly add obstacles\n",
    "    count = 0\n",
    "    while count < obstacles:\n",
    "        r, c = random.randint(0, size-1), random.randint(0, size-1)\n",
    "        if (r, c) not in path and grid[r][c] != 'O' and (r, c) != start_pos and (r, c) != exit_pos:\n",
    "            grid[r][c] = 'O'\n",
    "            count += 1\n",
    "            \n",
    "    return grid\n",
    "\n",
    "def print_grid(grid):\n",
    "    '''Prints the grid'''\n",
    "    for row in grid:\n",
    "        print(' '.join(row))\n",
    "        \n",
    "def check_valid_moves(cur_pos, grid, grid_size):\n",
    "    '''Checks the valid moves in the grid given the cur_pos and grid_size. Returns list of valid moves within action space of Up, Down, Left, Right, Stay'''\n",
    "    row, col = cur_pos\n",
    "    mapping = {'Up': (-1, 0), 'Down': (1, 0), 'Left': (0, -1), 'Right': (0, 1), 'Stay': (0, 0)}\n",
    "    valid_moves = []\n",
    "    for key, value in mapping.items():\n",
    "        row_offset, col_offset = value\n",
    "        # check if valid move\n",
    "        if 0 <= row+row_offset < grid_size and 0 <= col+col_offset < grid_size and grid[row+row_offset][col+col_offset] != 'O':\n",
    "            valid_moves.append(key)\n",
    "    return valid_moves\n",
    "\n",
    "def update_obstacles(cur_pos, grid, grid_size, known_obstacle_pos):\n",
    "    '''Returns the updated known obstacle positions in the current grid given the cur_pos and grid_size'''\n",
    "    row, col = cur_pos\n",
    "    mapping = {'Up': (-1, 0), 'Down': (1, 0), 'Left': (0, -1), 'Right': (0, 1), 'Stay': (0, 0)}\n",
    "    for key, value in mapping.items():\n",
    "        row_offset, col_offset = value\n",
    "        next_row, next_col = row+row_offset, col+col_offset\n",
    "        # check if valid move\n",
    "        if 0 <= next_row < grid_size and 0 <= next_col < grid_size:\n",
    "            # adds in obstacle if observed\n",
    "            if grid[next_row][next_col] == 'O' and (next_row, next_col) not in known_obstacle_pos:\n",
    "                known_obstacle_pos.append((next_row, next_col))\n",
    "            # remove obstacle that is not observed\n",
    "            elif grid[next_row][next_col] != 'O' and (next_row, next_col) in known_obstacle_pos:\n",
    "                known_obstacle_pos.remove((next_row, next_col))\n",
    "    return known_obstacle_pos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b38b44a2-f11d-4aca-9ca6-43bb64469ae8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def move(shared_variables, action: str):\n",
    "    ''' Moves the agent according to the action and outputs outcome of action '''\n",
    "    if action not in shared_variables[\"next_valid_moves\"]: \n",
    "        # if next move is not valid, reflect to agent\n",
    "        return f'The current action of {action} is not valid. You must choose one action from {shared_variables[\"next_valid_moves\"]}'\n",
    "    mapping = {'Up': (-1, 0), 'Down': (1, 0), 'Left': (0, -1), 'Right': (0, 1), 'Stay': (0, 0)}\n",
    "    \n",
    "    # Retrieve from shared variables\n",
    "    row, col = shared_variables[\"cur_pos\"]\n",
    "    grid = shared_variables[\"grid\"]\n",
    "    grid_size = shared_variables[\"grid_size\"]\n",
    "    known_obstacle_pos = shared_variables[\"known_obstacle_pos\"]\n",
    "    \n",
    "    # Do processing for the next action\n",
    "    row_offset, col_offset = mapping[action]\n",
    "    newpos = (row+row_offset, col+col_offset)\n",
    "    next_valid_moves = check_valid_moves(newpos, grid, grid_size)\n",
    "    known_obstacle_pos = update_obstacles(newpos, grid, grid_size, known_obstacle_pos)\n",
    "    \n",
    "    # shift the current agent position\n",
    "    grid[row][col] = ' '\n",
    "    grid[row+row_offset][col+col_offset] = 'S'\n",
    "    \n",
    "    # Store back into shared variables\n",
    "    shared_variables[\"cur_pos\"] = newpos\n",
    "    shared_variables[\"next_valid_moves\"] = next_valid_moves\n",
    "    shared_variables[\"known_obstacle_pos\"] = known_obstacle_pos\n",
    "    shared_variables[\"past_grid_states\"].append(newpos)\n",
    "    shared_variables[\"grid\"] = grid\n",
    "    \n",
    "    print_grid(grid)\n",
    "    \n",
    "    return f'Action successful. Agent moved from {(row, col)} to {newpos}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "050e8c7c-a667-4cae-a9cc-f1b7ef3ade30",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_global_context(agent):\n",
    "    ''' Outputs additional information to the agent '''\n",
    "    \n",
    "    # process additional context based on shared variables (this is what is called persistent variables - variables that will be updated each step)\n",
    "    global_context = f'''Agent position in grid: {agent.shared_variables[\"cur_pos\"]}\n",
    "Exit Position: {agent.shared_variables[\"exit_pos\"]}\n",
    "Last 10 Visited Grid Positions: {agent.shared_variables[\"past_grid_states\"][:10]}\n",
    "Known Obstacle Positions: {agent.shared_variables[\"known_obstacle_pos\"]}\n",
    "Next Valid Moves: {agent.shared_variables[\"next_valid_moves\"]}'''\n",
    "    \n",
    "    # you can also influence how the planner performs the plan with additional details\n",
    "    global_context += '''\n",
    "Try to specify the action in the Subtask\n",
    "Example Assigned Task: ```Navigate to (1, 1)```\n",
    "Example Subtasks taking one action at a time: 'Move Down', 'Move Right' ```\n",
    "'''\n",
    "    return global_context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9a8ec05f-518c-4e0f-a81d-70c2b5a6f6f6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import random\n",
    "# Customise your grid here\n",
    "# O means obstacles, S means start, E means end, blank means nothing there\n",
    "grid_size = 5  # Grid size\n",
    "start_pos = (random.randint(0, grid_size//2 - 1), random.randint(0, grid_size//2 - 1))  # Starting position\n",
    "exit_pos = (random.randint(grid_size//2, grid_size-1), random.randint(grid_size//2, grid_size-1))  # Exit position\n",
    "num_obstacles = 5  # Number of obstacles\n",
    "\n",
    "grid = generate_grid(grid_size, start_pos, exit_pos, num_obstacles)\n",
    "valid_moves = check_valid_moves(start_pos, grid, grid_size)\n",
    "\n",
    "# Assign your agent\n",
    "my_agent = Agent('Maze Navigator', \n",
    "                 f'''You are to move to the Exit Position of the maze. Task is completed when Agent's position is at Exit Position.\n",
    "At each step, you must output an action - Up, Down, Left, Right or Stay.\n",
    "From current position (row, col), this is what you end up with after doing actions:\n",
    "{{'Up': (row-1, col), 'Down': (row+1, col), 'Left': (row, col-1), 'Right': (row, col+1), 'Stay': (row, col)}}\n",
    "You can only move to cells without obstacles. Only take an action from Current Valid Actions.\n",
    "If your previous action is invalid, choose another action.\n",
    "Top left of grid is (0, 0), bottom right is {(grid_size, grid_size)}.''', \n",
    "                 shared_variables = {\n",
    "                    \"cur_pos\": start_pos,\n",
    "                    \"exit_pos\": exit_pos,\n",
    "                    \"past_grid_states\": [],\n",
    "                    \"next_valid_moves\": valid_moves,\n",
    "                    \"known_obstacle_pos\": [],\n",
    "                    \"grid_size\": grid_size,\n",
    "                    \"grid\": grid}, \n",
    "                 max_subtasks = 20,\n",
    "                 get_global_context = get_global_context, # let Agent know the persistent states\n",
    "                 default_to_llm = False,\n",
    "                 llm = llm).assign_functions([move])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e81f5f35-535f-4087-af10-64f5daebc462",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Starting Grid ###\n",
      "    O    \n",
      "  S      \n",
      "O     E  \n",
      "        O\n",
      "O O      \n",
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet, and the agent is currently at position (1, 1) with the exit at (2, 3).\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the exit at (2, 3), the agent can move down to (2, 1) or right to (1, 2). Both moves are valid and will bring the agent closer to the exit.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Down to (2, 1)\u001b[0m\n",
      "Calling function move with parameters {'action': 'Down'}\n",
      "    O    \n",
      "         \n",
      "O S   E  \n",
      "        O\n",
      "O O      \n",
      "> {'output_1': 'Action successful. Agent moved from (1, 1) to (2, 1)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The agent has successfully moved to position (2, 1) from (1, 1) by moving Down. The next target position is (2, 3).\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the Exit Position at (2, 3), the agent can move Right to (2, 2) and then again Right to (2, 3). The next valid move is Right, which will bring the agent closer to the exit.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Right to position (2, 2).\u001b[0m\n",
      "Calling function move with parameters {'action': 'Right'}\n",
      "    O    \n",
      "         \n",
      "O   S E  \n",
      "        O\n",
      "O O      \n",
      "> {'output_1': 'Action successful. Agent moved from (2, 1) to (2, 2)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The agent has successfully moved from (2, 1) to (2, 2) by moving right. The next step is to move to the exit position at (2, 3).\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: The agent can move right to reach the exit position directly since it is adjacent to the current position. This is a valid action and will complete the assigned task.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Right to (2, 3)\u001b[0m\n",
      "Calling function move with parameters {'action': 'Right'}\n",
      "    O    \n",
      "         \n",
      "O     S  \n",
      "        O\n",
      "O O      \n",
      "> {'output_1': 'Action successful. Agent moved from (2, 2) to (2, 3)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The agent has successfully moved to the Exit Position (2, 3) after completing the necessary actions.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since the agent is already at the Exit Position, there is no further action needed to complete the Assigned Task.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('### Starting Grid ###')\n",
    "print_grid(grid)\n",
    "output = my_agent.run(f'Navigate to {exit_pos}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "aed0df04-1d11-46e7-9068-18f3a1ddc513",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent Name: Maze Navigator\n",
      "Agent Description: You are to move to the Exit Position of the maze. Task is completed when Agent's position is at Exit Position.\n",
      "At each step, you must output an action - Up, Down, Left, Right or Stay.\n",
      "From current position (row, col), this is what you end up with after doing actions:\n",
      "{'Up': (row-1, col), 'Down': (row+1, col), 'Left': (row, col-1), 'Right': (row, col+1), 'Stay': (row, col)}\n",
      "You can only move to cells without obstacles. Only take an action from Current Valid Actions.\n",
      "If your previous action is invalid, choose another action.\n",
      "Top left of grid is (0, 0), bottom right is (5, 5).\n",
      "Available Functions: ['end_task', 'move']\n",
      "Shared Variables: ['cur_pos', 'exit_pos', 'past_grid_states', 'next_valid_moves', 'known_obstacle_pos', 'grid_size', 'grid', 'agent']\n",
      "\u001b[1m\u001b[32mTask: Navigate to (2, 3)\u001b[0m\n",
      "\u001b[1m\u001b[30mSubtasks Completed:\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask: move(action=\"Down\")\u001b[0m\n",
      "{'output_1': 'Action successful. Agent moved from (1, 1) to (2, 1)'}\n",
      "\n",
      "\u001b[1m\u001b[34mSubtask: move(action=\"Right\")\u001b[0m\n",
      "{'output_1': 'Action successful. Agent moved from (2, 1) to (2, 2)'}\n",
      "\n",
      "\u001b[1m\u001b[34mSubtask: move(action=\"Right\")(2)\u001b[0m\n",
      "{'output_1': 'Action successful. Agent moved from (2, 2) to (2, 3)'}\n",
      "\n",
      "Is Task Completed: True\n"
     ]
    }
   ],
   "source": [
    "my_agent.status()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ad5c4c68-c820-4e1d-90f3-e7ac9686ffab",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The assigned task is to navigate to the Exit Position at (2, 3). The agent is currently at (2, 3), which is the same as the Exit Position. Therefore, the task is already completed. The last three actions taken were: 1. Move Down from (1, 1) to (2, 1), 2. Move Right from (2, 1) to (2, 2), and 3. Move Right from (2, 2) to (2, 3). Since the agent is already at the Exit Position, no further actions are needed.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The assigned task is to navigate to the Exit Position at (2, 3). The agent is currently at (2, 3), which is the same as the Exit Position. Therefore, the task is already completed. The last three actions taken were: 1. Move Down from (1, 1) to (2, 1), 2. Move Right from (2, 1) to (2, 2), and 3. Move Right from (2, 2) to (2, 3). Since the agent is already at the Exit Position, no further actions are needed.'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The reply to user should include everything from start position to the final position\n",
    "my_agent.reply_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef2de60-dee8-4d8b-911f-158055cefbb4",
   "metadata": {},
   "source": [
    "## Recommended Practice: Reset Subtasks Completed and use Global Context to convey information\n",
    "\n",
    "- If you have multiple similar subtask names, then it is likely the Agent can be confused and think it has already done the subtask\n",
    "- In this case, you can disambiguate by resetting the agent and store the persistent information in `shared_variables` and provide it to the agent using `get_global_context`\n",
    "- Has the benefit of shifting the Start State closer to End State desired by resetting the Agent's planning cycle\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "23da50cb-4134-45f3-90ec-30fddb91e113",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_global_context(agent):\n",
    "    ''' Outputs additional information to the agent '''\n",
    "    \n",
    "    # process additional context based on shared variables (this is what is called persistent variables - variables that will be updated each step)\n",
    "    global_context = f'''Agent position in grid: {agent.shared_variables[\"cur_pos\"]}\n",
    "Exit Position: {agent.shared_variables[\"exit_pos\"]}\n",
    "Last 10 Visited Grid Positions: {agent.shared_variables[\"past_grid_states\"][:10]}\n",
    "Known Obstacle Positions: {agent.shared_variables[\"known_obstacle_pos\"]}\n",
    "Current Valid Actions: {agent.shared_variables[\"next_valid_moves\"]}'''\n",
    "    \n",
    "    # you can also influence how the planner performs the plan with additional details\n",
    "    global_context += '''\n",
    "Try to specify the action in the Subtask\n",
    "Example Assigned Task: ```Navigate to (1, 1)```\n",
    "Example Subtasks taking one action at a time: 'Move Down', 'Move Right' ```\n",
    "'''\n",
    "    \n",
    "    return global_context\n",
    "\n",
    "# Customise your grid here\n",
    "# O means obstacles, S means start, E means end, blank means nothing there\n",
    "grid_size = 5  # Grid size\n",
    "start_pos = (random.randint(0, grid_size//2 - 1), random.randint(0, grid_size//2 - 1))  # Starting position\n",
    "exit_pos = (random.randint(grid_size//2, grid_size-1), random.randint(grid_size//2, grid_size-1))  # Exit position\n",
    "num_obstacles = 5  # Number of obstacles\n",
    "\n",
    "grid = generate_grid(grid_size, start_pos, exit_pos, num_obstacles)\n",
    "valid_moves = check_valid_moves(start_pos, grid, grid_size)\n",
    "\n",
    "# Assign your agent\n",
    "my_agent = Agent('Maze Navigator', \n",
    "                 f'''You are to move to the Exit Position of the maze. Task is completed when Agent's position is at Exit Position.\n",
    "At each step, you must output an action - Up, Down, Left, Right or Stay.\n",
    "From current position (row, col), this is what you end up with after doing actions:\n",
    "{{'Up': (row-1, col), 'Down': (row+1, col), 'Left': (row, col-1), 'Right': (row, col+1), 'Stay': (row, col)}}\n",
    "You can only move to cells without obstacles. Only take an action from Current Valid Actions.\n",
    "If your previous action is invalid, choose another action.\n",
    "Top left of grid is (0, 0), bottom right is {(grid_size, grid_size)}.''', \n",
    "                 shared_variables = {\n",
    "                    \"cur_pos\": start_pos,\n",
    "                    \"exit_pos\": exit_pos,\n",
    "                    \"past_grid_states\": [],\n",
    "                    \"next_valid_moves\": valid_moves,\n",
    "                    \"known_obstacle_pos\": [],\n",
    "                    \"grid_size\": grid_size,\n",
    "                    \"grid\": grid}, \n",
    "                 max_subtasks = 20,\n",
    "                 get_global_context = get_global_context, # let Agent know the persistent states\n",
    "                 default_to_llm = False,\n",
    "                 llm = llm).assign_functions([move])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1c1f0bb7-a4a8-4d95-b262-dddfe51b8e84",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Starting Grid ###\n",
      "  S      \n",
      "        O\n",
      "O       E\n",
      "    O    \n",
      "  O     O\n",
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet, and the agent is currently at position (0, 1) with the exit at (2, 4).\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the exit, I need to move downwards or to the right, as those are the valid actions available. The most direct approach would be to move down first, as it brings me closer to the exit row.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Down\u001b[0m\n",
      "Calling function move with parameters {'action': 'Down'}\n",
      "         \n",
      "  S     O\n",
      "O       E\n",
      "    O    \n",
      "  O     O\n",
      "> {'output_1': 'Action successful. Agent moved from (0, 1) to (1, 1)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: Currently at position (1, 1) with the exit at (2, 4). No obstacles in the way and all actions are valid.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the exit, I need to move down to (2, 1) and then navigate towards (2, 4). The first step should be to move down.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Down to (2, 1)\u001b[0m\n",
      "Calling function move with parameters {'action': 'Down'}\n",
      "         \n",
      "        O\n",
      "O S     E\n",
      "    O    \n",
      "  O     O\n",
      "> {'output_1': 'Action successful. Agent moved from (1, 1) to (2, 1)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: Currently at position (2, 1) with the goal to reach (2, 4). The last visited position was (2, 1) and the only obstacle is at (2, 0).\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the exit at (2, 4), I can move right to (2, 2), then to (2, 3), and finally to (2, 4). The next step is to move right since it is a valid action.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Right to position (2, 2).\u001b[0m\n",
      "Calling function move with parameters {'action': 'Right'}\n",
      "         \n",
      "        O\n",
      "O   S   E\n",
      "    O    \n",
      "  O     O\n",
      "> {'output_1': 'Action successful. Agent moved from (2, 1) to (2, 2)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The agent is currently at position (2, 2) and needs to navigate to the exit position (2, 4). The last visited positions indicate that the agent has moved around but has not yet made progress towards the exit. The current valid actions include moving Up, Left, Right, or Staying.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To reach the exit position (2, 4), the agent can move Right to (2, 3) and then again Right to (2, 4) to complete the task. The next step should be to move Right since it is a valid action and leads directly towards the exit.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Right to (2, 3)\u001b[0m\n",
      "Calling function move with parameters {'action': 'Right'}\n",
      "         \n",
      "        O\n",
      "O     S E\n",
      "    O    \n",
      "  O     O\n",
      "> {'output_1': 'Action successful. Agent moved from (2, 2) to (2, 3)'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The agent is currently at position (2, 3) and needs to move to the exit position (2, 4). No obstacles are present in the immediate path.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the task, the agent should move to the right since that is the only valid action that leads directly to the exit position.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Move Right to position (2, 4)\u001b[0m\n",
      "Calling function move with parameters {'action': 'Right'}\n",
      "         \n",
      "        O\n",
      "O       S\n",
      "    O    \n",
      "  O     O\n",
      "> {'output_1': 'Action successful. Agent moved from (2, 3) to (2, 4)'}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print('### Starting Grid ###')\n",
    "print_grid(grid)\n",
    "\n",
    "num_moves = 0\n",
    "# Keep resetting subtask's history and changing start position to current position\n",
    "while num_moves < 20:\n",
    "    my_agent.reset()\n",
    "    my_agent.run(f\"Navigate to {my_agent.shared_variables['exit_pos']}\", num_subtasks = 1)\n",
    "    # use rule-based task checks as agent may not get it right all the time\n",
    "    if my_agent.shared_variables['cur_pos'] == my_agent.shared_variables['exit_pos']: \n",
    "        my_agent.task_completed = True\n",
    "        break\n",
    "    num_moves += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "fd2136df-e5c0-4c9a-9c9e-259d4d1ffe56",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent Name: Maze Navigator\n",
      "Agent Description: You are to move to the Exit Position of the maze. Task is completed when Agent's position is at Exit Position.\n",
      "At each step, you must output an action - Up, Down, Left, Right or Stay.\n",
      "From current position (row, col), this is what you end up with after doing actions:\n",
      "{'Up': (row-1, col), 'Down': (row+1, col), 'Left': (row, col-1), 'Right': (row, col+1), 'Stay': (row, col)}\n",
      "You can only move to cells without obstacles. Only take an action from Current Valid Actions.\n",
      "If your previous action is invalid, choose another action.\n",
      "Top left of grid is (0, 0), bottom right is (5, 5).\n",
      "Available Functions: ['end_task', 'move']\n",
      "Shared Variables: ['cur_pos', 'exit_pos', 'past_grid_states', 'next_valid_moves', 'known_obstacle_pos', 'grid_size', 'grid', 'agent']\n",
      "\u001b[1m\u001b[32mTask: Navigate to (2, 4)\u001b[0m\n",
      "\u001b[1m\u001b[30mSubtasks Completed:\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask: move(action=\"Right\")\u001b[0m\n",
      "{'output_1': 'Action successful. Agent moved from (2, 3) to (2, 4)'}\n",
      "\n",
      "Is Task Completed: True\n"
     ]
    }
   ],
   "source": [
    "my_agent.status()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4f4da96c-9b94-4e6d-a5b3-faff70bf5b95",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The agent is currently at position (2, 4), which is also the Exit Position. The last action taken was to move Right from (2, 3) to (2, 4), which was successful. Therefore, the task of navigating to (2, 4) is completed as the agent is already at the Exit Position.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The agent is currently at position (2, 4), which is also the Exit Position. The last action taken was to move Right from (2, 3) to (2, 4), which was successful. Therefore, the task of navigating to (2, 4) is completed as the agent is already at the Exit Position.'"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Note Agent's task is just one step away now, since we give a new task each step\n",
    "my_agent.reply_user()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fca84e41-3e23-4b51-aa9c-1fe4b5edb3c4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Legacy Support (Shared Variables via Function) - Calculator\n",
    "- `s_` at the start of the variable names means shared variables in `Function`\n",
    "    - For input, it means we take the variable from `shared_variables` instead of LLM generated input\n",
    "    - For output, it means we store the variable into `shared_variables` instead of storing it in `subtasks_completed`. If `subtasks_completed` output is empty, it will be output as `{'Status': 'Completed'}`\n",
    "- Example shared variables names: `s_sum`, `s_total`, `s_list_of_words`\n",
    "- Generally not a preferred approach because it can be confusing to the llm how to interpret these shared variables if it is placed in the function description directly\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "7346c6f3-d922-47db-bb85-8db7ddbeddad",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet for the assigned task of incrementing the total by 3 and then multiplying by 5.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the assigned task, I need to first increment the total by 3. Once that is done, I can then proceed to multiply the new total by 5.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Add 3 to the current total.\u001b[0m\n",
      "Calling function add_increment_to_total with parameters {'increment': 3}\n",
      "> {'Status': 'Completed'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: The increment of 3 has been successfully added to the total.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: The next step is to multiply the updated total by 5 to complete the assigned task.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Multiply the current total by 5 to finalize the computation.\u001b[0m\n",
      "Calling function multiply_s_total_by_factor with parameters {'factor': 5}\n",
      "> {'Status': 'Completed'}\n",
      "\n",
      "\u001b[1m\u001b[30mObservation: Both subtasks of incrementing the total by 3 and multiplying the total by 5 have been completed successfully.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since all parts of the assigned task have been executed, I can now finalize the process and provide the output to the user.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: End Task\u001b[0m\n",
      "Task completed successfully!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Function takes in increment (LLM generated) and s_total (retrieves from shared variable dict), and outputs to s_total (in shared variable dict)\n",
    "add = Function(fn_description = \"Add <increment: int> to <s_total>\", \n",
    "              output_format = {\"s_total\": \"Modified total\"},\n",
    "              llm = llm)\n",
    "\n",
    "# Function takes in factor (LLM generated) and s_total (retrieves from shared variable dict), and outputs to s_total (in shared variable dict)\n",
    "multiply = Function(fn_description = \"Multiply <s_total> by <factor: int>\",\n",
    "                    output_format = {\"s_total\": \"Modified total\"},\n",
    "                   llm = llm)\n",
    "\n",
    "# Define the calculator agent and the shared_variables - Note the naming convention of s_ at the start of the names for shared variables\n",
    "my_agent = Agent('Calculator', 'Does computations', \n",
    "                 shared_variables = {'s_total': 0},\n",
    "                 llm = llm).assign_functions([add, multiply])\n",
    "\n",
    "output = my_agent.run('Increment total by 3, then multiply by 5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "9661e89c-fed6-4208-9cd2-1610d6e574b7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current Total: 15\n"
     ]
    }
   ],
   "source": [
    "# This should be 15\n",
    "print('Current Total:', my_agent.shared_variables['s_total'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}