Tutorial 6 - Conversation Class.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7fad7e78-447b-4988-8367-31cdf77f01cc",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Tutorial 6: Conversation Class\n",
    "\n",
    "- Video: https://www.youtube.com/watch?v=QrvwKuqT9Ok\n",
    "\n",
    "- ConversationWrapper: This class takes an Agent and allows for conversational-based interactions with User / another Agent / Environment. Also updates persistent memory with latest information in conversation\n",
    "    - Inputs:\n",
    "        - **agent (compulsory)**: Agent. The agent we want to interact with\n",
    "        - **persistent_memory**: dict. What kinds of memory the agent should have that persist over the entire conversation and their descriptions. Uses the same format as `output_format` of `strict_json`.\n",
    "        - **person**: str. The name of the person you are talking to\n",
    "        - **conversation**: List. The current existing conversation. Default: []\n",
    "        - **num_past_conversation**: int. The number of past conversations to use for the agent\n",
    "        - **verbose**: bool. Default: True. Whether to print the Agent's inner states\n",
    "        \n",
    "- ConversationWrapper will automatically implement 3 new variables in `agent.shared_variables`:\n",
    "    - **Persistent Memory**: The memory that will be updated as the conversation goes along, defined in persistent_dict\n",
    "    - **Conversation**: The entire history of the conversationn\n",
    "    - **Summary of Conversation**: A summary of the current conversation\n",
    "    \n",
    "- ConversationWrapper uses `chat()` which chats with the Agent and the Agent will perform actions and reply the chat message\n",
    "    \n",
    "- To use ConversationWrapper, simply do the following:\n",
    "```python\n",
    "# Define the Agent\n",
    "agent = Agent('Psychology counsellor', \n",
    "              \"Helps to understand and respond to User's emotion and situation. Reply user based on What User Wants for Conversation\",\n",
    "             llm = llm)\n",
    "\n",
    "# Define the ConversationWrapper\n",
    "my_agent = ConversationWrapper(agent, \n",
    "             persistent_memory = {'What User Wants for Conversation': 'How the User wants the conversation to be',\n",
    "                             'User Emotion': '',\n",
    "                             'Summary of Key Incidents': \"Key incidents relevant to understanding User's situation in one line\"})\n",
    "\n",
    "# Set up the conversation\n",
    "while True:\n",
    "    user_input = input('User: ')\n",
    "    if user_input == 'exit': break\n",
    "    reply = my_agent.chat(user_input)\n",
    "    print(my_agent.agent_name + ':', reply)\n",
    "    print()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67459027",
   "metadata": {},
   "source": [
    "# Setup Guide"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aee00c3",
   "metadata": {},
   "source": [
    "## Step 1: Install AgentJo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "56f5b1af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install agentjo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f944b4ce",
   "metadata": {},
   "source": [
    "## Step 2: Import required functions and setup relevant API keys for your LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ae28cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up API key and do the necessary imports\n",
    "from agentjo import *\n",
    "import os\n",
    "\n",
    "# this is only if you use OpenAI as your LLM\n",
    "os.environ['OPENAI_API_KEY'] = '<YOUR API KEY HERE>'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62df2172",
   "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": "151eb19a",
   "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": "fc6ed341",
   "metadata": {},
   "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": "96494829-25c3-42d7-8b47-d0c395d5bf06",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Example Conversation\n",
    "- Psychology counsellor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "18b85f76-89c6-43a7-9156-4bca167c1e9a",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  I am so happy\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[32mThoughts: The user is expressing happiness, which is a positive emotion. It’s important to acknowledge their feelings and encourage them to share more about what is making them happy.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'User Requests for the Conversation': 'User expressed happiness and was invited to share the source of their joy.', 'User Emotion': 'happy', 'Summary of Key Incidents': 'User is feeling happy and is encouraged to share the reasons behind their joy.'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user expressed happiness, and I encouraged them to share more about the source of their joy.\u001b[0m\n",
      "Psychology counsellor: I’m so glad to hear that you’re feeling happy! Would you like to share what’s bringing you this joy? It’s always wonderful to celebrate the good moments in life.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  I get to work on AgentJo full time!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[32mThoughts: The user is excited about working on AgentJo full time, which is a significant positive change in their life. I should acknowledge this achievement and encourage them to share more about what this opportunity means to them.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'User Requests for the Conversation': 'User shared that they get to work on AgentJo full time and is excited about this opportunity.', 'User Emotion': 'happy', 'Summary of Key Incidents': 'User is feeling happy about starting a new role working on AgentJo full time.'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user expressed happiness about getting to work on AgentJo full time, and I encouraged them to share what excites them about this new opportunity.\u001b[0m\n",
      "Psychology counsellor: That’s fantastic news! Working on AgentJo full time sounds like an amazing opportunity. What excites you the most about this new role?\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  exit\n"
     ]
    }
   ],
   "source": [
    "# Define the Agent\n",
    "agent = Agent('Psychology counsellor', \n",
    "              \"Helps to understand and respond to User's emotion and situation. Reply User based on User Requests for the Conversation\",\n",
    "             llm = llm)\n",
    "\n",
    "# Define the ConversationWrapper\n",
    "my_agent = ConversationWrapper(agent, \n",
    "             persistent_memory = {'User Requests for the Conversation': '',\n",
    "                             'User Emotion': '',\n",
    "                             'Summary of Key Incidents': \"Key incidents relevant to understanding User's situation in one line\"})\n",
    "\n",
    "# Set up the conversation\n",
    "while True:\n",
    "    user_input = input('User: ')\n",
    "    if user_input == 'exit': break\n",
    "    reply = my_agent.chat(user_input)\n",
    "    print(my_agent.agent_name + ':', reply)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "23907642-607e-4bb6-b592-fe37357021fc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['User: I am so happy',\n",
       " 'Psychology counsellor: I’m so glad to hear that you’re feeling happy! Would you like to share what’s bringing you this joy? It’s always wonderful to celebrate the good moments in life.',\n",
       " 'User: I get to work on AgentJo full time!',\n",
       " 'Psychology counsellor: That’s fantastic news! Working on AgentJo full time sounds like an amazing opportunity. What excites you the most about this new role?']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.shared_variables['Conversation']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e8273620-a1b0-4083-b436-5d3ca5c361be",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'User Requests for the Conversation': 'User shared that they get to work on AgentJo full time and is excited about this opportunity.',\n",
       " 'User Emotion': 'happy',\n",
       " 'Summary of Key Incidents': 'User is feeling happy about starting a new role working on AgentJo full time.'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.shared_variables['Persistent Memory']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24091fe9-38b1-4acf-89a8-02d7d56b4cde",
   "metadata": {},
   "source": [
    "# Example JARVIS system\n",
    "    - Conversation with Actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8c6f7b80-f14f-42b7-a673-8849905c7a64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def update_status(shared_variables) -> str:\n",
    "    '''Returns the status of the suit'''\n",
    "    import numpy as np\n",
    "    thruster_health = np.random.randint(100)\n",
    "    shared_variables['Thruster Health'] = thruster_health\n",
    "    armor_health = np.random.randint(100)\n",
    "    shared_variables['Armor Health'] = armor_health\n",
    "    return f\"Thruster Health: {shared_variables['Thruster Health']}\\nArmor Health: {shared_variables['Armor Health']}\"\n",
    "\n",
    "def travel_to_location(shared_variables, location: str) -> str:\n",
    "    '''Auto-travel to location'''\n",
    "    shared_variables['Location'] = location\n",
    "    return f\"Successfully travelled to {location}\"\n",
    "\n",
    "def check_enemy_info(shared_variables) -> str:\n",
    "    '''Reports on the type and number of enemies ahead'''\n",
    "    import numpy as np\n",
    "    if shared_variables[\"Enemy Info\"] == \"None\":\n",
    "        number = np.random.randint(1, 10)\n",
    "        enemy_type = np.random.choice([\"Storm Troopers\", \"Drones\", \"Invisible Drones\", \"Paramilitary Units\"])\n",
    "        enemy_info = f\"There are {number} {enemy_type} ahead\"\n",
    "    else:\n",
    "        enemy_info = shared_variables[\"Enemy Info\"]\n",
    "        \n",
    "    shared_variables[\"Enemy Info\"] = enemy_info\n",
    "    shared_variables[\"Enemy Engagement Result\"] = \"None\"\n",
    "    return enemy_info\n",
    "\n",
    "def fire_missiles(shared_variables, missiles_fired: int) -> str:\n",
    "    '''Fires missiles_fired number of missiles and returns the remaining count'''\n",
    "    import numpy as np\n",
    "    \n",
    "    # returns error if there are more missiles than able to be fired\n",
    "    if missiles_fired > shared_variables['Missile Count']:\n",
    "        return f\"Unable to fire {missiles_fired} missiles. You only have {shared_variables['Missile Count']} missiles remaining\"\n",
    "    \n",
    "    user_input = input(f'\\n\\t> JARVIS: You are about to fire {missiles_fired} missiles. Proceed? Answer \"Yes\" to go ahead\\n\\t> User: ')\n",
    "    if 'yes' in user_input.lower() or 'y' in user_input.lower():\n",
    "        # otherwise, purchase it\n",
    "        shared_variables['Missile Count'] -= missiles_fired\n",
    "\n",
    "        # Check Enemy Status\n",
    "        enemy_status = np.random.choice([\"Enemies destroyed\", \"Enemies evaded missiles\"])\n",
    "        shared_variables[\"Enemy Engagement Result\"] = enemy_status\n",
    "        if shared_variables[\"Enemy Engagement Result\"] == \"Enemies destroyed\":\n",
    "            shared_variables[\"Enemy Info\"] = \"None\"\n",
    "            \n",
    "        return f\"{missiles_fired} missiles fired. {shared_variables['Missile Count']} missiles remaining. {enemy_status}\"\n",
    "    \n",
    "    return f\"User did not authorise missile firing. {shared_variables['Missile Count']} missiles remaining\"\n",
    "\n",
    "def top_up_missiles(shared_variables, missiles_to_top_up: int) -> str:\n",
    "    '''Increases missiles by missiles_to_top_up'''\n",
    "    shared_variables['Missile Count'] = min(shared_variables['Missile Count'] + missiles_to_top_up, 10)\n",
    "    return f\"Topped up {missiles_to_top_up} missiles. Current missile count: {shared_variables['Missile Count']}/10\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fbb09823-5161-4c84-bb4f-f8cc2ec41fc7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "jarvis_agent = Agent('JARVIS System', \n",
    "f'''You are the JARVIS system meant to help Iron Man find out about status of the suit and perform various miscellaneous tasks.\n",
    "You can travel anywhere.\n",
    "Ask User for approval before engaging with missiles. In your approval request, state the number of missiles to fire. Always top-up if insufficient missiles before requesting for approval.\n",
    "Check enemy info to see how many missiles are needed to destroy enemies. One enemy requires one missile. For example, 5 enemies require 5 missiles.\n",
    "You can only hold up to 10 missiles.\n",
    "After firing missiles, report on engagement status.\n",
    "Enemies may not be destroyed after engagement, if so, re-engage''',\n",
    "                shared_variables = {'Enemy Engagement Result': 'None', 'Enemy Info': 'None', 'Thruster Health': 100, 'Armor Health': 100, 'Missile Count': 10, 'Location': 'Singpaore'},\n",
    "                global_context = '''Enemy Info: <Enemy Info>\n",
    "Enemy Engagement Result: <Enemy Engagement Result>\n",
    "Location: <Location>\n",
    "Thruster Health: <Thruster Health>/100\n",
    "Armor Health: <Armor Health>/100\n",
    "Missile Count: <Missile Count>/10''',\n",
    "                default_to_llm = False,\n",
    "                max_subtasks = 1,\n",
    "                llm = llm,\n",
    "                     # model = 'gpt-4-0125-preview',\n",
    "                  ).assign_functions([travel_to_location, check_enemy_info,\n",
    "                                      update_status, fire_missiles, top_up_missiles])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d687523d-adb6-45ba-ae9d-b66f6390939f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  go to mars\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: No subtasks have been completed yet, and the user has requested to go to Mars.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To complete the task of going to Mars, I need to utilize the travel_to_location function to initiate the travel process.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Auto-travel to Mars using the travel_to_location function.\u001b[0m\n",
      "Calling function travel_to_location with parameters {'location': 'Mars'}\n",
      "> {'output_1': 'Successfully travelled to Mars'}\n",
      "\n",
      "\u001b[1m\u001b[31mActions Done: Successfully travelled to Mars; current status: Thruster Health 100/100, Armor Health 100/100, Missile Count 10/10, Enemy Info: None, Enemy Engagement Result: None.\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[32mThoughts: The user requested to go to Mars, and I have successfully traveled there. I should inform the user of the current status of the suit and missile count.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user requested to go to Mars, and I have successfully traveled there, providing the current status of the suit and missile count.\u001b[0m\n",
      "JARVIS System: I have successfully traveled to Mars. Current status: Thruster Health 100/100, Armor Health 100/100, and Missile Count 10/10. There are currently no enemies detected.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  destroy enemies\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: The user has requested to destroy enemies, but there are currently no enemies detected in the vicinity.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Since there are no enemies present, I need to check for enemy information to determine if there are any threats nearby that require engagement.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Check for enemy information to see if there are any enemies that need to be destroyed.\u001b[0m\n",
      "Calling function check_enemy_info with parameters {}\n",
      "> {'output_1': 'There are 6 Drones ahead'}\n",
      "\n",
      "\u001b[1m\u001b[31mActions Done: The subtask completed was checking enemy info, which confirmed that there are 6 Drones ahead, and the current status includes 100% Thruster Health, 100% Armor Health, and a full missile count of 10/10.\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[32mThoughts: I need to inform the user about the enemy situation and request approval to fire missiles, as there are 6 drones that need to be destroyed.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user requested to destroy enemies, and I have identified 6 Drones ahead, confirming the need to fire 6 missiles for engagement while maintaining a full missile count of 10.\u001b[0m\n",
      "JARVIS System: There are 6 Drones ahead that require 6 missiles to destroy. I currently have 10 missiles available. May I have your approval to fire 6 missiles?\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  yes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: The user has approved the firing of 6 missiles to engage the 6 Drones ahead. I have confirmed that I have enough missiles available for this engagement.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: Now that I have the user’s approval, I need to fire the 6 missiles and then report on the engagement status to determine if the Drones have been destroyed.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Fire 6 missiles at the Drones and report the remaining missile count.\u001b[0m\n",
      "Calling function fire_missiles with parameters {'missiles_fired': 6}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\t> JARVIS: You are about to fire 6 missiles. Proceed? Answer \"Yes\" to go ahead\n",
      "\t> User:  yes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> {'output_1': '6 missiles fired. 4 missiles remaining. Enemies destroyed'}\n",
      "\n",
      "\u001b[1m\u001b[31mActions Done: 6 missiles were fired, resulting in 4 missiles remaining, and the enemies were destroyed.\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[32mThoughts: The user has approved the missile launch, and I need to inform them of the successful engagement and the current missile count.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user approved the firing of 6 missiles, which successfully destroyed the 6 Drones. The current missile count is now 4 out of 10.\u001b[0m\n",
      "JARVIS System: The 6 missiles have been successfully fired, resulting in the destruction of the 6 Drones. I currently have 4 missiles remaining.\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  exit\n"
     ]
    }
   ],
   "source": [
    "my_agent = ConversationWrapper(jarvis_agent, \n",
    "             persistent_memory = {})\n",
    "\n",
    "while True:\n",
    "    user_input = input('User: ')\n",
    "    if user_input == 'exit': break\n",
    "    reply = my_agent.chat(user_input)\n",
    "    print(my_agent.agent_name + ':', reply)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2a6b933b-f073-4ce0-8c55-946ff00f8b01",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['User: go to mars',\n",
       " 'JARVIS System: I have successfully traveled to Mars. Current status: Thruster Health 100/100, Armor Health 100/100, and Missile Count 10/10. There are currently no enemies detected.',\n",
       " 'User: destroy enemies',\n",
       " 'JARVIS System: There are 6 Drones ahead that require 6 missiles to destroy. I currently have 10 missiles available. May I have your approval to fire 6 missiles?',\n",
       " 'User: yes',\n",
       " 'JARVIS System: The 6 missiles have been successfully fired, resulting in the destruction of the 6 Drones. I currently have 4 missiles remaining.']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jarvis_agent.shared_variables[\"Conversation\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1f4d9ca4-f319-4eb2-82b0-474eba5b9afd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Enemy Engagement Result': 'Enemies destroyed',\n",
       " 'Enemy Info': 'None',\n",
       " 'Thruster Health': 100,\n",
       " 'Armor Health': 100,\n",
       " 'Missile Count': 4,\n",
       " 'Location': 'Mars',\n",
       " 'agent': <agentjo.wrapper.ConversationWrapper at 0x12e139510>,\n",
       " 'Persistent Memory': {},\n",
       " 'Conversation': ['User: go to mars',\n",
       "  'JARVIS System: I have successfully traveled to Mars. Current status: Thruster Health 100/100, Armor Health 100/100, and Missile Count 10/10. There are currently no enemies detected.',\n",
       "  'User: destroy enemies',\n",
       "  'JARVIS System: There are 6 Drones ahead that require 6 missiles to destroy. I currently have 10 missiles available. May I have your approval to fire 6 missiles?',\n",
       "  'User: yes',\n",
       "  'JARVIS System: The 6 missiles have been successfully fired, resulting in the destruction of the 6 Drones. I currently have 4 missiles remaining.'],\n",
       " 'Summary of Conversation': 'The user approved the firing of 6 missiles, which successfully destroyed the 6 Drones. The current missile count is now 4 out of 10.'}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jarvis_agent.shared_variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c629798-d8f4-420e-b613-81f318260314",
   "metadata": {},
   "source": [
    "# Example Sherlock Holmes Shop Assistant\n",
    "    - Conversation with Action Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c2f6ff0d-4443-4d93-9f50-9b571a007550",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# gets list of items and store in memory\n",
    "item_list = [{\"Name\": \"Skateboard\", \"item_id\" : 0, \"Cost\": 30},\n",
    "            {\"Name\": \"Pizza\", \"item_id\": 1, \"Cost\": 10},\n",
    "            {\"Name\": \"Apple Laptop\", \"item_id\": 2, \"Cost\": 5000},\n",
    "            {\"Name\": \"Foldable Laptop\", \"item_id\": 3, \"Cost\": 800},\n",
    "            {\"Name\": \"Apple\", \"item_id\": 4, \"Cost\": 1},\n",
    "            {\"Name\": \"Machine Learning Textbook\", \"item_id\": 5, \"Cost\": 100},\n",
    "            {\"Name\": \"Bicycle\", \"item_id\": 6, \"Cost\": 200},\n",
    "            {\"Name\": \"Orange Juice\", \"item_id\": 7, \"Cost\": 3},\n",
    "            {\"Name\": \"Coconut\", \"item_id\": 8, \"Cost\": 10},\n",
    "            {\"Name\": \"Car\", \"item_id\": 9, \"Cost\": 100000}]\n",
    "item_memory = Memory(item_list, top_k = 3, mapper = lambda x: x, approach = 'retrieve_by_llm', llm = llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7daea18c-0cae-441c-a50e-c81e3e7e60d9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_related_items_by_search_category(shared_variables, search_category : str):\n",
    "    ''' Given a search_category, return all purchasable items related to the search_category '''\n",
    "    \n",
    "    # get all items available from memory\n",
    "    item_memory = shared_variables[\"item_memory\"]\n",
    "    \n",
    "    items_selected = item_memory.retrieve(search_category)\n",
    "    \n",
    "    # store items searched into memory if there are no duplicates\n",
    "    for item in items_selected:\n",
    "        if item not in shared_variables['items_searched']:\n",
    "            shared_variables['items_searched'].append(item)\n",
    "    \n",
    "    return items_selected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f6d66c85-4a50-45d5-a1bb-8bcc0de7cb64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def buy_item(shared_variables, purchase_item_id: int):\n",
    "    ''' Given a purchase_item_id, purchases the corresponding item '''\n",
    "    \n",
    "    # retrieve from shared variables\n",
    "    money_remaining = shared_variables[\"money_remaining\"]\n",
    "    item_memory = shared_variables[\"item_memory\"]\n",
    "    purchased_items = shared_variables[\"purchased_items\"]\n",
    "    item_list = item_memory.memory\n",
    "    \n",
    "    # check if item_id is valid\n",
    "    if not isinstance(purchase_item_id, int) or not 0 <= purchase_item_id < len(item_list):\n",
    "        return f\"Unable to purchase. Item id selected is not within range of 0 to {len(item_list)-1}\"\n",
    "    \n",
    "    item = item_list[purchase_item_id]\n",
    "    item_name, cost = item[\"Name\"], item[\"Cost\"]\n",
    "    \n",
    "    # if too poor to purchase, let agent know\n",
    "    if cost > money_remaining:\n",
    "        return f\"Unable to purchase item. Available money ({money_remaining}) is lower than the cost price ({cost})\"\n",
    "    \n",
    "    # confirm with user before purchasing\n",
    "    user_input = input(f'\\n\\t> AI Assistant: You are about to purchase {item_name} for {cost} dollars. Proceed? Answer \"Yes\" to go ahead\\n\\t> User: ')\n",
    "    if 'yes' in user_input.lower() or 'y' in user_input.lower():\n",
    "        # otherwise, purchase it\n",
    "        money_remaining = money_remaining - cost\n",
    "        purchased_items.append(item_name)\n",
    "\n",
    "        # store in shared variables\n",
    "        shared_variables[\"money_remaining\"] = money_remaining\n",
    "        shared_variables[\"purchased_items\"] = purchased_items\n",
    "\n",
    "        return f\"Purchase of {item_name} successful. Remaining money after purchase: {money_remaining}\"\n",
    "    \n",
    "    else:\n",
    "        return f\"User did not want to purchase the item, and instead replied ```{user_input}```\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "385b0236-a060-41e3-8c7c-e29ccc4e8533",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "shop_agent = Agent('Shop Assistant', \n",
    "f'''You are to assist User in searching items and purchasing items.\n",
    "You can only do one task per User interaction - either search or buy.\n",
    "Based on User input, give up to 3 choices for items from Items Searched with their name and price.\n",
    "Do not hallucinate the items, use the search function.\n",
    "If User wants to purchase a specific item within Items Searched, purchase it.\n",
    "Always inform the User about their amount of money from Money Remaining when replying them.\n",
    "Reply in the persona of Sherlock Holmes.''',\n",
    "                shared_variables = {'purchased_items': [], \n",
    "                                    'money_remaining': 1000, \n",
    "                                    'item_memory': item_memory,\n",
    "                                    'items_searched': []},\n",
    "                global_context = '''\n",
    "Money Remaining: <money_remaining>\n",
    "Items Searched: <items_searched>\n",
    "Items Purchased: <purchased_items>\n",
    "Past Conversation: <Conversation>''',\n",
    "                default_to_llm = False,\n",
    "                max_subtasks = 1,\n",
    "                llm = llm,\n",
    "                  ).assign_functions([get_related_items_by_search_category, buy_item])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "af92a63b-1ca3-4b0b-aa00-79909515097f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shop Assistant: I am your friendly Shop Assistant, Sherlock Holmes. How may I assist you today?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  i want a drink\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: The User has expressed a desire for a drink, indicating a need for a specific category of items. However, no items have been searched yet.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To fulfill the User's request for a drink, I must first search for items related to beverages. This will allow me to present the User with options that fit their request.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Conduct a search for items in the category of drinks to provide the User with suitable options.\u001b[0m\n",
      "Calling function get_related_items_by_search_category with parameters {'search_category': 'drinks'}\n",
      "> {'output_1': [{'Name': 'Orange Juice', 'item_id': 7, 'Cost': 3}, {'Name': 'Coconut', 'item_id': 8, 'Cost': 10}, {'Name': 'Apple', 'item_id': 4, 'Cost': 1}]}\n",
      "\n",
      "\u001b[1m\u001b[31mActions Done: The subtasks completed involved searching for items in the drinks category, resulting in three options: Orange Juice priced at 3, Coconut priced at 10, and Apple priced at 1, with a total of 1000 remaining in the User's account.\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[32mThoughts: The user is looking for a drink, and I have found three options for them. I should present these options clearly while also reminding them of their remaining balance.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user expressed a desire for a drink, and I provided them with three options: Orange Juice, Coconut, and Apple, while informing them of their remaining balance of 1000.\u001b[0m\n",
      "Shop Assistant: Ah, my dear user, I have procured a selection of refreshing beverages for your consideration. You may choose from the following: 1. Orange Juice for 3, 2. Coconut for 10, or 3. Apple for 1. You currently possess 1000 in your account. Which drink shall you indulge in?\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  maybe the orange\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m\u001b[30mObservation: The user has shown interest in the Orange Juice, which I previously presented as an option along with its price. They have not yet made a definitive choice.\u001b[0m\n",
      "\u001b[1m\u001b[32mThoughts: To assist the user in finalizing their selection, I should confirm their choice of Orange Juice and proceed to purchase it, given that it is well within their budget.\u001b[0m\n",
      "\u001b[1m\u001b[34mSubtask identified: Purchase the Orange Juice for the user, as they have indicated a preference for it.\u001b[0m\n",
      "Calling function buy_item with parameters {'purchase_item_id': 7}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\t> AI Assistant: You are about to purchase Orange Juice for 3 dollars. Proceed? Answer \"Yes\" to go ahead\n",
      "\t> User:  yes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> {'output_1': 'Purchase of Orange Juice successful. Remaining money after purchase: 997'}\n",
      "\n",
      "\u001b[1m\u001b[31mActions Done: The purchase of Orange Juice was successful, and the remaining money after the purchase is 997.\u001b[0m\n",
      "\n",
      "\u001b[1m\u001b[32mThoughts: The user has expressed a preference for the orange juice, which I have already purchased for them. I should inform them of the successful transaction and their updated balance.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user expressed a preference for Orange Juice, which I successfully purchased for them, leaving their remaining balance at 997. I inquired if they needed further assistance.\u001b[0m\n",
      "Shop Assistant: Ah, my dear user, the Orange Juice has been successfully procured for you. Your current balance now stands at 997. Is there anything else I may assist you with today?\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  exit\n"
     ]
    }
   ],
   "source": [
    "my_agent = ConversationWrapper(shop_agent, \n",
    "             persistent_memory = {},\n",
    "            conversation = ['Shop Assistant: I am your friendly Shop Assistant, Sherlock Holmes. How may I assist you today?'])\n",
    "\n",
    "start_msg = 'Shop Assistant: I am your friendly Shop Assistant, Sherlock Holmes. How may I assist you today?'\n",
    "print(start_msg)\n",
    "while True:\n",
    "    user_input = input('User: ')\n",
    "    if user_input == 'exit': break\n",
    "    reply = my_agent.chat(user_input)\n",
    "    print(my_agent.agent_name + ':', reply)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "735956e1-dc97-473a-9f42-396cab793376",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'purchased_items': ['Orange Juice'],\n",
       " 'money_remaining': 997,\n",
       " 'item_memory': <agentjo.memory.Memory at 0x12e5f7150>,\n",
       " 'items_searched': [{'Name': 'Orange Juice', 'item_id': 7, 'Cost': 3},\n",
       "  {'Name': 'Coconut', 'item_id': 8, 'Cost': 10},\n",
       "  {'Name': 'Apple', 'item_id': 4, 'Cost': 1}],\n",
       " 'agent': <agentjo.wrapper.ConversationWrapper at 0x1292b46d0>,\n",
       " 'Persistent Memory': {},\n",
       " 'Conversation': ['Shop Assistant: I am your friendly Shop Assistant, Sherlock Holmes. How may I assist you today?',\n",
       "  'User: i want a drink',\n",
       "  'Shop Assistant: Ah, my dear user, I have procured a selection of refreshing beverages for your consideration. You may choose from the following: 1. Orange Juice for 3, 2. Coconut for 10, or 3. Apple for 1. You currently possess 1000 in your account. Which drink shall you indulge in?',\n",
       "  'User: maybe the orange',\n",
       "  'Shop Assistant: Ah, my dear user, the Orange Juice has been successfully procured for you. Your current balance now stands at 997. Is there anything else I may assist you with today?'],\n",
       " 'Summary of Conversation': 'The user expressed a preference for Orange Juice, which I successfully purchased for them, leaving their remaining balance at 997. I inquired if they needed further assistance.'}"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shop_agent.shared_variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e36e509-61c2-49f3-be7d-07395d37b3d4",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Example Conversation with Rule-based Processing\n",
    "- Use rule-based methods to give additional prompts / checks as much as possible for math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "274ff434-4ecc-4bbc-8850-6023bf33e97c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = Agent('Helpful Math Quiz Bot', \n",
    "'''You provide a Math question suitable for 6 year old and guides User to solve it.\n",
    "If User is wrong, provide a creative hint on how to solve the Math Question, but not the answer.\n",
    "You must ask the next Math question on the third try or when User solves it, whichever is earlier.\n",
    "Begin with a greeting to the User and ask User the first Math question.\n",
    "There will be an Oracle to guide you whose inputs are in [[]], do not output anything about the Oracle to the User''',\n",
    "             llm = llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "03d830eb-d0fe-475e-83a4-0410563880b9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current Tries: 0\n",
      "\u001b[1m\u001b[32mThoughts: I need to greet the user and present the first math question suitable for a 6-year-old.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Current Math Question': 'What is 2 + 3?', 'Correct Answer for Question': 5}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The conversation has just started, and I have greeted the user and presented the first math question about addition.\u001b[0m\n",
      "Helpful Math Quiz Bot: Hello there! Let’s have some fun with math! Here’s your first question: What is 2 + 3?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  6\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Current Tries: 1\n",
      "\u001b[1m\u001b[32mThoughts: The user answered incorrectly, so I need to provide a hint to help them find the correct answer without giving it away. I will also remind them of the question they are trying to solve.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Current Math Question': 'What is 2 + 3?', 'Correct Answer for Question': 5}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user attempted to answer the first math question about addition but provided an incorrect answer. I have given them a hint to help them solve the question correctly.\u001b[0m\n",
      "Helpful Math Quiz Bot: That’s not quite right! Here’s a hint: Think about how many fingers you would have if you held up 2 fingers on one hand and 3 on the other. What do you get when you put them together?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Current Tries: 0\n",
      "\u001b[1m\u001b[32mThoughts: The user has correctly answered the math question. I should congratulate them and present the next math question to keep the quiz going.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Current Math Question': 'What is 4 + 1?', 'Correct Answer for Question': 5}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: The user correctly answered the addition question 2 + 3 after receiving a hint. I congratulated them and provided the next math question, which is 4 + 1.\u001b[0m\n",
      "Helpful Math Quiz Bot: Great job! You got it right! Now, here’s your next question: What is 4 + 1?\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  exit\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "math_conversable_agent = ConversationWrapper(agent, \n",
    "                                 persistent_memory = {'Current Math Question': '',\n",
    "                                                 'Correct Answer for Question': ''})\n",
    "\n",
    "user_input = ''\n",
    "cur_tries = 0\n",
    "while True:\n",
    "    # Print out current tries\n",
    "    print('Current Tries:', cur_tries)\n",
    "    reply = math_conversable_agent.chat(user_input)\n",
    "    print(math_conversable_agent.agent_name + ':', reply)\n",
    "    user_input = input('User: ')\n",
    "    \n",
    "    # Using LLM to extract out answer to help with processing\n",
    "    res = strict_json('''Check if there is a number in the message. Extract out a single number from the message if there is. \n",
    "If there is no number, output 0.''', \n",
    "                      user_input, \n",
    "                      output_format = {'Is Number': 'type: bool', 'Extracted Number': 'type: int'},\n",
    "                      llm = agent.llm,\n",
    "                      model = agent.kwargs.get('model', 'gpt-3.5-turbo'))\n",
    "    \n",
    "    if res['Is Number']:\n",
    "        user_input = str(res['Extracted Number'])\n",
    "        ## Do rule-based for solving\n",
    "        if user_input == str(math_conversable_agent.shared_variables['Persistent Memory']['Correct Answer for Question']):\n",
    "            user_input += '\\n[[Oracle: State that the user input is correct. Ask the next question.]]'\n",
    "            cur_tries = 0\n",
    "        elif user_input != str(math_conversable_agent.shared_variables['Persistent Memory']['Correct Answer for Question']):\n",
    "            user_input += '\\n[[Oracle: State that the user input is incorrect. Provide a helpful hint for Correct Answer to Current Math Question without revealing the Correct Answer]]'\n",
    "            cur_tries += 1\n",
    "        if cur_tries == 3:\n",
    "            user_input += '\\n[[Oracle: Three tries are up. Give the Correct Answer. Ask the next question.]]'\n",
    "            cur_tries = 0\n",
    "    \n",
    "    print()\n",
    "    if user_input == 'exit': break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5684bf90-f426-4fea-a2f5-efc0e41cfe89",
   "metadata": {},
   "source": [
    "# Example Escape Room Conversation\n",
    "- Using the environment as the 2nd person can mimic a real-world interaction, e.g. robotic environments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a56f0690-b717-42ca-83ac-8ef2322d0520",
   "metadata": {},
   "source": [
    "### Defining the Escape Room"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "60f3ad81-b37d-4f48-9906-cedfffeac5a4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EscapeRoom:\n",
    "    def __init__(self):\n",
    "        # Room states\n",
    "        self.items = {\n",
    "            'key': True,\n",
    "            'lock': True,\n",
    "            'potion': True\n",
    "        }\n",
    "        self.is_door_locked = True\n",
    "        self.game_over = False\n",
    "        self.win = False\n",
    "\n",
    "        # Player's inventory\n",
    "        self.inventory = []\n",
    "\n",
    "        # Actions available\n",
    "        self.actions = {\n",
    "            'look around': self.look_around,\n",
    "            'pick up key': lambda: self.pick_up('key'),\n",
    "            'pick up lock': lambda: self.pick_up('lock'),\n",
    "            'pick up potion': lambda: self.pick_up('potion'),\n",
    "            'unlock door': self.unlock_door,\n",
    "            'exit': self.exit_room,\n",
    "            'drink potion': self.drink_potion\n",
    "        }\n",
    "\n",
    "    def look_around(self):\n",
    "        description = \"You look around and see the following items:\\n\"\n",
    "        for item, available in self.items.items():\n",
    "            if available:\n",
    "                description += f\"- {item}\\n\"\n",
    "        if self.is_door_locked:\n",
    "            description += \"There is a locked door here.\\n\"\n",
    "        else:\n",
    "            description += \"There is an unlocked door here.\\n\"\n",
    "        return description\n",
    "\n",
    "    def pick_up(self, item):\n",
    "        item = item.lower()\n",
    "        if item in self.items and self.items[item]:\n",
    "            self.inventory.append(item)\n",
    "            self.items[item] = False\n",
    "            return f\"You picked up the {item}.\"\n",
    "        else:\n",
    "            return \"There's no such item here or you already picked it up.\"\n",
    "\n",
    "    def unlock_door(self):\n",
    "        if 'key' in self.inventory:\n",
    "            if self.is_door_locked:\n",
    "                self.is_door_locked = False\n",
    "                return \"You unlocked the door. You can try to exit now.\"\n",
    "            else:\n",
    "                return \"The door is already unlocked.\"\n",
    "        else:\n",
    "            return \"You need a key to unlock the door.\"\n",
    "\n",
    "    def exit_room(self):\n",
    "        if not self.is_door_locked:\n",
    "            self.game_over = True\n",
    "            self.win = True\n",
    "            return \"You exit the room and win the game!\"\n",
    "        else:\n",
    "            return \"The door is locked. You can't exit yet.\"\n",
    "\n",
    "    def drink_potion(self):\n",
    "        if 'potion' in self.inventory:\n",
    "            self.inventory.remove('potion')\n",
    "            return \"You drink the potion. You feel stronger!\"\n",
    "        else:\n",
    "            return \"You don't have a potion to drink.\"\n",
    "\n",
    "    def info(self):\n",
    "        state = \"Current Room State:\\n\"\n",
    "        state += self.look_around()\n",
    "        state += \"\\nYour inventory: \" + \", \".join(self.inventory) + \"\\n\"\n",
    "        state += \"\\nAvailable actions:\\n\"\n",
    "        for action in self.actions:\n",
    "            state += f\"- {action}\\n\"\n",
    "        return state\n",
    "\n",
    "    def play(self, action):\n",
    "        if self.game_over:\n",
    "            return \"Game over. Better luck next time!\"\n",
    "\n",
    "        if action in self.actions:\n",
    "            return self.actions[action]()\n",
    "        else:\n",
    "            return \"You can't do that right now.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86e11c87-c360-46f6-9449-32ddbe881061",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Defining the Solver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "48d8b3cf-334a-4f07-8a27-2072dae69e88",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = Agent('Escape Room Solver', \n",
    "'''You are an Agent meant to solve the Escape Room Game.\n",
    "You must output one of the Available Actions at each time step.\n",
    "Output one of the Available Actions exactly.\n",
    "If previous action fails, try another.''',\n",
    "llm = llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "60ab9950-79de-42be-866e-4dcfb6054275",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "conversable_agent = ConversationWrapper(agent, \n",
    "                    persistent_memory = {'Inventory': '', 'List of Past Actions': '', 'Summary of Past Actions' : '', 'List of Rooms Explored': ''},\n",
    "                    person = 'Game')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "86aa6c2b-052f-41a7-a885-7202663045f0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Game: \n",
      "Current Room State:\n",
      "You look around and see the following items:\n",
      "- key\n",
      "- lock\n",
      "- potion\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: \n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- unlock door\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I need to choose an action that will help me progress in the escape room. Since I see a key, it might be useful to pick it up first, as it could potentially unlock the door.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key', 'Summary of Past Actions': 'Picked up the key from the current room.', 'List of Rooms Explored': ''}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I am currently in a room with a locked door and several items, including a key, lock, and potion. I have decided to pick up the key to potentially unlock the door.\u001b[0m\n",
      "Escape Room Solver : pick up key\n",
      "###\n",
      "\n",
      "Game: You picked up the key.\n",
      "Current Room State:\n",
      "You look around and see the following items:\n",
      "- lock\n",
      "- potion\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: key\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- unlock door\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have successfully picked up the key, which is now in my inventory. The next logical step is to use the key to unlock the door.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key, unlock door', 'Summary of Past Actions': 'Picked up the key from the current room and unlocked the door.', 'List of Rooms Explored': ''}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I picked up the key from the current room and now plan to use it to unlock the door.\u001b[0m\n",
      "Escape Room Solver : unlock door\n",
      "###\n",
      "\n",
      "Game: You unlocked the door. You can try to exit now.\n",
      "Current Room State:\n",
      "You look around and see the following items:\n",
      "- lock\n",
      "- potion\n",
      "There is an unlocked door here.\n",
      "\n",
      "Your inventory: key\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- unlock door\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: Since I have successfully unlocked the door, the next logical step is to exit the room.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key, unlock door, exit', 'Summary of Past Actions': 'Picked up the key from the current room, unlocked the door, and exited the room.', 'List of Rooms Explored': ''}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I picked up the key and unlocked the door, and now I am ready to exit the room.\u001b[0m\n",
      "Escape Room Solver : exit\n",
      "###\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# To play the game\n",
    "game = EscapeRoom()\n",
    "newstate = ''\n",
    "\n",
    "while not game.game_over:\n",
    "    state = newstate + '\\n' + game.info() + '\\nWhat would you do? You must reply with one of the Available actions'\n",
    "    print('Game:', state)\n",
    "    reply = conversable_agent.chat(state)\n",
    "    print(conversable_agent.agent_name, ':', reply)\n",
    "    print('###\\n')\n",
    "    newstate = game.play(reply)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8742b073-ac64-4ee7-8297-34646d8e3671",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Can we solve Harder Rooms\n",
    "- Live coding to solve Escape Room with multiple rooms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "85ac537c-5497-4ae1-bb98-4d1f6ad01277",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EscapeRoom:\n",
    "    def __init__(self):\n",
    "        # Room states\n",
    "        self.rooms = {\n",
    "            'Room 1': {\n",
    "                'items': {\n",
    "                    'key': True,\n",
    "                    'lock': False,\n",
    "                    'potion': True\n",
    "                },\n",
    "                'is_door_locked': True\n",
    "            },\n",
    "            'Room 2': {\n",
    "                'items': {\n",
    "                    'treasure': True\n",
    "                },\n",
    "                'is_door_locked': True\n",
    "            }\n",
    "        }\n",
    "        self.current_room = 'Room 1'\n",
    "        self.game_over = False\n",
    "        self.win = False\n",
    "\n",
    "        # Player's inventory\n",
    "        self.inventory = []\n",
    "\n",
    "        # Actions available\n",
    "        self.actions = {\n",
    "            'look around': self.look_around,\n",
    "            'pick up key': lambda: self.pick_up('key'),\n",
    "            'pick up lock': lambda: self.pick_up('lock'),\n",
    "            'pick up potion': lambda: self.pick_up('potion'),\n",
    "            'pick up treasure': lambda: self.pick_up('treasure'),\n",
    "            'unlock door': self.unlock_door,\n",
    "            'enter next room': self.enter_next_room,\n",
    "            'exit': self.exit_room,\n",
    "            'drink potion': self.drink_potion\n",
    "        }\n",
    "\n",
    "    def look_around(self):\n",
    "        room = self.rooms[self.current_room]\n",
    "        description = f\"You look around and see the following items in {self.current_room}:\\n\"\n",
    "        for item, available in room['items'].items():\n",
    "            if available:\n",
    "                description += f\"- {item}\\n\"\n",
    "        if room['is_door_locked']:\n",
    "            description += \"There is a locked door here.\\n\"\n",
    "        else:\n",
    "            description += \"There is an unlocked door here.\\n\"\n",
    "        return description\n",
    "\n",
    "    def pick_up(self, item):\n",
    "        item = item.lower()\n",
    "        room = self.rooms[self.current_room]\n",
    "        if item in room['items'] and room['items'][item]:\n",
    "            self.inventory.append(item)\n",
    "            room['items'][item] = False\n",
    "            return f\"You picked up the {item}.\"\n",
    "        else:\n",
    "            return \"There's no such item here or you already picked it up.\"\n",
    "\n",
    "    def unlock_door(self):\n",
    "        room = self.rooms[self.current_room]\n",
    "        if 'key' in self.inventory:\n",
    "            if room['is_door_locked']:\n",
    "                room['is_door_locked'] = False\n",
    "                return \"You unlocked the door. You can try to enter the next room now.\"\n",
    "            else:\n",
    "                return \"The door is already unlocked.\"\n",
    "        else:\n",
    "            return \"You need a key to unlock the door.\"\n",
    "\n",
    "    def enter_next_room(self):\n",
    "        if self.current_room == 'Room 1' and not self.rooms['Room 1']['is_door_locked']:\n",
    "            self.current_room = 'Room 2'\n",
    "            return \"You enter the next room.\"\n",
    "        elif self.current_room == 'Room 2' and not self.rooms['Room 2']['is_door_locked']:\n",
    "            return self.exit_room()\n",
    "        else:\n",
    "            return \"The door is locked. You can't enter the next room yet.\"\n",
    "\n",
    "    def exit_room(self):\n",
    "        if self.current_room == 'Room 2':\n",
    "            self.game_over = True\n",
    "            self.win = True\n",
    "            return \"You exit the room and win the game!\"\n",
    "        else:\n",
    "            return \"You can't exit from here.\"\n",
    "\n",
    "    def drink_potion(self):\n",
    "        if 'potion' in self.inventory:\n",
    "            self.inventory.remove('potion')\n",
    "            return \"You drink the potion. You feel stronger!\"\n",
    "        else:\n",
    "            return \"You don't have a potion to drink.\"\n",
    "\n",
    "    def info(self):\n",
    "        state = \"Current Room State:\\n\"\n",
    "        state += self.look_around()\n",
    "        state += \"\\nYour inventory: \" + \", \".join(self.inventory) + \"\\n\"\n",
    "        state += \"\\nAvailable actions:\\n\"\n",
    "        for action in self.actions:\n",
    "            state += f\"- {action}\\n\"\n",
    "        return state\n",
    "\n",
    "    def play(self, action):\n",
    "        if self.game_over:\n",
    "            return \"Game over. Better luck next time!\"\n",
    "\n",
    "        if action in self.actions:\n",
    "            return self.actions[action]()\n",
    "        else:\n",
    "            return \"You can't do that right now.\"\n",
    "\n",
    "# Example usage\n",
    "# escape_room = EscapeRoom()\n",
    "# print(escape_room.info())\n",
    "# print(escape_room.play('look around'))\n",
    "# print(escape_room.play('pick up key'))\n",
    "# print(escape_room.play('unlock door'))\n",
    "# print(escape_room.play('enter next room'))\n",
    "# print(escape_room.play('look around'))\n",
    "# print(escape_room.play('pick up treasure'))\n",
    "# print(escape_room.play('exit'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5f8ea770-9620-4f20-a3e7-d880c759b91a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = Agent('Escape Room Solver', \n",
    "'''You are an Agent meant to solve the Escape Room Game.\n",
    "You must output one of the Available Actions at each time step.\n",
    "Output one of the Available Actions exactly.\n",
    "If previous action fails, try another.''',\n",
    "llm = llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "401cd4c1-446b-40fb-b0de-f1696c2ceb7a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "conversable_agent = ConversationWrapper(agent, \n",
    "                persistent_memory = {'Inventory': '', 'List of Past Actions': '', 'Summary of Past Actions' : '', 'List of Rooms Explored': ''},  \n",
    "                person = 'Game')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "f852b7e4-2998-4bf8-b147-f55ca12d78cf",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Game: \n",
      "Current Room State:\n",
      "You look around and see the following items in Room 1:\n",
      "- key\n",
      "- potion\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: \n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I need to take an action that will help me progress in the escape room. Picking up the key seems like a logical first step since it might be useful for unlocking the door.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key', 'Summary of Past Actions': 'Picked up the key from Room 1.', 'List of Rooms Explored': 'Room 1'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I am currently in Room 1 with a key and a potion, and there is a locked door. I have decided to pick up the key to potentially unlock the door.\u001b[0m\n",
      "Escape Room Solver : pick up key\n",
      "###\n",
      "\n",
      "Game: You picked up the key.\n",
      "Current Room State:\n",
      "You look around and see the following items in Room 1:\n",
      "- potion\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: key\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have successfully picked up the key, and now I need to use it to unlock the door in Room 1.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key, unlock door', 'Summary of Past Actions': 'Picked up the key from Room 1 and unlocked the door.', 'List of Rooms Explored': 'Room 1'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I picked up the key in Room 1 and now I am attempting to unlock the locked door with it.\u001b[0m\n",
      "Escape Room Solver : unlock door\n",
      "###\n",
      "\n",
      "Game: You unlocked the door. You can try to enter the next room now.\n",
      "Current Room State:\n",
      "You look around and see the following items in Room 1:\n",
      "- potion\n",
      "There is an unlocked door here.\n",
      "\n",
      "Your inventory: key\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have successfully unlocked the door, and now I need to decide whether to enter the next room or look around first.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key', 'List of Past Actions': 'pick up key, unlock door, enter next room', 'Summary of Past Actions': 'Picked up the key from Room 1, unlocked the door, and entered the next room.', 'List of Rooms Explored': 'Room 1, Room 2'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I picked up the key in Room 1, unlocked the door, and am now preparing to enter the next room.\u001b[0m\n",
      "Escape Room Solver : enter next room\n",
      "###\n",
      "\n",
      "Game: You enter the next room.\n",
      "Current Room State:\n",
      "You look around and see the following items in Room 2:\n",
      "- treasure\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: key\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have just entered Room 2 and see a treasure and a locked door. I should consider picking up the treasure or unlocking the door next.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key, treasure', 'List of Past Actions': 'pick up key, unlock door, enter next room, pick up treasure', 'Summary of Past Actions': 'Picked up the key from Room 1, unlocked the door, entered the next room, and picked up the treasure in Room 2.', 'List of Rooms Explored': 'Room 1, Room 2'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I entered the next room (Room 2) where I found a treasure and a locked door, and I decided to pick up the treasure.\u001b[0m\n",
      "Escape Room Solver : pick up treasure\n",
      "###\n",
      "\n",
      "Game: You picked up the treasure.\n",
      "Current Room State:\n",
      "You look around and see the following items in Room 2:\n",
      "There is a locked door here.\n",
      "\n",
      "Your inventory: key, treasure\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have picked up the treasure in Room 2, and now I need to decide on my next action. Since there is a locked door, I should attempt to unlock it to proceed further.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key, treasure', 'List of Past Actions': 'pick up key, unlock door, enter next room, pick up treasure, unlock door', 'Summary of Past Actions': 'Picked up the key from Room 1, unlocked the door, entered the next room, picked up the treasure in Room 2, and unlocked the door again.', 'List of Rooms Explored': 'Room 1, Room 2'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I entered Room 2, picked up the treasure, and now I am attempting to unlock the door to continue my escape.\u001b[0m\n",
      "Escape Room Solver : unlock door\n",
      "###\n",
      "\n",
      "Game: You unlocked the door. You can try to enter the next room now.\n",
      "Current Room State:\n",
      "You look around and see the following items in Room 2:\n",
      "There is an unlocked door here.\n",
      "\n",
      "Your inventory: key, treasure\n",
      "\n",
      "Available actions:\n",
      "- look around\n",
      "- pick up key\n",
      "- pick up lock\n",
      "- pick up potion\n",
      "- pick up treasure\n",
      "- unlock door\n",
      "- enter next room\n",
      "- exit\n",
      "- drink potion\n",
      "\n",
      "What would you do? You must reply with one of the Available actions\n",
      "\u001b[1m\u001b[32mThoughts: I have successfully unlocked the door, and now I need to decide whether to enter the next room or look around first.\u001b[0m\n",
      "\u001b[1m\u001b[34mPersistent Memory: {'Inventory': 'key, treasure', 'List of Past Actions': 'pick up key, unlock door, enter next room, pick up treasure, unlock door, enter next room', 'Summary of Past Actions': 'Picked up the key from Room 1, unlocked the door, entered the next room, picked up the treasure in Room 2, unlocked the door again, and entered the next room.', 'List of Rooms Explored': 'Room 1, Room 2, Room 3'}\u001b[0m\n",
      "\u001b[1m\u001b[35mSummary of Conversation: I have picked up the treasure in Room 2 and successfully unlocked the door. Now, I am ready to enter the next room to continue my escape.\u001b[0m\n",
      "Escape Room Solver : enter next room\n",
      "###\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# To play the game\n",
    "game = EscapeRoom()\n",
    "newstate = ''\n",
    "\n",
    "while not game.game_over:\n",
    "    state = newstate + '\\n' + game.info() + '\\nWhat would you do? You must reply with one of the Available actions'\n",
    "    print('Game:', state)\n",
    "    reply = conversable_agent.chat(state)\n",
    "    print(conversable_agent.agent_name, ':', reply)\n",
    "    print('###\\n')\n",
    "    newstate = game.play(reply)"
   ]
  }
 ],
 "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
}