{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cac06555-9ce8-4f01-bbef-3f8407f4b54d",
   "metadata": {},
   "source": [
    "# Customizing Memory Management \n",
    "\n",
    "> Make sure you run the Letta server before running this example using `letta server`\n",
    "\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": 1,
   "id": "7ccd43f2-164b-4d25-8465-894a3bb54c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta_client import CreateBlock, Letta, MessageCreate\n",
    "\n",
    "client = Letta(base_url=\"http://localhost:8283\")"
   ]
  },
  {
   "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": 3,
   "id": "a0c20727-89b8-4820-88bc-a7daa79be1d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from letta_client import ChatMemory "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5a41d77a-dcf2-445a-bdb9-16012b752510",
   "metadata": {},
   "outputs": [],
   "source": [
    "human_memory_block = client.blocks.create(\n",
    "    label=\"human\",\n",
    "    value=\"Name: Bob\",\n",
    ")\n",
    "persona_memory_block = client.blocks.create(\n",
    "    label=\"persona\",\n",
    "    value=\"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": 5,
   "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-92112694-b5ab-4210-9af6-ccb9acad3456'),\n",
       " Block(value='Name: Bob', limit=2000, template_name=None, template=False, label='human', description=None, metadata_={}, user_id=None, id='block-776d96df-7c07-4db1-b76a-1a8f1879c358')]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.blocks.list()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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-776d96df-7c07-4db1-b76a-1a8f1879c358')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.blocks.list(label=\"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": 7,
   "id": "3472325b-46eb-46ae-8909-0d8d10168076",
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "from letta.functions.function_sets.base import core_memory_append"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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(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": 9,
   "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": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_memory.get_prompt_template()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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": 10,
     "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 task related tools\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7808912f-831b-4cdc-8606-40052eb809b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, List, TYPE_CHECKING\n",
    "import json\n",
    "\n",
    "if TYPE_CHECKING:\n",
    "    from letta import AgentState\n",
    "\n",
    "def task_queue_push(agent_state: \"AgentState\", 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(agent_state.memory.get_block(\"tasks\").value)\n",
    "    tasks.append(task_description)\n",
    "    agent_state.memory.update_block_value(\"tasks\", json.dumps(tasks))\n",
    "    return None\n",
    "\n",
    "def task_queue_pop(agent_state: \"AgentState\"):\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(agent_state.memory.get_block(\"tasks\").value)\n",
    "    if len(tasks) == 0: \n",
    "        return None\n",
    "    task = tasks[0]\n",
    "    print(\"CURRENT TASKS: \", tasks)\n",
    "    agent_state.memory.update_block_value(\"tasks\", json.dumps(tasks[1:]))\n",
    "    return task\n",
    "\n",
    "push_task_tool = client.tools.upsert_from_function(func=task_queue_push)\n",
    "pop_task_tool = client.tools.upsert_from_function(func=task_queue_pop)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4182a134-65d2-423b-9c4b-731f55eca5aa",
   "metadata": {},
   "source": [
    "### Creating an agent with custom `TaskMemory`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135fcf3e-59c4-4da3-b86b-dbffb21aa343",
   "metadata": {},
   "outputs": [],
   "source": [
    "task_agent_name = \"task_agent\"\n",
    "\n",
    "# delete agent if exists \n",
    "agents = client.agents.list(name=task_agent_name)\n",
    "if len(agents) > 0: \n",
    "    client.agents.delete(agent_id=agents[0].id)\n",
    "\n",
    "task_agent_state = client.agents.create(\n",
    "    name=task_agent_name, \n",
    "    system = open(\"data/task_queue_system_prompt.txt\", \"r\").read(),\n",
    "    memory_blocks=[\n",
    "        CreateBlock(\n",
    "            label=\"human\",\n",
    "            value=\"My name is Sarah\",\n",
    "        ),\n",
    "        CreateBlock(\n",
    "            label=\"persona\",\n",
    "            value=\"You are an agent that must clear its tasks.\",\n",
    "        ),\n",
    "        CreateBlock(\n",
    "            label=\"tasks\",\n",
    "            value=\"[]\",\n",
    "        ),\n",
    "    ],\n",
    "    tool_ids=[push_task_tool.id, pop_task_tool.id],\n",
    "    model=\"letta/letta-free\",\n",
    "    embedding=\"letta/letta-free\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4de79aea-dc3d-47a3-ac7f-1f4ce399d314",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CURRENT TASKS:  ['start calling me Charles', 'tell me a haiku about my name']\n",
      "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\">User wants to add &#x27;start calling me Charles&#x27; and a haiku about the name as tasks.</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-13 05:48:34 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 next task for the haiku about the name.</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-13 05:48:36 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 will now remove the first task from the queue: &#x27;start calling me Charles&#x27;.</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-13 05:48:37 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\">Next, I will complete the task about the 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_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-13 05:48:40 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\">Task queue is empty now. Ready to respond and complete the haiku request!</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\">\"Charles, a strong name\\nWhispers of noble echoes\\nStrength in every step.\"</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-13 05:48:41 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\">224</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">14235</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">14459</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"step_count\"</span>: <span class=\"json-number\">5</span><br>}</div>\n",
       "            </div>\n",
       "        </div>\n",
       "        "
      ],
      "text/plain": [
       "LettaResponse(messages=[InternalMonologue(id='message-34a1bb2c-3bc4-4269-8f76-c9888f18c435', date=datetime.datetime(2024, 11, 14, 1, 48, 34, 670884, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"User wants to add 'start calling me Charles' and a haiku about the name as tasks.\"), FunctionCallMessage(id='message-34a1bb2c-3bc4-4269-8f76-c9888f18c435', date=datetime.datetime(2024, 11, 14, 1, 48, 34, 670884, 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_zOqq1dOBwpO1j5j1f0ch1zU2')), FunctionReturn(id='message-6934a04d-0e93-450f-9a0f-139f8022bbbe', date=datetime.datetime(2024, 11, 14, 1, 48, 34, 672396, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-13 05:48:34 PM PST-0800\"\\n}', status='success', function_call_id='call_zOqq1dOBwpO1j5j1f0ch1zU2'), InternalMonologue(id='message-66c68a60-bd23-4659-95da-a3e25bb7883e', date=datetime.datetime(2024, 11, 14, 1, 48, 36, 394958, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"Now I'll add the next task for the haiku about the name.\"), FunctionCallMessage(id='message-66c68a60-bd23-4659-95da-a3e25bb7883e', date=datetime.datetime(2024, 11, 14, 1, 48, 36, 394958, 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_6fklGb62YHrXKtcYcgHseLpv')), FunctionReturn(id='message-28a1802b-1474-456f-b5ca-c706fd50f1fc', date=datetime.datetime(2024, 11, 14, 1, 48, 36, 396303, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-13 05:48:36 PM PST-0800\"\\n}', status='success', function_call_id='call_6fklGb62YHrXKtcYcgHseLpv'), InternalMonologue(id='message-8bf666a4-5ca1-4b76-b625-27410cefe2b3', date=datetime.datetime(2024, 11, 14, 1, 48, 37, 549545, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue=\"I will now remove the first task from the queue: 'start calling me Charles'.\"), FunctionCallMessage(id='message-8bf666a4-5ca1-4b76-b625-27410cefe2b3', date=datetime.datetime(2024, 11, 14, 1, 48, 37, 549545, 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_p28SUN7cOlgXV6tyGUtGkczG')), FunctionReturn(id='message-f19be3d8-1df2-4ac5-a134-9e6f04a8b93e', date=datetime.datetime(2024, 11, 14, 1, 48, 37, 553595, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"start calling me Charles\",\\n  \"time\": \"2024-11-13 05:48:37 PM PST-0800\"\\n}', status='success', function_call_id='call_p28SUN7cOlgXV6tyGUtGkczG'), InternalMonologue(id='message-d81b056d-69f2-49e9-9448-97d39c31fd8e', date=datetime.datetime(2024, 11, 14, 1, 48, 40, 191574, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='Next, I will complete the task about the haiku.'), FunctionCallMessage(id='message-d81b056d-69f2-49e9-9448-97d39c31fd8e', date=datetime.datetime(2024, 11, 14, 1, 48, 40, 191574, 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_bfl2RvzYj0zrpgiIzRYF8Wgc')), FunctionReturn(id='message-ac09ca1e-0cee-4260-8fe6-9fce1978f49e', date=datetime.datetime(2024, 11, 14, 1, 48, 40, 196240, 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-13 05:48:40 PM PST-0800\"\\n}', status='success', function_call_id='call_bfl2RvzYj0zrpgiIzRYF8Wgc'), InternalMonologue(id='message-be9151a5-ba67-4816-8c5f-bd3346b73756', date=datetime.datetime(2024, 11, 14, 1, 48, 41, 855182, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='Task queue is empty now. Ready to respond and complete the haiku request!'), FunctionCallMessage(id='message-be9151a5-ba67-4816-8c5f-bd3346b73756', date=datetime.datetime(2024, 11, 14, 1, 48, 41, 855182, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"Charles, a strong name\\\\nWhispers of noble echoes\\\\nStrength in every step.\"\\n}', function_call_id='call_37cVdqCSCfa3XzmrMvmAnPCM')), FunctionReturn(id='message-3b21e720-67ec-4e02-a4d5-533945cf896b', date=datetime.datetime(2024, 11, 14, 1, 48, 41, 856185, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-13 05:48:41 PM PST-0800\"\\n}', status='success', function_call_id='call_37cVdqCSCfa3XzmrMvmAnPCM')], usage=LettaUsageStatistics(completion_tokens=224, prompt_tokens=14235, total_tokens=14459, step_count=5))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.agents.messages.create(\n",
    "    agent_id=task_agent_state.id, \n",
    "    messages=[\n",
    "        MessageCreate(\n",
    "            role=\"user\",\n",
    "            content=\"Add 'start calling me Charles' and 'tell me a haiku about my name' as two separate tasks.\",\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6b54eab5-6220-4bb1-9e82-0cf21e81eb47",
   "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\">User is confused. I need to clarify that tasks are already completed.</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\">\"I've completed all your tasks, Charles! If there's anything else you'd like to do or ask, just let me know!\"</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-13 05:48:43 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\">56</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">3297</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">3353</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-3e24b340-977d-433d-a8fd-05b916bcf67f', date=datetime.datetime(2024, 11, 14, 1, 48, 43, 388438, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='User is confused. I need to clarify that tasks are already completed.'), FunctionCallMessage(id='message-3e24b340-977d-433d-a8fd-05b916bcf67f', date=datetime.datetime(2024, 11, 14, 1, 48, 43, 388438, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"I\\'ve completed all your tasks, Charles! If there\\'s anything else you\\'d like to do or ask, just let me know!\"\\n}', function_call_id='call_Leb5MXlO15Yn7V715O5Pb3Q0')), FunctionReturn(id='message-e5aeb5c8-c1c9-40b6-87cf-92ff33b61020', date=datetime.datetime(2024, 11, 14, 1, 48, 43, 389280, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-13 05:48:43 PM PST-0800\"\\n}', status='success', function_call_id='call_Leb5MXlO15Yn7V715O5Pb3Q0')], usage=LettaUsageStatistics(completion_tokens=56, prompt_tokens=3297, total_tokens=3353, step_count=1))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.agents.messages.create(\n",
    "    agent_id=task_agent_state.id, \n",
    "    messages=[\n",
    "        MessageCreate(\n",
    "            role=\"user\",\n",
    "            content=\"complete your tasks\",\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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\">User wants to keep the conversation going. Maybe I could ask a question or suggest something fun to talk about.</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\">\"Sure! What would you like to chat about next? We can dive into hobbies, favorite books, or whatever's on your mind!\"</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-13 05:48:45 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\">67</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"prompt_tokens\"</span>: <span class=\"json-number\">3446</span>,<br>&nbsp;&nbsp;<span class=\"json-key\">\"total_tokens\"</span>: <span class=\"json-number\">3513</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-67635cfd-bf4b-4025-a67c-3061c1b78651', date=datetime.datetime(2024, 11, 14, 1, 48, 45, 923304, tzinfo=datetime.timezone.utc), message_type='internal_monologue', internal_monologue='User wants to keep the conversation going. Maybe I could ask a question or suggest something fun to talk about.'), FunctionCallMessage(id='message-67635cfd-bf4b-4025-a67c-3061c1b78651', date=datetime.datetime(2024, 11, 14, 1, 48, 45, 923304, tzinfo=datetime.timezone.utc), message_type='function_call', function_call=FunctionCall(name='send_message', arguments='{\\n  \"message\": \"Sure! What would you like to chat about next? We can dive into hobbies, favorite books, or whatever\\'s on your mind!\"\\n}', function_call_id='call_pM4j4LZDovPvOwk4Up4xlsnG')), FunctionReturn(id='message-e6f02189-b330-4ad6-b427-52f143791d8d', date=datetime.datetime(2024, 11, 14, 1, 48, 45, 924171, tzinfo=datetime.timezone.utc), message_type='function_return', function_return='{\\n  \"status\": \"OK\",\\n  \"message\": \"None\",\\n  \"time\": \"2024-11-13 05:48:45 PM PST-0800\"\\n}', status='success', function_call_id='call_pM4j4LZDovPvOwk4Up4xlsnG')], usage=LettaUsageStatistics(completion_tokens=67, prompt_tokens=3446, total_tokens=3513, step_count=1))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = client.agents.messages.create(\n",
    "    agent_id=task_agent_state.id, \n",
    "    messages=[\n",
    "        MessageCreate(\n",
    "            role=\"user\",\n",
    "            content=\"keep going\",\n",
    "        )\n",
    "    ],\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "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-406ae267-2b00-4ff5-8df5-38c73ca88e45')"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.agents.core_memory.retrieve_block(agent_id=task_agent_state.id, block_label=\"tasks\")"
   ]
  }
 ],
 "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
}