From 746981ed9e6b7f31e4a0641e40f79deb2e59b281 Mon Sep 17 00:00:00 2001
From: Cogent Apps <cogentapps@fastmail.com>
Date: Mon, 20 Mar 2023 14:13:43 +0000
Subject: [PATCH] update trimmer logic

---
 app/src/tokenizer/chat-history-trimmer.ts | 27 +++++++++++++++++++----
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/app/src/tokenizer/chat-history-trimmer.ts b/app/src/tokenizer/chat-history-trimmer.ts
index 62c6897..a019feb 100644
--- a/app/src/tokenizer/chat-history-trimmer.ts
+++ b/app/src/tokenizer/chat-history-trimmer.ts
@@ -46,10 +46,29 @@ export class ChatHistoryTrimmer {
             return this.output;
         }
 
+        const systemPrompt = this.messages.find(m => m.role === 'system');
+        const firstUserMessage = this.messages.find(m => m.role === 'user');
         const last = this.messages[this.messages.length - 1];
+
+        this.output = [
+            tokenizer.truncateMessage(systemPrompt!, 100),
+        ];
+
+        if (firstUserMessage === last) {
+            this.output.push(tokenizer.truncateMessage(firstUserMessage, this.options.maxTokens - 100));
+        } else {
+            this.output.push(tokenizer.truncateMessage(firstUserMessage!, 100));
+            this.output.push(tokenizer.truncateMessage(last, this.options.maxTokens - 200));
+        }
+
+        excessTokens = this.countExcessTokens();
+        if (excessTokens === 0) {
+            return this.output;
+        }
+
         this.output = [
             tokenizer.truncateMessage(last, this.options.maxTokens),
-        ]
+        ];
 
         return this.output;
     }
@@ -82,13 +101,13 @@ export class ChatHistoryTrimmer {
         const output: OpenAIMessage[] = [...this.output];
 
         for (let i = 0; i < this.output.length && tokenizer.countTokensForMessages(output) > this.options.maxTokens; i++) {
-            if (i == lastMessageIndex) {
+            if (i === lastMessageIndex) {
                 continue;
             }
-            if (i !== systemPromptIndex && !this.options.preserveSystemPrompt) {
+            if (i === systemPromptIndex || this.options.preserveSystemPrompt) {
                 continue;
             }
-            if (i !== firstUserMessageIndex && this.options.preserveFirstUserMessage) {
+            if (i === firstUserMessageIndex || this.options.preserveFirstUserMessage) {
                 continue;
             }
             output[i].content = '';