From 4db0222cd9be94d100ecd47e032caef3687e0740 Mon Sep 17 00:00:00 2001
From: smathot <s.mathot@cogsci.nl>
Date: Mon, 22 Jul 2024 16:40:59 +0200
Subject: [PATCH] Fix the message history from becoming too large due to images
 - Bump to 0.22.1

---
 sigmund/__init__.py              | 2 +-
 sigmund/config.py                | 3 +++
 sigmund/messages.py              | 6 ++++--
 sigmund/sigmund.py               | 6 ++++--
 sigmund/tools/_generate_image.py | 2 +-
 5 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/sigmund/__init__.py b/sigmund/__init__.py
index c7d7b5a..49e602c 100644
--- a/sigmund/__init__.py
+++ b/sigmund/__init__.py
@@ -1,3 +1,3 @@
 """AI-based chatbot that provides sensible answers based on documentation"""
 
-__version__ = '0.22.0'
+__version__ = '0.22.1'
diff --git a/sigmund/config.py b/sigmund/config.py
index 59c26bb..6c856e4 100644
--- a/sigmund/config.py
+++ b/sigmund/config.py
@@ -66,6 +66,9 @@
 # The number of previous messages for which tool results should be 
 # retained.
 keep_tool_results = 4
+# Tool results larger than this are not included in the prompt used for
+# generating
+large_tool_result_length = 1024
 
 # RATE LIMITS
 #
diff --git a/sigmund/messages.py b/sigmund/messages.py
index 133abe7..7a138a3 100644
--- a/sigmund/messages.py
+++ b/sigmund/messages.py
@@ -83,7 +83,7 @@ def delete(self, message_id):
         if self._persistent:
             self.save()
             
-    def prompt(self, system_prompt=None):
+    def prompt(self, system_prompt=None, skip_large_tool_results=False):
         """The prompt consists of the system prompt followed by a sequence of
         AI, user, and tool/ function messages.
         
@@ -104,6 +104,8 @@ def prompt(self, system_prompt=None):
             elif role == 'tool':
                 if msg_nr + config.keep_tool_results < msg_len:
                     continue
+                if len(content) > config.large_tool_result_length:
+                    continue
                 model_prompt.append(FunctionMessage(content=content,
                                                     name='tool_function'))
             else:
@@ -232,4 +234,4 @@ def save(self):
             'condensed_message_history': self._condensed_message_history,
             'title': self._conversation_title,
         }
-        self._sigmund.database.update_active_conversation(conversation)
\ No newline at end of file
+        self._sigmund.database.update_active_conversation(conversation)
diff --git a/sigmund/sigmund.py b/sigmund/sigmund.py
index e9d0263..044b423 100644
--- a/sigmund/sigmund.py
+++ b/sigmund/sigmund.py
@@ -121,7 +121,8 @@ def _search(self, message: str) -> GeneratorType:
         # Then search based on the search-model queries derived from the user
         # question
         reply = self.search_model.predict(self.messages.prompt(
-            system_prompt=prompt.SYSTEM_PROMPT_SEARCH))
+            system_prompt=prompt.SYSTEM_PROMPT_SEARCH,
+            skip_large_tool_results=True))
         if config.log_replies:
             logger.info(f'[search state] reply: {reply}')
         if callable(reply):
@@ -140,7 +141,8 @@ def _answer(self, state: str = 'answer') -> GeneratorType:
         # We first collect a regular reply to the user message. While doing so
         # we also keep track of the number of tokens consumed.
         tokens_consumed_before = self.answer_model.total_tokens_consumed
-        reply = self.answer_model.predict(self.messages.prompt())
+        reply = self.answer_model.predict(self.messages.prompt(
+            skip_large_tool_results=True))
         tokens_consumed = self.answer_model.total_tokens_consumed \
             - tokens_consumed_before
         logger.info(f'tokens consumed: {tokens_consumed}')
diff --git a/sigmund/tools/_generate_image.py b/sigmund/tools/_generate_image.py
index 2389b73..d96ed5d 100644
--- a/sigmund/tools/_generate_image.py
+++ b/sigmund/tools/_generate_image.py
@@ -15,7 +15,7 @@ class generate_image(BaseTool):
         },
         'quality': {
             'type': 'string',
-            'description': 'The quality of the image',
+            'description': 'The quality of the image. Use standard unless there is a good reason to use HD.',
             'enum': ['standard', 'hd']
         },
         'size': {