From e76fc771896bd636f733b69e7f398dcb28943d1c Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Sat, 6 May 2023 20:31:53 +1000
Subject: [PATCH] fixes (#53)

* Minor... use username suggester in case username already exists

* FIX: ensure we truncate long prompts

Previously we

1. Used raw length instead of token counts for counting length
2. We totally dropped a prompt if it was too long

New implementation will truncate "raw" if it gets too long maintaining
meaning.
---
 app/models/completion_prompt.rb       | 17 +++++++++++++----
 db/fixtures/ai_bot/602_bot_users.rb   |  3 +--
 lib/shared/tokenizer/tokenizer.rb     |  5 ++++-
 spec/models/completion_prompt_spec.rb | 13 +++++++++++++
 4 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/app/models/completion_prompt.rb b/app/models/completion_prompt.rb
index c8227359..135408d7 100644
--- a/app/models/completion_prompt.rb
+++ b/app/models/completion_prompt.rb
@@ -4,7 +4,8 @@ class CompletionPrompt < ActiveRecord::Base
   # TODO(roman): Remove sept 2023.
   self.ignored_columns = ["value"]
 
-  MAX_PROMPT_LENGTH = 3000
+  # GPT 3.5 allows 4000 tokens
+  MAX_PROMPT_TOKENS = 3500
 
   enum :prompt_type, { text: 0, list: 1, diff: 2 }
 
@@ -22,11 +23,19 @@ class CompletionPrompt < ActiveRecord::Base
         .order("post_number desc")
         .pluck(:raw, :username)
 
-    total_prompt_length = 0
+    total_prompt_tokens = 0
     messages =
       conversation.reduce([]) do |memo, (raw, username)|
-        total_prompt_length += raw.length
-        break(memo) if total_prompt_length > MAX_PROMPT_LENGTH
+        break(memo) if total_prompt_tokens >= MAX_PROMPT_TOKENS
+
+        tokens = DiscourseAi::Tokenizer.tokenize(raw)
+
+        if tokens.length + total_prompt_tokens > MAX_PROMPT_TOKENS
+          tokens = tokens[0...(MAX_PROMPT_TOKENS - total_prompt_tokens)]
+          raw = tokens.join(" ")
+        end
+
+        total_prompt_tokens += tokens.length
         role = username == Discourse.gpt_bot.username ? "system" : "user"
 
         memo.unshift({ role: role, content: raw })
diff --git a/db/fixtures/ai_bot/602_bot_users.rb b/db/fixtures/ai_bot/602_bot_users.rb
index 8c0ff990..7beb41b2 100644
--- a/db/fixtures/ai_bot/602_bot_users.rb
+++ b/db/fixtures/ai_bot/602_bot_users.rb
@@ -10,8 +10,7 @@ end
 User.seed do |u|
   u.id = -110
   u.name = "GPT Bot"
-  u.username = "gpt_bot"
-  u.username_lower = "gpt_bot"
+  u.username = UserNameSuggester.suggest("gpt_bot")
   u.password = SecureRandom.hex
   u.active = true
   u.admin = true
diff --git a/lib/shared/tokenizer/tokenizer.rb b/lib/shared/tokenizer/tokenizer.rb
index 17f55fdc..ea80a18e 100644
--- a/lib/shared/tokenizer/tokenizer.rb
+++ b/lib/shared/tokenizer/tokenizer.rb
@@ -7,8 +7,11 @@ module DiscourseAi
         Tokenizers.from_file("./plugins/discourse-ai/tokenizers/bert-base-uncased.json")
     end
 
+    def self.tokenize(text)
+      tokenizer.encode(text).tokens
+    end
     def self.size(text)
-      tokenizer.encode(text).tokens.size
+      tokenize(text).size
     end
   end
 end
diff --git a/spec/models/completion_prompt_spec.rb b/spec/models/completion_prompt_spec.rb
index c4f6cb57..9fe29457 100644
--- a/spec/models/completion_prompt_spec.rb
+++ b/spec/models/completion_prompt_spec.rb
@@ -39,6 +39,19 @@ RSpec.describe CompletionPrompt do
       end
     end
 
+    context "when prompt gets very long" do
+      fab!(:post_1) { Fabricate(:post, topic: topic, raw: "test " * 6000, post_number: 1) }
+
+      it "trims the prompt" do
+        prompt_messages = described_class.bot_prompt_with_topic_context(post_1)
+
+        expect(prompt_messages[0][:role]).to eq("system")
+        expect(prompt_messages[1][:role]).to eq("user")
+        expected_length = ("test " * (CompletionPrompt::MAX_PROMPT_TOKENS)).length
+        expect(prompt_messages[1][:content].length).to eq(expected_length)
+      end
+    end
+
     context "when the topic has multiple posts" do
       fab!(:post_1) { Fabricate(:post, topic: topic, raw: post_body(1), post_number: 1) }
       fab!(:post_2) do