Conversation History Trimming with max_tokens and max_history_length

agentai/conversation.py

@@ -1,5 +1,7 @@
+import threading
 from typing import List, Optional
 
+import tiktoken
 from pydantic import BaseModel
 from termcolor import colored
 
@@ -11,7 +13,7 @@ class Message(BaseModel):
 
 
 class Conversation:
-    def __init__(self, history: List[Message] = [], id: Optional[str] = None):
+    def __init__(self, history: List[Message] = [], id: Optional[str] = None, max_history_tokens: int = 200):
         self.history: List[Message] = history
         self.role_to_color = {
             "system": "red",
@@ -20,19 +22,42 @@ def __init__(self, history: List[Message] = [], id: Optional[str] = None):
             "function": "magenta",
         }
         self.id = id
+        self.max_tokens = max_tokens
 
-    def add_message(self, role: str, content: str, name: Optional[str] = None):
+    def add_message(self, role: str, content: str, name: Optional[str] = None) -> None:
         message_dict = {"role": role, "content": content}
-        if name is not None:
+        if name:
             message_dict["name"] = name
         message = Message(**message_dict)
         self.history.append(message)
 
-    def display_conversation(self):
+    def display_conversation(self) -> None:
         for message in self.history:
             print(
                 colored(
                     f"{message.role}: {message.content}\n\n",
-                    self.role_to_color[message.role],
+                    self.role_to_color[message.role.lower()],
                 )
             )
+
+    def get_history(self) -> List[Message]:
+        """Function to get the conversation history based on the number of tokens"""
+        local = threading.local()
+        try:
+            enc = local.gpt2enc
+        except AttributeError:
+            enc = tiktoken.get_encoding("gpt2")
+            local.gpt2enc = enc
+
+        total_tokens = 0
+        for i in range(len(self.history) - 1, -1, -1):
+            # Iterate over the messages in reverse order - from the latest to the oldest messages
+            message = self.history[i]  # Message(role='User', content='I appreciate that. Take care too!', name=None)
+            content = message.content
+            tokens = len(enc.encode(content))
+            total_tokens += tokens
+            if total_tokens > self.max_tokens:
+                # Trim the history inplace to keep the total tokens under max_tokens
+                self.history = self.history[i + 1 :]
+                break
+        return self.history

docs/05_Conversation_History.ipynb

@@ -0,0 +1,103 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from agentai.conversation import Conversation\n",
+    "\n",
+    "# Create a conversation object with a maximum history length of 20 messages\n",
+    "conversation = Conversation(max_tokens=100)\n",
+    "\n",
+    "# Chatgpt Simulated conversations of different length\n",
+    "conversation.add_message(\"User\", \"Hello!\")\n",
+    "conversation.add_message(\"Assistant\", \"Hi there!\")\n",
+    "conversation.add_message(\"User\", \"How are you?\")\n",
+    "conversation.add_message(\"Assistant\", \"I'm doing well, thank you!\")\n",
+    "conversation.add_message(\"User\", \"That's great to hear!\")\n",
+    "conversation.add_message(\"Assistant\", \"Thank you! How can I assist you today?\")\n",
+    "conversation.add_message(\"User\", \"I have a recipe conundrum. I want to make something special for dinner.\")\n",
+    "conversation.add_message(\"Assistant\", \"That sounds fun! What type of cuisine are you interested in?\")\n",
+    "conversation.add_message(\"User\", \"I'd like to try making a Thai dish.\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"Excellent choice! How about trying to make Pad Thai? It's a popular Thai noodle dish.\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"Pad Thai sounds delicious! Do you have a recipe?\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"Certainly! Here's a simple Pad Thai recipe for you:\\n\\nIngredients:\\n- Rice noodles\\n- Tofu or your choice of protein\\n- Eggs\\n- Bean sprouts\\n- Chopped green onions\\n- Chopped peanuts\\n- Lime wedges\\n- Pad Thai sauce (made from tamarind paste, fish sauce, sugar, and chili flakes)\\n\\nInstructions:\\n1. Soak the rice noodles in warm water until they soften. Drain and set aside.\\n2. In a wok or large pan, stir-fry the tofu or protein until cooked. Push it to one side of the pan.\\n3. Crack the eggs into the pan and scramble them until they're nearly cooked.\\n4. Add the softened rice noodles to the pan and pour the Pad Thai sauce over them.\\n5. Toss everything together until the noodles are coated in the sauce.\\n6. Add bean sprouts and chopped green onions, continuing to stir-fry until they're tender.\\n7. Serve hot, garnished with chopped peanuts and lime wedges.\\n\\nEnjoy your Pad Thai!\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"Wow, thank you! I can't wait to try making it.\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"You're welcome! Cooking can be a delightful experience. If you have any more questions or need assistance with anything else, feel free to ask.\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"Sure thing! By the way, what's the weather like today?\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"I apologize, but as an AI language model, I don't have access to real-time data, including weather updates. You can check the weather forecast on a reliable website or a weather app for accurate information.\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"No problem. I have another question then. Can you tell me a joke?\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"Of course! Here's one for you: Why don't scientists trust atoms? Because they make up everything!\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"Haha, that's a good one! I needed that laugh. Thanks!\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"You're welcome! Laughter is the best medicine. If you ever need more jokes or anything else, I'm here to brighten your day!\",\n",
+    ")\n",
+    "conversation.add_message(\n",
+    "    \"User\",\n",
+    "    \"That's very kind of you. Actually, I've been feeling a bit stressed lately.\",\n",
+    ")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"I'm sorry to hear that. Stress can be challenging to deal with. Is there something specific bothering you, or do you need some general tips to manage stress?\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"I'd appreciate some tips on how to manage stress.\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"Of course! Here are some helpful tips:\\n1. Practice deep breathing or meditation to relax your mind and body.\\n2. Engage in physical activities like walking, jogging, or yoga to release endorphins.\\n3. Take breaks and do something you enjoy, like reading, listening to music, or spending time in nature.\\n4. Prioritize your tasks and break them into smaller, manageable steps.\\n5. Reach out to friends or family for support and talk about what's on your mind.\\n6. Get enough sleep to allow your body to recover and rejuvenate.\\n7. Consider trying relaxation techniques like progressive muscle relaxation or visualization.\\n8. Limit caffeine and alcohol consumption, as they can increase stress levels.\\n9. Learn to say no to tasks or commitments when you feel overwhelmed.\\n10. Seek professional help if you find that stress is significantly affecting your daily life.\\n\\nRemember, it's essential to take care of yourself, and don't hesitate to seek support when needed.\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"Thank you for the wonderful advice. I'll definitely give these tips a try.\")\n",
+    "conversation.add_message(\n",
+    "    \"Assistant\",\n",
+    "    \"You're welcome! I hope these tips help you find some relief from stress. If you have any more questions or need further assistance, don't hesitate to ask. Take care!\",\n",
+    ")\n",
+    "conversation.add_message(\"User\", \"I appreciate that. Take care too!\")\n",
+    "conversation.add_message(\"Assistant\", \"Thank you! Have a fantastic day!\")\n",
+    "\n",
+    "# Get the current conversation history with a maximum token limit of 500\n",
+    "history = conversation.get_history()\n",
+    "conversation.display_conversation()\n",
+    "print(history)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "aditya",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}