{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cac06555-9ce8-4f01-bbef-3f8407f4b54d",
   "metadata": {},
   "source": [
    "# Customizing Memory Management \n",
    "This tutorial goes over how to implement a custom memory class in Letta, which allows you to customize how memory is organized (via `Block` objects) and also how memory is maintained (through memory editing tools). \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad3a8cc-d17a-4da1-b621-ecc93c9e2106",
   "metadata": {},
   "source": [
    "## Section 0: Setup a MemGPT client "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7ccd43f2-164b-4d25-8465-894a3bb54c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta import create_client \n",
    "\n",
    "client = create_client() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a28e38a-7dbe-4530-8260-202322a8458e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta import LLMConfig, EmbeddingConfig\n",
    "\n",
    "client.set_default_llm_config(LLMConfig.default_config(\"gpt-4o-mini\")) \n",
    "client.set_default_embedding_config(EmbeddingConfig.default_config(\"text-embedding-ada-002\")) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65bf0dc2-d1ac-4d4c-8674-f3156eeb611d",
   "metadata": {},
   "source": [
    "## Section 1: Memory Blocks \n",
    "Core memory consists of multiple memory *blocks*. A block represents a section of the LLM's context window, reservered to store the block's value (with an associated character limit). Blocks are persisted in the DB, so can be re-used or also shared accross agents. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce43919c-bd54-4da7-9b19-2e5a3f6bb66a",
   "metadata": {},
   "source": [
    "## Understanding `ChatMemory`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a0c20727-89b8-4820-88bc-a7daa79be1d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta import ChatMemory "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5a41d77a-dcf2-445a-bdb9-16012b752510",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_memory = ChatMemory(\n",
    "    human=\"Name: Bob\", \n",
    "    persona=\"You are a helpful assistant\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fbda842-0f66-4afb-b4d7-c65b9fe4c87e",
   "metadata": {},
   "source": [
    "#### Memory blocks \n",
    "A memory class consists of a list of `Block` objects (labeled with a block name), as well as function definitions to edit these blocks. These blocks each represent a section of the context window reserved for memory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f66c25e6-d119-49af-a972-723f4c0c4415",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Block(value='You are a helpful assistant', limit=2000, template_name=None, template=False, label='persona', description=None, metadata_={}, user_id=None, id='block-865bef7d-ab60-4e73-a376-2f34357cfaa0'),\n",
       " Block(value='Name: Bob', limit=2000, template_name=None, template=False, label='human', description=None, metadata_={}, user_id=None, id='block-45401bef-cd7c-492e-ae7e-50ab501c0c6f')]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_memory.get_blocks()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "845b027e-13de-46c6-a075-601d32f45d39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Block(value='Name: Bob', limit=2000, template_name=None, template=False, label='human', description=None, metadata_={}, user_id=None, id='block-45401bef-cd7c-492e-ae7e-50ab501c0c6f')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_memory.get_block(\"human\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "676e11d0-fad6-4683-99fe-7ae4435b617e",
   "metadata": {},
   "source": [
    "#### Memory editing functions  \n",
    "The `Memory` class also consists of functions for editing memory, which are provided as tools to the agent (so it can call them to edit memory). The `ChatMemory` class provides `core_memory_append` and `core_memory_append` functions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3472325b-46eb-46ae-8909-0d8d10168076",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4a79d810-6b48-445f-a2a1-5a5e55809581",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    def core_memory_append(self: \"Agent\", label: str, content: str) -> Optional[str]:  # type: ignore\n",
      "        \"\"\"\n",
      "        Append to the contents of core memory.\n",
      "\n",
      "        Args:\n",
      "            label (str): Section of the memory to be edited (persona or human).\n",
      "            content (str): Content to write to the memory. All unicode (including emojis) are supported.\n",
      "\n",
      "        Returns:\n",
      "            Optional[str]: None is always returned as this function does not produce a response.\n",
      "        \"\"\"\n",
      "        current_value = str(self.memory.get_block(label).value)\n",
      "        new_value = current_value + \"\\n\" + str(content)\n",
      "        self.memory.update_block_value(label=label, value=new_value)\n",
      "        return None\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(inspect.getsource(chat_memory.core_memory_append))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42f25de0-d4f9-4954-a581-ca8125e13968",
   "metadata": {},
   "source": [
    "#### Context compilation \n",
    "Each time the LLM is called (for each reasoning step of the agent), the memory is \"compiled\" into a context window representation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "34da47e1-a988-4995-afc9-e01881d36a11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'{% for block in memory.values() %}<{{ block.label }} characters=\"{{ block.value|length }}/{{ block.limit }}\">\\n{{ block.value }}\\n</{{ block.label }}>{% if not loop.last %}\\n{% endif %}{% endfor %}'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_memory.get_prompt_template()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3c71e302-11e0-4252-a3a9-65a65421f5fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<persona characters=\"27/2000\">\\nYou are a helpful assistant\\n</persona>\\n<human characters=\"9/2000\">\\nName: Bob\\n</human>'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_memory.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ec227fc-55ea-4bc2-87b9-0bc385aa5ae3",
   "metadata": {},
   "source": [
    "## Section 2: Defining a custom memory module \n",
    "In the previous example, we used a built in `ChatMemory` class which has a `human` and `persona` field in the memory to allow the agent to save important information in a 1:1 chat, and also used the `BasicBlockMemory` to customize the memory blocks. \n",
    "\n",
    "In the section, we'll go over how to define a custom memory class, including how to implement memory editing tools. We'll do this by implementing a `TaskMemory` class, which has a section of memory that is reserved for a list of tasks that can be pushed and popped form. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbdc9b6e-8bd5-4c42-970e-473da4adb2f2",
   "metadata": {},
   "source": [
    "### Defining a memory module\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7808912f-831b-4cdc-8606-40052eb809b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta import ChatMemory, Block \n",
    "from typing import Optional, List\n",
    "import json\n",
    "\n",
    "class TaskMemory(ChatMemory): \n",
    "\n",
    "    def __init__(self, human: str, persona: str, tasks: List[str]): \n",
    "        super().__init__(human=human, persona=persona, limit=2000) \n",
    "        self.link_block( \n",
    "            Block(\n",
    "                limit=2000, \n",
    "                value=json.dumps(tasks), \n",
    "                label=\"tasks\"\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def task_queue_push(self: \"Agent\", task_description: str):\n",
    "        \"\"\"\n",
    "        Push to a task queue stored in core memory. \n",
    "\n",
    "        Args:\n",
    "            task_description (str): A description of the next task you must accomplish. \n",
    "            \n",
    "        Returns:\n",
    "            Optional[str]: None is always returned as this function \n",
    "            does not produce a response.\n",
    "        \"\"\"\n",
    "        import json\n",
    "        tasks = json.loads(self.memory.get_block(\"tasks\").value)\n",
    "        tasks.append(task_description)\n",
    "        self.memory.update_block_value(\"tasks\", json.dumps(tasks))\n",
    "        return None\n",
    "\n",
    "    def task_queue_pop(self: \"Agent\"):\n",
    "        \"\"\"\n",
    "        Get the next task from the task queue \n",
    " \n",
    "        Returns:\n",
    "            Optional[str]: The description of the task popped from the \n",
    "            queue, if there are still tasks in queue. Otherwise, returns\n",
    "            None (the task queue is empty)\n",
    "        \"\"\"\n",
    "        import json\n",
    "        tasks = json.loads(self.memory.get_block(\"tasks\").value)\n",
    "        if len(tasks) == 0: \n",
    "            return None\n",
    "        task = tasks[0]\n",
    "        print(\"CURRENT TASKS: \", tasks)\n",
    "        self.memory.update_block_value(\"tasks\", json.dumps(tasks[1:]))\n",
    "        return task\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4182a134-65d2-423b-9c4b-731f55eca5aa",
   "metadata": {},
   "source": [
    "### Creating an agent with custom `TaskMemory`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "135fcf3e-59c4-4da3-b86b-dbffb21aa343",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_agent_name = \"task_agent\"\n",
    "\n",
    "task_agent_state = client.create_agent(\n",
    "    name=task_agent_name, \n",
    "    system = open(\"data/task_queue_system_prompt.txt\", \"r\").read(),\n",
    "    memory=TaskMemory(\n",
    "        human=\"My name is Sarah\", \n",
    "        persona=\"You are an agent that must clear its tasks.\", \n",
    "        tasks=[]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4de79aea-dc3d-47a3-ac7f-1f4ce399d314",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/sarahwooders/repos/letta/letta/helpers/tool_rule_solver.py:70: UserWarning: User provided tool rules and execution state resolved to no more possible tool calls.\n",
      "  warnings.warn(message)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CURRENT TASKS:  ['start calling me Charles', 'tell me a haiku about my name']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <style>\n",
       "            .message-container, .usage-container {\n",
       "                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;\n",
       "                max-width: 800px;\n",
       "                margin: 20px auto;\n",
       "                background-color: #1e1e1e;\n",
       "                border-radius: 8px;\n",
       "                overflow: hidden;\n",
       "                color: #d4d4d4;\n",
       "            }\n",
       "            .message, .usage-stats {\n",
       "                padding: 10px 15px;\n",
       "                border-bottom: 1px solid #3a3a3a;\n",
       "            }\n",
       "            .message:last-child, .usage-stats:last-child {\n",
       "                border-bottom: none;\n",
       "            }\n",
       "            .title {\n",
       "                font-weight: bold;\n",
       "                margin-bottom: 5px;\n",
       "                color: #ffffff;\n",
       "                text-transform: uppercase;\n",
       "                font-size: 0.9em;\n",
       "            }\n",
       "            .content {\n",
       "                background-color: #2d2d2d;\n",
       "                border-radius: 4px;\n",
       "                padding: 5px 10px;\n",
       "                font-family: 'Consolas', 'Courier New', monospace;\n",
       "                white-space: pre-wrap;\n",
       "            }\n",
       "            .json-key, .function-name, .json-boolean { color: #9cdcfe; }\n",
       "            .json-string { color: #ce9178; }\n",
       "            .json-number { color: #b5cea8; }\n",
       "            .internal-monologue { font-style: italic; }\n",
       "        </style>\n",
       "        <div class=\"message-container\">\n",
       "        \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">User requested two new tasks. I&#x27;ll add them one by one.</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">task_queue_push</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"task_description\"</span>: <span class=\"json-key\">\"start calling me Charles\",<br>&nbsp;&nbsp;\"request_heartbeat\"</span>: <span class=\"json-boolean\">true</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"None\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:06 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">Now I&#x27;ll add the second task about a haiku.</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">task_queue_push</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"task_description\"</span>: <span class=\"json-key\">\"tell me a haiku about my name\",<br>&nbsp;&nbsp;\"request_heartbeat\"</span>: <span class=\"json-boolean\">true</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"None\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:08 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">Tasks added. Now I&#x27;ll clear the task queue next.</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">task_queue_pop</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"request_heartbeat\"</span>: <span class=\"json-boolean\">true</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"start calling me Charles\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:09 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">I need to adjust how I address the user next.</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">send_message</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"message\"</span>: <span class=\"json-string\">\"Got it, Charles! From now on, I'll call you by that name.\"</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"None\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:11 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            </div>\n",
       "        <div class=\"usage-container\">\n",
       "            <div class=\"usage-stats\">\n",
       "                <div class=\"title\">USAGE STATISTICS</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"completion_tokens\"</span>: <span class=\"json-number\">166</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">11064</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">11230</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"step_count\"</span>: <span class=\"json-number\">4</span><br>}</div>\n",
       "            </div>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "LettaResponse(messages=[InternalMonologue(id='message-985cc68a-3276-4444-a1ad-e591ac224cf0', date=datetime.datetime(2024, 11, 7, 4, 37, 6, 166293, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"User requested two new tasks. I'll add them one by one.\"), FunctionCallMessage(id='message-985cc68a-3276-4444-a1ad-e591ac224cf0', date=datetime.datetime(2024, 11, 7, 4, 37, 6, 166293, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='task_queue_push', arguments='{\\n  \"task_description\": \"start calling me Charles\",\\n  \"request_heartbeat\": true\\n}', function_call_id='call_ejclsdkkYxd1mCeqOv2nd5PP')), FunctionReturn(id='message-a049b96e-aba0-4f84-85bb-3e644704036e', date=datetime.datetime(2024, 11, 7, 4, 37, 6, 167388, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-06 08:37:06 PM PST-0800\"\\n}', status='success', function_call_id='call_ejclsdkkYxd1mCeqOv2nd5PP'), InternalMonologue(id='message-1ef7118b-0e40-4827-8bb8-f2d828f6e14d', date=datetime.datetime(2024, 11, 7, 4, 37, 8, 830449, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"Now I'll add the second task about a haiku.\"), FunctionCallMessage(id='message-1ef7118b-0e40-4827-8bb8-f2d828f6e14d', date=datetime.datetime(2024, 11, 7, 4, 37, 8, 830449, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='task_queue_push', arguments='{\\n  \"task_description\": \"tell me a haiku about my name\",\\n  \"request_heartbeat\": true\\n}', function_call_id='call_fAUwIS8LMdIXSYl13dZMHAH5')), FunctionReturn(id='message-5dd1ecc9-2c04-40e4-8d90-a0009d43e5fe', date=datetime.datetime(2024, 11, 7, 4, 37, 8, 832851, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-06 08:37:08 PM PST-0800\"\\n}', status='success', function_call_id='call_fAUwIS8LMdIXSYl13dZMHAH5'), InternalMonologue(id='message-0687755d-180f-4399-83c3-0ac2493f7341', date=datetime.datetime(2024, 11, 7, 4, 37, 9, 840806, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"Tasks added. Now I'll clear the task queue next.\"), FunctionCallMessage(id='message-0687755d-180f-4399-83c3-0ac2493f7341', date=datetime.datetime(2024, 11, 7, 4, 37, 9, 840806, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='task_queue_pop', arguments='{\\n  \"request_heartbeat\": true\\n}', function_call_id='call_x44aL8FIGcMcJlkuO5MeYoqo')), FunctionReturn(id='message-b68af297-3d9d-451c-a073-313474a5c911', date=datetime.datetime(2024, 11, 7, 4, 37, 9, 847964, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"start calling me Charles\",\\n  \"time\": \"2024-11-06 08:37:09 PM PST-0800\"\\n}', status='success', function_call_id='call_x44aL8FIGcMcJlkuO5MeYoqo'), InternalMonologue(id='message-e7685454-3424-4d79-8294-07b2c21e911d', date=datetime.datetime(2024, 11, 7, 4, 37, 11, 76376, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='I need to adjust how I address the user next.'), FunctionCallMessage(id='message-e7685454-3424-4d79-8294-07b2c21e911d', date=datetime.datetime(2024, 11, 7, 4, 37, 11, 76376, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"Got it, Charles! From now on, I\\'ll call you by that name.\"\\n}', function_call_id='call_592kDLiCB5Rt0nY4nHFteE3r')), FunctionReturn(id='message-dab1e366-1d89-4c71-b94e-d4ae66e37402', date=datetime.datetime(2024, 11, 7, 4, 37, 11, 77104, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-06 08:37:11 PM PST-0800\"\\n}', status='success', function_call_id='call_592kDLiCB5Rt0nY4nHFteE3r')], usage=LettaUsageStatistics(completion_tokens=166, prompt_tokens=11064, total_tokens=11230, step_count=4))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.send_message(\n",
    "    agent_id=task_agent_state.id, \n",
    "    role=\"user\", \n",
    "    message=\"Add 'start calling me Charles' and 'tell me a haiku about my name' as two seperate tasks.\"\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6b54eab5-6220-4bb1-9e82-0cf21e81eb47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CURRENT TASKS:  ['tell me a haiku about my name']\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <style>\n",
       "            .message-container, .usage-container {\n",
       "                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;\n",
       "                max-width: 800px;\n",
       "                margin: 20px auto;\n",
       "                background-color: #1e1e1e;\n",
       "                border-radius: 8px;\n",
       "                overflow: hidden;\n",
       "                color: #d4d4d4;\n",
       "            }\n",
       "            .message, .usage-stats {\n",
       "                padding: 10px 15px;\n",
       "                border-bottom: 1px solid #3a3a3a;\n",
       "            }\n",
       "            .message:last-child, .usage-stats:last-child {\n",
       "                border-bottom: none;\n",
       "            }\n",
       "            .title {\n",
       "                font-weight: bold;\n",
       "                margin-bottom: 5px;\n",
       "                color: #ffffff;\n",
       "                text-transform: uppercase;\n",
       "                font-size: 0.9em;\n",
       "            }\n",
       "            .content {\n",
       "                background-color: #2d2d2d;\n",
       "                border-radius: 4px;\n",
       "                padding: 5px 10px;\n",
       "                font-family: 'Consolas', 'Courier New', monospace;\n",
       "                white-space: pre-wrap;\n",
       "            }\n",
       "            .json-key, .function-name, .json-boolean { color: #9cdcfe; }\n",
       "            .json-string { color: #ce9178; }\n",
       "            .json-number { color: #b5cea8; }\n",
       "            .internal-monologue { font-style: italic; }\n",
       "        </style>\n",
       "        <div class=\"message-container\">\n",
       "        \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">Checking next task to complete from the queue.</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">task_queue_pop</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"request_heartbeat\"</span>: <span class=\"json-boolean\">true</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"tell me a haiku about my name\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:13 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">Completing the haiku task for Charles. Here goes!</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">send_message</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"message\"</span>: <span class=\"json-string\">\"Here’s a haiku for you, Charles:\\n\\nWith strength, you embrace,\\nWhispers of your name surround,\\nCharles, calm like the sea.\"</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"None\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:14 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            </div>\n",
       "        <div class=\"usage-container\">\n",
       "            <div class=\"usage-stats\">\n",
       "                <div class=\"title\">USAGE STATISTICS</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"completion_tokens\"</span>: <span class=\"json-number\">96</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">6409</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">6505</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"step_count\"</span>: <span class=\"json-number\">2</span><br>}</div>\n",
       "            </div>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "LettaResponse(messages=[InternalMonologue(id='message-27ca6f0a-0751-4090-aac6-68ae38f5ad35', date=datetime.datetime(2024, 11, 7, 4, 37, 13, 96373, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='Checking next task to complete from the queue.'), FunctionCallMessage(id='message-27ca6f0a-0751-4090-aac6-68ae38f5ad35', date=datetime.datetime(2024, 11, 7, 4, 37, 13, 96373, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='task_queue_pop', arguments='{\\n  \"request_heartbeat\": true\\n}', function_call_id='call_UzwHlQkPuyQUyBecvU5cVvab')), FunctionReturn(id='message-f8cecbb0-bdf1-46c3-8d2e-9cfe35fd392e', date=datetime.datetime(2024, 11, 7, 4, 37, 13, 102275, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"tell me a haiku about my name\",\\n  \"time\": \"2024-11-06 08:37:13 PM PST-0800\"\\n}', status='success', function_call_id='call_UzwHlQkPuyQUyBecvU5cVvab'), InternalMonologue(id='message-e4e3dfcc-56c6-437e-8e1f-4e14eb1e7548', date=datetime.datetime(2024, 11, 7, 4, 37, 14, 444273, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='Completing the haiku task for Charles. Here goes!'), FunctionCallMessage(id='message-e4e3dfcc-56c6-437e-8e1f-4e14eb1e7548', date=datetime.datetime(2024, 11, 7, 4, 37, 14, 444273, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"Here’s a haiku for you, Charles:\\\\n\\\\nWith strength, you embrace,\\\\nWhispers of your name surround,\\\\nCharles, calm like the sea.\"\\n}', function_call_id='call_ykUJiQpCi0zLe4XPuQyWPvJ3')), FunctionReturn(id='message-65f2a5de-5577-47e3-994a-e3dc17bd7fdc', date=datetime.datetime(2024, 11, 7, 4, 37, 14, 447766, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-06 08:37:14 PM PST-0800\"\\n}', status='success', function_call_id='call_ykUJiQpCi0zLe4XPuQyWPvJ3')], usage=LettaUsageStatistics(completion_tokens=96, prompt_tokens=6409, total_tokens=6505, step_count=2))"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.send_message(\n",
    "    agent_id=task_agent_state.id, \n",
    "    role=\"user\", \n",
    "    message=\"complete your tasks\"\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "b104fe56-4ff3-439f-9e2b-1e2d24261be0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "        <style>\n",
       "            .message-container, .usage-container {\n",
       "                font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;\n",
       "                max-width: 800px;\n",
       "                margin: 20px auto;\n",
       "                background-color: #1e1e1e;\n",
       "                border-radius: 8px;\n",
       "                overflow: hidden;\n",
       "                color: #d4d4d4;\n",
       "            }\n",
       "            .message, .usage-stats {\n",
       "                padding: 10px 15px;\n",
       "                border-bottom: 1px solid #3a3a3a;\n",
       "            }\n",
       "            .message:last-child, .usage-stats:last-child {\n",
       "                border-bottom: none;\n",
       "            }\n",
       "            .title {\n",
       "                font-weight: bold;\n",
       "                margin-bottom: 5px;\n",
       "                color: #ffffff;\n",
       "                text-transform: uppercase;\n",
       "                font-size: 0.9em;\n",
       "            }\n",
       "            .content {\n",
       "                background-color: #2d2d2d;\n",
       "                border-radius: 4px;\n",
       "                padding: 5px 10px;\n",
       "                font-family: 'Consolas', 'Courier New', monospace;\n",
       "                white-space: pre-wrap;\n",
       "            }\n",
       "            .json-key, .function-name, .json-boolean { color: #9cdcfe; }\n",
       "            .json-string { color: #ce9178; }\n",
       "            .json-number { color: #b5cea8; }\n",
       "            .internal-monologue { font-style: italic; }\n",
       "        </style>\n",
       "        <div class=\"message-container\">\n",
       "        \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">INTERNAL MONOLOGUE</div>\n",
       "                <div class=\"content\"><span class=\"internal-monologue\">Let&#x27;s keep the conversation flowing!</span></div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION CALL</div>\n",
       "                <div class=\"content\"><span class=\"function-name\">send_message</span>({<br>&nbsp;&nbsp;<span class=\"json-key\">\"message\"</span>: <span class=\"json-string\">\"What would you like to talk about next, Charles? I’m all ears!\"</span><br>})</div>\n",
       "            </div>\n",
       "            \n",
       "            <div class=\"message\">\n",
       "                <div class=\"title\">FUNCTION RETURN</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"status\"</span>: <span class=\"json-key\">\"OK\",<br>&nbsp;&nbsp;\"message\"</span>: <span class=\"json-key\">\"None\",<br>&nbsp;&nbsp;\"time\"</span>: <span class=\"json-string\">\"2024-11-06 08:37:16 PM PST-0800\"</span><br>}</div>\n",
       "            </div>\n",
       "            </div>\n",
       "        <div class=\"usage-container\">\n",
       "            <div class=\"usage-stats\">\n",
       "                <div class=\"title\">USAGE STATISTICS</div>\n",
       "                <div class=\"content\">{<br>&nbsp;&nbsp;<span class=\"json-key\">\"completion_tokens\"</span>: <span class=\"json-number\">41</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">3431</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">3472</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"step_count\"</span>: <span class=\"json-number\">1</span><br>}</div>\n",
       "            </div>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "LettaResponse(messages=[InternalMonologue(id='message-b9ff3fe6-4850-40f6-9c75-c37b0bad100f', date=datetime.datetime(2024, 11, 7, 4, 37, 16, 138739, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"Let's keep the conversation flowing!\"), FunctionCallMessage(id='message-b9ff3fe6-4850-40f6-9c75-c37b0bad100f', date=datetime.datetime(2024, 11, 7, 4, 37, 16, 138739, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"What would you like to talk about next, Charles? I’m all ears!\"\\n}', function_call_id='call_ZTuH5CZlz6At9Y1ltVBttNNj')), FunctionReturn(id='message-c3b31082-ce78-42ba-9434-7edec821c3dc', date=datetime.datetime(2024, 11, 7, 4, 37, 16, 146847, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-06 08:37:16 PM PST-0800\"\\n}', status='success', function_call_id='call_ZTuH5CZlz6At9Y1ltVBttNNj')], usage=LettaUsageStatistics(completion_tokens=41, prompt_tokens=3431, total_tokens=3472, step_count=1))"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.send_message(\n",
    "    agent_id=task_agent_state.id, \n",
    "    role=\"user\", \n",
    "    message=\"keep going\"\\\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "bfac7677-5136-4a2d-8ce3-08cb3d4dfd8a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Block(value='[]', limit=2000, template_name=None, template=False, label='tasks', description=None, metadata_={}, user_id=None, id='block-288d04a9-e5c3-4da8-8746-89a728130b9a')"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.get_in_context_memory(task_agent_state.id).get_block(\"tasks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfb41f81-26e0-4bb7-8a49-b90a2e8b9ec6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "letta",
   "language": "python",
   "name": "letta"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}