From 4cfbdac7d694f5e2b28a9e873f18d784d8008b65 Mon Sep 17 00:00:00 2001
From: Limour <93720049+Limour-dev@users.noreply.github.com>
Date: Sat, 6 May 2023 11:15:12 +0800
Subject: [PATCH] Set a system role as a persistent role. (#134)

* set a system role as a persistent role

Fixes https://github.com/ztjhz/BetterChatGPT/issues/130

* handle system message exceed token limit + improve logic

---------

Co-authored-by: Jing Hua <tohjinghua123@gmail.com>
---
 src/utils/messageUtils.ts | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/utils/messageUtils.ts b/src/utils/messageUtils.ts
index 2d7fc62..e4aacde 100644
--- a/src/utils/messageUtils.ts
+++ b/src/utils/messageUtils.ts
@@ -51,13 +51,39 @@ export const limitMessageTokens = (
   const limitedMessages: MessageInterface[] = [];
   let tokenCount = 0;
 
-  for (let i = messages.length - 1; i >= 0; i--) {
+  const isSystemFirstMessage = messages[0]?.role === 'system';
+  let retainSystemMessage = false;
+
+  // Check if the first message is a system message and if it fits within the token limit
+  if (isSystemFirstMessage) {
+    const systemTokenCount = countTokens([messages[0]], model);
+    if (systemTokenCount < limit) {
+      tokenCount += systemTokenCount;
+      retainSystemMessage = true;
+    }
+  }
+
+  // Iterate through messages in reverse order, adding them to the limitedMessages array
+  // until the token limit is reached (excludes first message)
+  for (let i = messages.length - 1; i >= 1; i--) {
     const count = countTokens([messages[i]], model);
     if (count + tokenCount > limit) break;
     tokenCount += count;
     limitedMessages.unshift({ ...messages[i] });
   }
 
+  // Process first message
+  if (retainSystemMessage) {
+    // Insert the system message in the third position from the end
+    limitedMessages.splice(-3, 0, { ...messages[0] });
+  } else if (!isSystemFirstMessage) {
+    // Check if the first message (non-system) can fit within the limit
+    const firstMessageTokenCount = countTokens([messages[0]], model);
+    if (firstMessageTokenCount + tokenCount < limit) {
+      limitedMessages.unshift({ ...messages[0] });
+    }
+  }
+
   return limitedMessages;
 };