From f40c35fd149965dc50e45c9f45b8813853fd2bce Mon Sep 17 00:00:00 2001
From: Siddharth Balyan <balyan.sid@gmail.com>
Date: Fri, 3 May 2024 14:35:07 +0530
Subject: [PATCH] summarization step added.

---
 agents-api/notebooks/03-summarise.ipynb | 762 ++++++++++++++++++++++--
 1 file changed, 724 insertions(+), 38 deletions(-)

diff --git a/agents-api/notebooks/03-summarise.ipynb b/agents-api/notebooks/03-summarise.ipynb
index 88357151d..a934fd1b9 100644
--- a/agents-api/notebooks/03-summarise.ipynb
+++ b/agents-api/notebooks/03-summarise.ipynb
@@ -33,6 +33,13 @@
     "- REMOVE: Safely remove messages from the session"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Helper functions and Setup"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -50,20 +57,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "make_chatml = lambda content, role=\"system\", name=None, **_: {\n",
+    "make_chatml = lambda content, role=\"system\", name=None, index=None, **_: {\n",
     "    key: value\n",
     "    for key, value in dict(role=role, name=name, content=content).items()\n",
     "    if value is not None\n",
     "}\n",
     "\n",
-    "user = lambda content, name=None: make_chatml(role=\"user\", content=content, name=name)\n",
-    "assistant = lambda content, name=None: make_chatml(\n",
-    "    role=\"assistant\", content=content, name=name\n",
+    "user = lambda content, name=None, index=None: make_chatml(\n",
+    "    role=\"user\", content=content, name=name, index=index\n",
+    ")\n",
+    "assistant = lambda content, name=None, index=None: make_chatml(\n",
+    "    role=\"assistant\", content=content, name=name, index=index\n",
+    ")\n",
+    "system = lambda content, name=None, index=None: make_chatml(\n",
+    "    content, name=name, index=index\n",
     ")\n",
-    "system = lambda content, name=None: make_chatml(content, name=name)\n",
     "thought = lambda content, name=None: make_chatml(content, name=\"thought\")\n",
     "information = lambda content: system(content, name=\"information\")\n",
-    "summary = lambda content: system(content, name=\"summary\")\n",
+    "summary = lambda content, index=None: system(content, name=\"summary\", index=index)\n",
     "entities = lambda content: system(content, name=\"entity\")"
    ]
   },
@@ -89,7 +100,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -104,9 +115,16 @@
     "        chat_session.append(make_chatml(**result))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## One-Shot Example Chat"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -292,12 +310,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Entity CoT"
+    "## Entity Extraction One-Shot Example"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -323,12 +341,21 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Trim CoT"
+    "## Trim Messages One-Shot Example"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "*TODO*\n",
+    "- only output the index and the content of the message to be changed.\n",
+    "- [!] Index was not in the output\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -344,14 +371,18 @@
     "\n",
     "example_trimmed_messages = [\n",
     "    *example_chat[0:8],\n",
-    "    assistant(\"Haha accurate. He does push boundaries and keep things interesting!\"),\n",
+    "    assistant(\n",
+    "        \"Haha accurate. He does push boundaries and keep things interesting!\", index=8\n",
+    "    ),\n",
     "    *example_chat[9:10],\n",
     "    assistant(\n",
-    "        \"For sure, he's definitely made a huge impact in different industries. Excited to see what he'll come up with next\"\n",
+    "        \"For sure, he's definitely made a huge impact in different industries. Excited to see what he'll come up with next\",\n",
+    "        index=10,\n",
     "    ),\n",
     "    *example_chat[11:18],\n",
     "    assistant(\n",
-    "        \"Glad I could make you laugh! It's always nice to share a joke or two. What else is up?\"\n",
+    "        \"Glad I could make you laugh! It's always nice to share a joke or two. What else is up?\",\n",
+    "        index=18,\n",
     "    ),\n",
     "    *example_chat[19:21],\n",
     "    user(\n",
@@ -368,7 +399,8 @@
     "- Though a tinge of guilt lingers for this routine coming to an end, akin to secretly devouring ice cream without my younger sibling\n",
     "- For my next foray, I'm resolute about \"buying an ice cream\" and sharing the entire journey with them this time around\n",
     "</trim>\n",
-    "\"\"\"\n",
+    "\"\"\",\n",
+    "        index=21,\n",
     "    ),\n",
     "    *example_chat[22:34],\n",
     "    assistant(\n",
@@ -382,7 +414,8 @@
     "- Highly engaging and enjoyable interactions\n",
     "- Look forward to more amusing conversations!\n",
     "</trim>\n",
-    "\"\"\"\n",
+    "\"\"\",\n",
+    "        index=34,\n",
     "    ),\n",
     "]"
    ]
@@ -391,16 +424,52 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Summarise CoT"
+    "## Summarise One-Shot Example"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 47,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# TODO:"
+    "# TODO: add and optimise the Chain of Thought\n",
+    "summarize_chain_of_thought_example = \"\"\"\\\n",
+    "Planning step by step:\n",
+    "- To add context for future entries, let's add one at the top which outlines the main entities in the session.\n",
+    "- We can replace entries with index 1,2,3,4 with a summary of those messages.\n",
+    "- We can replace entries with index 5,6,7,8 similarly.\n",
+    "- It could be disruptive to remove messages with index 9-16 because that might lose the joke's context.\n",
+    "- We can safely summarize entries with index 17,18 and remove them.\n",
+    "- We can safely remove message with index 20.\n",
+    "- We should keep entry with index 21 because it's given by the user and they might ask about it again.\n",
+    "- We should keep the assistant's response in message with index 22 to keep the context.\n",
+    "- Messages with index 23-32 are repetitive and should be summarized.\n",
+    "- We should retain message with index 33 since it's a direct request from the user.\n",
+    "- We can safely summarize message with index 34's essay into just the salient points only.\n",
+    "\"\"\"\n",
+    "\n",
+    "example_summarized_messages = [\n",
+    "    summary(\n",
+    "        \"Event: Camille says hi to JaneBot and they reply that they are working and trying to stay busy.\"\n",
+    "    ),\n",
+    "    summary(\n",
+    "        \"Event: They discuss Elon Musk and agree that he can be a polarizing personality.\"\n",
+    "    ),\n",
+    "    *example_chat[9:17],\n",
+    "    summary(\"Event: Camille appreciates JaneBot's sense of humor.\"),\n",
+    "    *example_chat[21:23],\n",
+    "    summary(\n",
+    "        'Event: Camille asks JaneBot to play \"What number I am thinking of?\" game but she keeps saying \"no\" to all guesses which JaneBot finds really funny.'\n",
+    "    ),\n",
+    "    example_chat[33],\n",
+    "    summary(\n",
+    "        \"\"\"Event: JaneBot wrote an essay about Camille. Summary of the essay:\n",
+    "1. You have a great sense of humor and enjoy playful, lighthearted banter in conversations.\n",
+    "2. Your creativity is evident in how you approach topics in unexpected ways and keep conversations interesting.\n",
+    "3. Chatting with you is enjoyable due to the casual, friendly tone you bring to interactions.\"\"\"\n",
+    "    ),\n",
+    "]"
    ]
   },
   {
@@ -409,28 +478,35 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "chat_session = [system(\"you are a friend who likes to give life advice\")]"
+    "# chat_session = [system(\"you are a friend who likes to give life advice\")]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Loading Test Chat JSON"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 19,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Load test chat from json file\n",
+    "# Test Chat File\n",
     "import json\n",
     "\n",
     "with open(\"./test-chat.json\", \"r\") as f:\n",
     "    chat_session = json.load(f)\n",
     "\n",
-    "for message, index in zip(chat_session, range(1000)):\n",
-    "    message[\"index\"] = index"
+    "for index, message in enumerate(chat_session):\n",
+    "    message[\"index\"] = index + 1"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -439,7 +515,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -456,7 +532,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -479,7 +555,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 48,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -541,7 +617,6 @@
     "\n",
     "    messages.append(cot_result)\n",
     "    start_message = \"\"\"\\\n",
-    "\n",
     "Begin! Directly write the entities as a Markdown formatted list, don't write any thoughts or delimiters.\n",
     "\n",
     "<ct:entities>\n",
@@ -550,11 +625,69 @@
     "    messages.append(user(start_message))\n",
     "\n",
     "    print(\"Starting chatml generation\")\n",
-    "    result = generate(messages, model=\"gpt-4-turbo\", temperature=0.1, stop=[\"</ct\"])\n",
+    "    entity_result = generate(\n",
+    "        messages, model=\"gpt-4-turbo\", temperature=0.1, stop=[\"</ct\"]\n",
+    "    )\n",
     "    print(\"End chatml generation\")\n",
-    "    messages.append(result)\n",
+    "    messages.append(entity_result)\n",
     "\n",
-    "    return make_chatml(**result)"
+    "    return make_chatml(**entity_result)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting CoT generation\n",
+      "{'content': 'To provide a comprehensive overview of the entities in the '\n",
+      "            'session, we need to focus on the main people, places, things, and '\n",
+      "            'concepts that have been discussed:\\n'\n",
+      "            '\\n'\n",
+      "            '1. **User**: Engages in conversation about video games, '\n",
+      "            'specifically mentioning experiences with \"Red Dead Redemption 2\" '\n",
+      "            'and \"Helldivers 2.\" Also, shares technical issues related to '\n",
+      "            'Nvidia graphics drivers on Linux.\\n'\n",
+      "            '\\n'\n",
+      "            '2. **Assistant**: Offers insights and engages in discussions '\n",
+      "            'about the video games mentioned by the user, provides gameplay '\n",
+      "            'strategies, and offers technical advice regarding Nvidia driver '\n",
+      "            'issues on Linux.\\n'\n",
+      "            '\\n'\n",
+      "            '3. **Red Dead Redemption 2 (RDR2)**: A video game discussed '\n",
+      "            'extensively in the conversation. Key aspects such as the '\n",
+      "            'character Arthur, missions like \"Blood Feuds, Ancient and '\n",
+      "            'Modern,\" and side quests are highlighted.\\n'\n",
+      "            '\\n'\n",
+      "            '4. **Helldivers 2**: Another video game mentioned by the user. '\n",
+      "            'The conversation covers gameplay elements like the laser cannon, '\n",
+      "            'guard dog, and challenges like the Charger enemy.\\n'\n",
+      "            '\\n'\n",
+      "            '5. **Nvidia Drivers**: A technical issue brought up by the user, '\n",
+      "            'particularly focusing on compatibility problems with Linux. The '\n",
+      "            'assistant provides troubleshooting advice and recommendations for '\n",
+      "            'dealing with these issues.\\n'\n",
+      "            '\\n'\n",
+      "            '6. **Linux**: Mentioned in the context of having compatibility '\n",
+      "            \"issues with Nvidia drivers, highlighting the user's struggle with \"\n",
+      "            'technical aspects of gaming on this operating system.\\n'\n",
+      "            '\\n'\n",
+      "            'These entities encapsulate the core topics and issues discussed '\n",
+      "            \"in the session, providing a clear view of the conversation's \"\n",
+      "            'focus areas.',\n",
+      " 'role': 'assistant'}\n",
+      "End CoT generation\n",
+      "Starting chatml generation\n",
+      "End chatml generation\n"
+     ]
+    }
+   ],
+   "source": [
+    "entity_info = get_entities(chat_session=chat_session)"
    ]
   },
   {
@@ -566,7 +699,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -597,7 +730,7 @@
     "</ct:example-history-for-demonstration>\n",
     "\n",
     "<ct:example-plan>\n",
-    "{entity_chain_of_thought}\n",
+    "{summarize_chain_of_thought_example}\n",
     "</ct:example-plan>\n",
     "\n",
     "<ct:example-trimmed-messages>\n",
@@ -633,11 +766,564 @@
     "    messages.append(user(start_message))\n",
     "\n",
     "    print(\"Starting chatml generation\")\n",
-    "    result = generate(messages, model=\"gpt-4-turbo\", temperature=0.1, stop=[\"</ct\"])\n",
+    "    trim_result = generate(messages, model=\"gpt-4-turbo\", temperature=0.1, stop=[\"</ct\"])\n",
     "    print(\"End chatml generation\")\n",
-    "    messages.append(result)\n",
+    "    messages.append(trim_result)\n",
     "\n",
-    "    return make_chatml(**result)"
+    "    return make_chatml(**trim_result)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting CoT generation\n",
+      "{'content': 'The conversation between the user and the assistant revolves '\n",
+      "            'around video games and technical issues. The assistant engages '\n",
+      "            'enthusiastically about the user’s gaming experiences, providing '\n",
+      "            'tips and engaging on a technical level about Nvidia drivers on '\n",
+      "            'Linux towards the end. The user is generally responsive and '\n",
+      "            'shares specifics about their gaming preferences and technical '\n",
+      "            'difficulties. The tone is casual and friendly throughout, with '\n",
+      "            'the assistant providing support and showing interest in the '\n",
+      "            \"user's activities. \\n\"\n",
+      "            '\\n'\n",
+      "            'No trimming is required as the messages are well-paced and '\n",
+      "            \"relevant to the users' interests, fostering an engaging and \"\n",
+      "            'informative dialogue.',\n",
+      " 'role': 'assistant'}\n",
+      "End CoT generation\n",
+      "Starting chatml generation\n",
+      "End chatml generation\n"
+     ]
+    }
+   ],
+   "source": [
+    "trimmed_messages = trim_messages(chat_session)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "('[\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 2,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"hey dude! whats good\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 3,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Hey! Just chilling, man. How about you? Played any cool '\n",
+      " 'games recently?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 4,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"i just got done with finishing red dead redemption 2 '\n",
+      " 'recently. the last mission was so fkng coolll\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 5,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Oh man, RDR2 is epic! That ending is intense, right? What '\n",
+      " \"did you think about Arthur's journey throughout the game? Got any favorite \"\n",
+      " 'moments?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 6,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"blood feuds mission was amazing\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 7,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Absolutely, the \\\\\"Blood Feuds, Ancient and Modern\\\\\" '\n",
+      " 'mission is a standout! The whole shootout at Braithwaite Manor is intense. '\n",
+      " 'The atmosphere Rockstar created during that mission with the burning '\n",
+      " 'building in the background is just cinematic perfection. Did you side more '\n",
+      " 'with Arthur or Dutch throughout the game?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 8,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"arthur of course. dutch just kept becoming crazier and '\n",
+      " 'crazier.\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 9,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Totally get that. Arthur\\'s character development is really '\n",
+      " \"something else. He starts questioning Dutch's decisions and looking at \"\n",
+      " 'things more critically, which adds so much depth to his character. Watching '\n",
+      " \"Dutch's descent as the gang falls apart is pretty intense. Have you thought \"\n",
+      " 'about checking out any of the side quests or exploring more of the open '\n",
+      " 'world now that you\\'ve finished the main story?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 10,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"oh i\\'ve done all the side quests. the widow missions are '\n",
+      " 'great but i also loved the bounty missions (where u have to talk to '\n",
+      " 'gangsters for an author)\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 11,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Ah, the bounty missions for the author are a blast! Hunting '\n",
+      " 'down those gunslingers for the book gives you such a great tour of the '\n",
+      " \"game's diverse characters. And Charlotte's missions, the widow, they're \"\n",
+      " 'pretty touching, right? It\\\\u2019s nice to see Arthur helping her adapt to '\n",
+      " 'wilderness life. Adds a softer side to the game amidst all the '\n",
+      " \"chaos.\\\\n\\\\nNow that you've conquered RDR2, any other games on your \"\n",
+      " 'radar?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 12,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"im playing helldivers 2. do u know of it\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 13,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Yes, Helldivers 2! It\\'s a great choice if you\\'re into '\n",
+      " 'intense co-op action. The game really ramps up with its twin-stick shooter '\n",
+      " 'gameplay, and it\\\\u2019s super fun but also quite challenging. Working '\n",
+      " 'together with teammates to complete those missions can get pretty hectic. '\n",
+      " 'How are you finding the chaos and team dynamics so far?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 14,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"amazing. my fav build is laser cannon + guard dog. what '\n",
+      " 'about u\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 15,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Nice combo! The laser cannon is awesome for tearing through '\n",
+      " 'those enemy swarms, and the guard dog drone really helps in keeping the '\n",
+      " 'smaller threats at bay. I personally love using the railgun with a shield '\n",
+      " 'backpack. It gives a good balance of offense and defense, especially when '\n",
+      " \"things get really chaotic. What's been the toughest mission for you so \"\n",
+      " 'far?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 16,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"dude the charger is so difficult to kill.\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 17,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Yeah, the Charger can be brutal! They really test your '\n",
+      " \"reflexes and teamwork. It's all about timing and coordination with your team \"\n",
+      " 'to bring those heavy hitters down. Have you tried using more strategic traps '\n",
+      " 'or mines? Sometimes setting up the right position with static fields can '\n",
+      " 'really make a difference against them. How do your squadmates handle it?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 18,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"we just run and gun\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 19,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Classic strategy! That can definitely be fun, especially in '\n",
+      " 'a game as intense as Helldivers 2 where things get chaotic fast. Keeping on '\n",
+      " \"the move helps a ton. Just make sure someone's always got your back! Have \"\n",
+      " 'you guys tried any of the higher difficulty levels yet, or are you sticking '\n",
+      " 'to a certain comfort zone for now?\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 20,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"i think i have to go work soon. also have this nvidia driver '\n",
+      " 'issue that\\'s bugging me\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 21,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"Ah, bummer about the driver issue \\\\u2014 it\\'s annoying '\n",
+      " 'when tech stuff gets in the way of gaming. Maybe try doing a clean install '\n",
+      " \"of the drivers if you haven't already? It can sometimes clear up persistent \"\n",
+      " \"issues.\\\\n\\\\nAlso, good luck with work! Hopefully, you'll get some time \"\n",
+      " 'later to sort out the driver thing and jump back into some gaming. Let me '\n",
+      " \"know how it goes, and if you need more gaming chats or tech tips, I'm \"\n",
+      " 'here!\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 22,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"nvidia and linux don\\'t go well together\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 23,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"That\\'s true, Nvidia drivers can be tricky to handle on '\n",
+      " 'Linux due to compatibility issues sometimes. You might want to check out the '\n",
+      " 'latest Nvidia driver from their official site or a stable version known to '\n",
+      " 'work well with your specific Linux distribution. Using proprietary drivers '\n",
+      " 'instead of Nouveau could also potentially solve a lot of issues. '\n",
+      " 'Additionally, forums and community support boards like those on Ubuntu, '\n",
+      " 'Arch, or even broader communities like Reddit\\\\u2019s r/linux_gaming can be '\n",
+      " 'really helpful.\\\\n\\\\nHope that helps, and your driver woes get sorted soon '\n",
+      " 'so you can enjoy your gaming hassle-free!\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 24,\\n'\n",
+      " '    \"role\": \"user\",\\n'\n",
+      " '    \"content\": \"yep. see u bye!\"\\n'\n",
+      " '  },\\n'\n",
+      " '  {\\n'\n",
+      " '    \"index\": 25,\\n'\n",
+      " '    \"role\": \"assistant\",\\n'\n",
+      " '    \"content\": \"See you, take care! Bye!\"\\n'\n",
+      " '  }\\n'\n",
+      " ']')\n"
+     ]
+    }
+   ],
+   "source": [
+    "pprint(trimmed_messages[\"content\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Step 3: Summarize Messages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from pprint import pprint\n",
+    "import time\n",
+    "\n",
+    "\n",
+    "def summarize_messages(chat_session, entity_info):\n",
+    "    _sys_msg, *history_to_compact = get_leaf_nodes(chat_session)\n",
+    "    history_to_compact = [\n",
+    "        {\"index\": i + 1, **msg} for i, msg in enumerate(history_to_compact)\n",
+    "    ]\n",
+    "    system_prompt = f\"\"\"\\\n",
+    "You are given a session history of a chat between a user and a gpt-4-turbo large language model made by OpenAI.\n",
+    "\n",
+    "Your goal is to compress the history by removing or coalescing redundant or irrelevant information in order to reduce its size and save costs. However, you must ensure that no important information is lost that might be relevant in the future.\n",
+    "\n",
+    "<ct:instructions>\n",
+    "- You will be provided with \"Entities\" which are important characters, people, places, things, objects relevant to the conversation.\n",
+    "- You will make sure to preserve the context for these \"Entities\".\n",
+    "- You may remove entries in the history that do not contain any crucial information such as small-talk or repeated questions/responses.\n",
+    "- You will combine multiple entries into a summary strictly as long as that doesn't disrupt the structure of the session.\n",
+    "- Do not remove content that the user shared if it might be relevant to future messages.\n",
+    "\n",
+    "</ct:instructions>\n",
+    "\n",
+    "<ct:example-entities>\n",
+    "{json.dumps(example_entity_messages, indent=2)}\n",
+    "</ct:example-entities>\n",
+    "\n",
+    "<ct:example-history-for-demonstration>\n",
+    "{json.dumps(example_chat, indent=2)}\n",
+    "</ct:example-history-for-demonstration>\n",
+    "\n",
+    "<ct:example-plan>\n",
+    "{summarize_chain_of_thought_example}\n",
+    "</ct:example-plan>\n",
+    "\n",
+    "<ct:example-trimmed-messages>\n",
+    "{json.dumps(example_summarized_messages, indent=2)}\n",
+    "</ct:example-trimmed-messages>\n",
+    "\"\"\".strip()\n",
+    "\n",
+    "    user_message = f\"\"\"\\\n",
+    "<ct:entities>\n",
+    "{json.dumps(entity_info)}\n",
+    "\n",
+    "<ct:history>\n",
+    "{json.dumps(history_to_compact, indent=2)}\n",
+    "</ct:history>\n",
+    "\n",
+    "<ct:plan>\n",
+    "\"\"\"\n",
+    "    messages = [system(system_prompt), user(user_message)]\n",
+    "\n",
+    "    print(\"Starting CoT generation\")\n",
+    "    cot_result = generate(messages, model=\"gpt-4-turbo\", stop=[\"</ct\"], temperature=0.7)\n",
+    "\n",
+    "    cot_result = make_chatml(**cot_result)\n",
+    "    pprint(cot_result)\n",
+    "\n",
+    "    print(\"End CoT generation\")\n",
+    "\n",
+    "    messages.append(cot_result)\n",
+    "    start_message = \"\"\"\\\n",
+    "\n",
+    "Begin! Rewrite the session history as instructed, as a JSON list, don't write any thoughts or delimiters.\n",
+    "\n",
+    "<ct:entities>\n",
+    "\"\"\".strip()\n",
+    "\n",
+    "    messages.append(user(start_message))\n",
+    "\n",
+    "    print(\"Starting chatml generation\")\n",
+    "    trim_result = generate(\n",
+    "        messages, model=\"gpt-4-turbo\", temperature=0.1, stop=[\"</ct\"]\n",
+    "    )\n",
+    "    print(\"End chatml generation\")\n",
+    "    messages.append(trim_result)\n",
+    "\n",
+    "    return make_chatml(**trim_result)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Starting CoT generation\n",
+      "{'content': '- Combine initial greetings and the start of the discussion about '\n",
+      "            '\"Red Dead Redemption 2 (RDR2)\" into a summary.\\n'\n",
+      "            '- Summarize the detailed conversation about specific missions and '\n",
+      "            'characters in \"RDR2\".\\n'\n",
+      "            '- Retain individual entries about \"Helldivers 2\" due to the '\n",
+      "            \"user's direct interaction and game preferences.\\n\"\n",
+      "            '- Summarize discussions around \"Nvidia Drivers\" and \"Linux\" '\n",
+      "            'compatibility, emphasizing the solution offered.\\n'\n",
+      "            '- Retain the closing conversation as it indicates the end of the '\n",
+      "            'session. \\n'\n",
+      "            '\\n'\n",
+      "            '<ct:trimmed-history>\\n'\n",
+      "            '[\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"role\": \"system\",\\n'\n",
+      "            '    \"name\": \"summary\",\\n'\n",
+      "            '    \"content\": \"Event: The user discussed finishing \\'Red Dead '\n",
+      "            \"Redemption 2' and particularly enjoyed the last mission. They \"\n",
+      "            'also talked about favorite missions and characters, specifically '\n",
+      "            'Arthur over Dutch.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 12,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"im playing helldivers 2. do u know of it\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 13,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Yes, Helldivers 2! It\\'s a great choice if '\n",
+      "            \"you're into intense co-op action. The game really ramps up with \"\n",
+      "            'its twin-stick shooter gameplay, and it\\\\u2019s super fun but '\n",
+      "            'also quite challenging. Working together with teammates to '\n",
+      "            'complete those missions can get pretty hectic. How are you '\n",
+      "            'finding the chaos and team dynamics so far?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 14,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"amazing. my fav build is laser cannon + guard '\n",
+      "            'dog. what about u\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 15,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Nice combo! The laser cannon is awesome for '\n",
+      "            'tearing through those enemy swarms, and the guard dog drone '\n",
+      "            'really helps in keeping the smaller threats at bay. I personally '\n",
+      "            'love using the railgun with a shield backpack. It gives a good '\n",
+      "            'balance of offense and defense, especially when things get really '\n",
+      "            'chaotic. What\\'s been the toughest mission for you so far?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 16,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"dude the charger is so difficult to kill.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 17,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Yeah, the Charger can be brutal! They really test '\n",
+      "            \"your reflexes and teamwork. It's all about timing and \"\n",
+      "            'coordination with your team to bring those heavy hitters down. '\n",
+      "            'Have you tried using more strategic traps or mines? Sometimes '\n",
+      "            'setting up the right position with static fields can really make '\n",
+      "            'a difference against them. How do your squadmates handle it?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 18,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"we just run and gun\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 19,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Classic strategy! That can definitely be fun, '\n",
+      "            'especially in a game as intense as Helldivers 2 where things get '\n",
+      "            'chaotic fast. Keeping on the move helps a ton. Just make sure '\n",
+      "            \"someone's always got your back! Have you guys tried any of the \"\n",
+      "            'higher difficulty levels yet, or are you sticking to a certain '\n",
+      "            'comfort zone for now?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"role\": \"system\",\\n'\n",
+      "            '    \"name\": \"summary\",\\n'\n",
+      "            '    \"content\": \"Event: The user discussed an Nvidia driver issue, '\n",
+      "            'particularly its incompatibility with Linux, and the assistant '\n",
+      "            'provided potential solutions and resources for assistance.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 24,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"yep. see u bye!\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 25,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"See you, take care! Bye!\"\\n'\n",
+      "            '  }\\n'\n",
+      "            ']\\n',\n",
+      " 'role': 'assistant'}\n",
+      "End CoT generation\n",
+      "Starting chatml generation\n",
+      "End chatml generation\n"
+     ]
+    }
+   ],
+   "source": [
+    "trimmed_session = json.loads(trimmed_messages[\"content\"])\n",
+    "entity_info = entities(entity_info[\"content\"])\n",
+    "\n",
+    "\n",
+    "summarized_messages = summarize_messages(\n",
+    "    entity_info=entity_info, chat_session=trimmed_session\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 56,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'content': '[\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"role\": \"system\",\\n'\n",
+      "            '    \"name\": \"summary\",\\n'\n",
+      "            '    \"content\": \"Event: The user discussed finishing \\'Red Dead '\n",
+      "            \"Redemption 2' and particularly enjoyed the last mission. They \"\n",
+      "            'also talked about favorite missions and characters, specifically '\n",
+      "            'Arthur over Dutch.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 12,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"im playing helldivers 2. do u know of it\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 13,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Yes, Helldivers 2! It\\'s a great choice if '\n",
+      "            \"you're into intense co-op action. The game really ramps up with \"\n",
+      "            'its twin-stick shooter gameplay, and it\\\\u2019s super fun but '\n",
+      "            'also quite challenging. Working together with teammates to '\n",
+      "            'complete those missions can get pretty hectic. How are you '\n",
+      "            'finding the chaos and team dynamics so far?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 14,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"amazing. my fav build is laser cannon + guard '\n",
+      "            'dog. what about u\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 15,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Nice combo! The laser cannon is awesome for '\n",
+      "            'tearing through those enemy swarms, and the guard dog drone '\n",
+      "            'really helps in keeping the smaller threats at bay. I personally '\n",
+      "            'love using the railgun with a shield backpack. It gives a good '\n",
+      "            'balance of offense and defense, especially when things get really '\n",
+      "            'chaotic. What\\'s been the toughest mission for you so far?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 16,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"dude the charger is so difficult to kill.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 17,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Yeah, the Charger can be brutal! They really test '\n",
+      "            \"your reflexes and teamwork. It's all about timing and \"\n",
+      "            'coordination with your team to bring those heavy hitters down. '\n",
+      "            'Have you tried using more strategic traps or mines? Sometimes '\n",
+      "            'setting up the right position with static fields can really make '\n",
+      "            'a difference against them. How do your squadmates handle it?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 18,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"we just run and gun\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 19,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"Classic strategy! That can definitely be fun, '\n",
+      "            'especially in a game as intense as Helldivers 2 where things get '\n",
+      "            'chaotic fast. Keeping on the move helps a ton. Just make sure '\n",
+      "            \"someone's always got your back! Have you guys tried any of the \"\n",
+      "            'higher difficulty levels yet, or are you sticking to a certain '\n",
+      "            'comfort zone for now?\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"role\": \"system\",\\n'\n",
+      "            '    \"name\": \"summary\",\\n'\n",
+      "            '    \"content\": \"Event: The user discussed an Nvidia driver issue, '\n",
+      "            'particularly its incompatibility with Linux, and the assistant '\n",
+      "            'provided potential solutions and resources for assistance.\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 24,\\n'\n",
+      "            '    \"role\": \"user\",\\n'\n",
+      "            '    \"content\": \"yep. see u bye!\"\\n'\n",
+      "            '  },\\n'\n",
+      "            '  {\\n'\n",
+      "            '    \"index\": 25,\\n'\n",
+      "            '    \"role\": \"assistant\",\\n'\n",
+      "            '    \"content\": \"See you, take care! Bye!\"\\n'\n",
+      "            '  }\\n'\n",
+      "            ']',\n",
+      " 'role': 'assistant'}\n"
+     ]
+    }
+   ],
+   "source": [
+    "pprint(summarized_messages)"
    ]
   }
  ],