From a3c827efcc097ec183a666f0cfeaeb0da970ec8e Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Sat, 3 Feb 2024 07:09:34 +1100
Subject: [PATCH] FEATURE: allow personas to supply top_p and temperature
 params (#459)

* FEATURE: allow personas to supply top_p and temperature params

Code assistance generally are more focused at a lower temperature
This amends it so SQL Helper runs at 0.2 temperature vs the more
common default across LLMs of 1.0.

Reduced temperature leads to more focused, concise and predictable
answers for the SQL Helper

* fix tests

* This is not perfect, but far better than what we do today

Instead of fishing for

1. Draft sequence
2. Draft body

We skip (2), this means the composer "only" needs 1 http request to
open, we also want to eliminate (1) but it is a bit of a trickier
core change, may figure out how to pull it off (defer it to first draft save)

Value of bot drafts < value of opening bot conversations really fast
---
 .../admin/ai_personas_controller.rb           |  2 +
 app/models/ai_persona.rb                      | 17 ++++++--
 .../localized_ai_persona_serializer.rb        |  4 +-
 .../discourse/admin/models/ai-persona.js      |  2 +
 .../components/ai-persona-editor.gjs          | 38 ++++++++++++++++
 .../discourse/lib/ai-bot-helper.js            |  3 +-
 config/locales/client.en.yml                  |  4 ++
 db/fixtures/ai_bot/603_bot_ai_personas.rb     |  2 +
 ...52_add_temperature_top_p_to_ai_personas.rb |  8 ++++
 lib/ai_bot/bot.rb                             |  8 +++-
 lib/ai_bot/personas/persona.rb                |  9 +++-
 lib/ai_bot/personas/sql_helper.rb             |  4 ++
 lib/completions/endpoints/fake.rb             | 10 +++++
 spec/lib/modules/ai_bot/bot_spec.rb           | 36 ++++++++++++++++
 .../admin/ai_personas_controller_spec.rb      | 43 ++++++++++++++++++-
 .../unit/models/ai-persona-test.js            |  4 ++
 16 files changed, 186 insertions(+), 8 deletions(-)
 create mode 100644 db/migrate/20240202010752_add_temperature_top_p_to_ai_personas.rb

diff --git a/app/controllers/discourse_ai/admin/ai_personas_controller.rb b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
index 120b7c32..1a994aa2 100644
--- a/app/controllers/discourse_ai/admin/ai_personas_controller.rb
+++ b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
@@ -62,6 +62,8 @@ module DiscourseAi
             :enabled,
             :system_prompt,
             :priority,
+            :top_p,
+            :temperature,
             allowed_group_ids: [],
           )
 
diff --git a/app/models/ai_persona.rb b/app/models/ai_persona.rb
index e49aff62..54e54020 100644
--- a/app/models/ai_persona.rb
+++ b/app/models/ai_persona.rb
@@ -157,6 +157,14 @@ class AiPersona < ActiveRecord::Base
         options
       end
 
+      define_method :temperature do
+        @ai_persona&.temperature
+      end
+
+      define_method :top_p do
+        @ai_persona&.top_p
+      end
+
       define_method :system_prompt do
         @ai_persona&.system_prompt || "You are a helpful bot."
       end
@@ -166,7 +174,8 @@ class AiPersona < ActiveRecord::Base
   private
 
   def system_persona_unchangeable
-    if system_prompt_changed? || commands_changed? || name_changed? || description_changed?
+    if top_p_changed? || temperature_changed? || system_prompt_changed? || commands_changed? ||
+         name_changed? || description_changed?
       errors.add(:base, I18n.t("discourse_ai.ai_bot.personas.cannot_edit_system_persona"))
     end
   end
@@ -186,7 +195,7 @@ end
 #  id                :bigint           not null, primary key
 #  name              :string(100)      not null
 #  description       :string(2000)     not null
-#  commands          :string           default([]), not null, is an Array
+#  commands          :json             not null
 #  system_prompt     :string(10000000) not null
 #  allowed_group_ids :integer          default([]), not null, is an Array
 #  created_by_id     :integer
@@ -194,7 +203,9 @@ end
 #  created_at        :datetime         not null
 #  updated_at        :datetime         not null
 #  system            :boolean          default(FALSE), not null
-#  priority          :integer          default(0), not null
+#  priority          :boolean          default(FALSE), not null
+#  temperature       :float
+#  top_p             :float
 #
 # Indexes
 #
diff --git a/app/serializers/localized_ai_persona_serializer.rb b/app/serializers/localized_ai_persona_serializer.rb
index 1ecd3239..fca5f6d2 100644
--- a/app/serializers/localized_ai_persona_serializer.rb
+++ b/app/serializers/localized_ai_persona_serializer.rb
@@ -11,7 +11,9 @@ class LocalizedAiPersonaSerializer < ApplicationSerializer
              :priority,
              :commands,
              :system_prompt,
-             :allowed_group_ids
+             :allowed_group_ids,
+             :temperature,
+             :top_p
 
   def name
     object.class_instance.name
diff --git a/assets/javascripts/discourse/admin/models/ai-persona.js b/assets/javascripts/discourse/admin/models/ai-persona.js
index d9152cd9..8d7c4171 100644
--- a/assets/javascripts/discourse/admin/models/ai-persona.js
+++ b/assets/javascripts/discourse/admin/models/ai-persona.js
@@ -10,6 +10,8 @@ const ATTRIBUTES = [
   "enabled",
   "system",
   "priority",
+  "top_p",
+  "temperature",
 ];
 
 class CommandOption {
diff --git a/assets/javascripts/discourse/components/ai-persona-editor.gjs b/assets/javascripts/discourse/components/ai-persona-editor.gjs
index be922cdb..b172d0de 100644
--- a/assets/javascripts/discourse/components/ai-persona-editor.gjs
+++ b/assets/javascripts/discourse/components/ai-persona-editor.gjs
@@ -73,6 +73,14 @@ export default class PersonaEditor extends Component {
     }
   }
 
+  get showTemperature() {
+    return this.editingModel?.temperature || !this.editingModel?.system;
+  }
+
+  get showTopP() {
+    return this.editingModel?.top_p || !this.editingModel?.system;
+  }
+
   @action
   delete() {
     return this.dialog.confirm({
@@ -213,6 +221,36 @@ export default class PersonaEditor extends Component {
           disabled={{this.editingModel.system}}
         />
       </div>
+      {{#if this.showTemperature}}
+        <div class="control-group">
+          <label>{{I18n.t "discourse_ai.ai_persona.temperature"}}</label>
+          <Input
+            @type="number"
+            class="ai-persona-editor__temperature"
+            @value={{this.editingModel.temperature}}
+            disabled={{this.editingModel.system}}
+          />
+          <DTooltip
+            @icon="question-circle"
+            @content={{I18n.t "discourse_ai.ai_persona.temperature_help"}}
+          />
+        </div>
+      {{/if}}
+      {{#if this.showTopP}}
+        <div class="control-group">
+          <label>{{I18n.t "discourse_ai.ai_persona.top_p"}}</label>
+          <Input
+            @type="number"
+            class="ai-persona-editor__top_p"
+            @value={{this.editingModel.top_p}}
+            disabled={{this.editingModel.system}}
+          />
+          <DTooltip
+            @icon="question-circle"
+            @content={{I18n.t "discourse_ai.ai_persona.top_p_help"}}
+          />
+        </div>
+      {{/if}}
       <div class="control-group ai-persona-editor__action_panel">
         <DButton
           class="btn-primary ai-persona-editor__save"
diff --git a/assets/javascripts/discourse/lib/ai-bot-helper.js b/assets/javascripts/discourse/lib/ai-bot-helper.js
index 5ed0e19d..cc292c9c 100644
--- a/assets/javascripts/discourse/lib/ai-bot-helper.js
+++ b/assets/javascripts/discourse/lib/ai-bot-helper.js
@@ -15,9 +15,10 @@ export function composeAiBotMessage(targetBot, composer) {
       recipients: botUsername,
       topicTitle: I18n.t("discourse_ai.ai_bot.default_pm_prefix"),
       archetypeId: "private_message",
-      draftKey: Composer.NEW_PRIVATE_MESSAGE_KEY,
+      draftKey: "private_message_ai",
       hasGroups: false,
       warningsDisabled: true,
+      skipDraftCheck: true,
     },
   });
 }
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index 2450795c..4d3e9a61 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -118,6 +118,10 @@ en:
         new: New
         title: "AI Personas"
         delete: Delete
+        temperature: Temperature
+        temperature_help: Temperature to use for the LLM, increase to increase creativity (leave empty to use model default, generally a value from 0.0 to 2.0)
+        top_p: Top P
+        top_p_help: Top P to use for the LLM, increase to increase randomness (leave empty to use model default, generally a value from 0.0 to 1.0)
         priority: Priority
         priority_help: Priority personas are displayed to users at the top of the persona list. If multiple personas have priority, they will be sorted alphabetically.
         command_options: "Command Options"
diff --git a/db/fixtures/ai_bot/603_bot_ai_personas.rb b/db/fixtures/ai_bot/603_bot_ai_personas.rb
index 7c110dba..30d4ef2d 100644
--- a/db/fixtures/ai_bot/603_bot_ai_personas.rb
+++ b/db/fixtures/ai_bot/603_bot_ai_personas.rb
@@ -34,5 +34,7 @@ DiscourseAi::AiBot::Personas::Persona.system_personas.each do |persona_class, id
   instance = persona_class.new
   persona.commands = instance.tools.map { |tool| tool.to_s.split("::").last }
   persona.system_prompt = instance.system_prompt
+  persona.top_p = instance.top_p
+  persona.temperature = instance.temperature
   persona.save!(validate: false)
 end
diff --git a/db/migrate/20240202010752_add_temperature_top_p_to_ai_personas.rb b/db/migrate/20240202010752_add_temperature_top_p_to_ai_personas.rb
new file mode 100644
index 00000000..a101338f
--- /dev/null
+++ b/db/migrate/20240202010752_add_temperature_top_p_to_ai_personas.rb
@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+class AddTemperatureTopPToAiPersonas < ActiveRecord::Migration[7.0]
+  def change
+    add_column :ai_personas, :temperature, :float, null: true
+    add_column :ai_personas, :top_p, :float, null: true
+  end
+end
diff --git a/lib/ai_bot/bot.rb b/lib/ai_bot/bot.rb
index a5a2f329..655f1c8a 100644
--- a/lib/ai_bot/bot.rb
+++ b/lib/ai_bot/bot.rb
@@ -48,13 +48,19 @@ module DiscourseAi
         low_cost = false
         raw_context = []
 
+        user = context[:user]
+
+        llm_kwargs = { user: user }
+        llm_kwargs[:temperature] = persona.temperature if persona.temperature
+        llm_kwargs[:top_p] = persona.top_p if persona.top_p
+
         while total_completions <= MAX_COMPLETIONS && ongoing_chain
           current_model = model(prefer_low_cost: low_cost)
           llm = DiscourseAi::Completions::Llm.proxy(current_model)
           tool_found = false
 
           result =
-            llm.generate(prompt, user: context[:user]) do |partial, cancel|
+            llm.generate(prompt, **llm_kwargs) do |partial, cancel|
               if (tool = persona.find_tool(partial))
                 tool_found = true
                 ongoing_chain = tool.chain_next_response?
diff --git a/lib/ai_bot/personas/persona.rb b/lib/ai_bot/personas/persona.rb
index b5f9d662..e56d3362 100644
--- a/lib/ai_bot/personas/persona.rb
+++ b/lib/ai_bot/personas/persona.rb
@@ -23,7 +23,6 @@ module DiscourseAi
 
           def all(user:)
             # listing tools has to be dynamic cause site settings may change
-
             AiPersona.all_personas.filter do |persona|
               next false if !user.in_any_groups?(persona.allowed_group_ids)
 
@@ -85,6 +84,14 @@ module DiscourseAi
           []
         end
 
+        def temperature
+          nil
+        end
+
+        def top_p
+          nil
+        end
+
         def options
           {}
         end
diff --git a/lib/ai_bot/personas/sql_helper.rb b/lib/ai_bot/personas/sql_helper.rb
index b048d7e3..6afd76a8 100644
--- a/lib/ai_bot/personas/sql_helper.rb
+++ b/lib/ai_bot/personas/sql_helper.rb
@@ -31,6 +31,10 @@ module DiscourseAi
           [Tools::DbSchema]
         end
 
+        def temperature
+          0.2
+        end
+
         def system_prompt
           <<~PROMPT
             You are a PostgreSQL expert.
diff --git a/lib/completions/endpoints/fake.rb b/lib/completions/endpoints/fake.rb
index 45e2cea7..b166cfa2 100644
--- a/lib/completions/endpoints/fake.rb
+++ b/lib/completions/endpoints/fake.rb
@@ -95,7 +95,17 @@ module DiscourseAi
           @chunk_count = chunk_count
         end
 
+        def self.last_call
+          @last_call
+        end
+
+        def self.last_call=(params)
+          @last_call = params
+        end
+
         def perform_completion!(dialect, user, model_params = {})
+          self.class.last_call = { dialect: dialect, user: user, model_params: model_params }
+
           content = self.class.fake_content
 
           if block_given?
diff --git a/spec/lib/modules/ai_bot/bot_spec.rb b/spec/lib/modules/ai_bot/bot_spec.rb
index d5faf055..2f2d7520 100644
--- a/spec/lib/modules/ai_bot/bot_spec.rb
+++ b/spec/lib/modules/ai_bot/bot_spec.rb
@@ -3,6 +3,8 @@
 RSpec.describe DiscourseAi::AiBot::Bot do
   subject(:bot) { described_class.as(bot_user) }
 
+  fab!(:admin) { Fabricate(:admin) }
+
   before do
     SiteSetting.ai_bot_enabled_chat_bots = "gpt-4"
     SiteSetting.ai_bot_enabled = true
@@ -25,6 +27,40 @@ RSpec.describe DiscourseAi::AiBot::Bot do
   let(:llm_responses) { [function_call, response] }
 
   describe "#reply" do
+    it "sets top_p and temperature params" do
+      # full integration test so we have certainty it is passed through
+
+      DiscourseAi::Completions::Endpoints::Fake.delays = []
+      DiscourseAi::Completions::Endpoints::Fake.last_call = nil
+
+      SiteSetting.ai_bot_enabled_chat_bots = "fake"
+      SiteSetting.ai_bot_enabled = true
+      Group.refresh_automatic_groups!
+
+      bot_user = User.find(DiscourseAi::AiBot::EntryPoint::FAKE_ID)
+      AiPersona.create!(
+        name: "TestPersona",
+        top_p: 0.5,
+        temperature: 0.4,
+        system_prompt: "test",
+        description: "test",
+        allowed_group_ids: [Group::AUTO_GROUPS[:trust_level_0]],
+      )
+
+      personaClass = DiscourseAi::AiBot::Personas::Persona.find_by(user: admin, name: "TestPersona")
+
+      bot = DiscourseAi::AiBot::Bot.as(bot_user, persona: personaClass.new)
+      bot.reply(
+        { conversation_context: [{ type: :user, content: "test" }] },
+      ) do |_partial, _cancel, _placeholder|
+        # we just need the block so bot has something to call with results
+      end
+
+      last_call = DiscourseAi::Completions::Endpoints::Fake.last_call
+      expect(last_call[:model_params][:top_p]).to eq(0.5)
+      expect(last_call[:model_params][:temperature]).to eq(0.4)
+    end
+
     context "when using function chaining" do
       it "yields a loading placeholder while proceeds to invoke the command" do
         tool = DiscourseAi::AiBot::Tools::ListCategories.new({})
diff --git a/spec/requests/admin/ai_personas_controller_spec.rb b/spec/requests/admin/ai_personas_controller_spec.rb
index d68bdeb6..3136cc8f 100644
--- a/spec/requests/admin/ai_personas_controller_spec.rb
+++ b/spec/requests/admin/ai_personas_controller_spec.rb
@@ -116,6 +116,8 @@ RSpec.describe DiscourseAi::Admin::AiPersonasController do
           description: "Assists with tasks",
           system_prompt: "you are a helpful bot",
           commands: [["search", { "base_query" => "test" }]],
+          top_p: 0.1,
+          temperature: 0.5,
         }
       end
 
@@ -127,8 +129,17 @@ RSpec.describe DiscourseAi::Admin::AiPersonasController do
                  "CONTENT_TYPE" => "application/json",
                }
           expect(response).to be_successful
-          persona = AiPersona.find(response.parsed_body["ai_persona"]["id"])
+          persona_json = response.parsed_body["ai_persona"]
+
+          expect(persona_json["name"]).to eq("superbot")
+          expect(persona_json["top_p"]).to eq(0.1)
+          expect(persona_json["temperature"]).to eq(0.5)
+
+          persona = AiPersona.find(persona_json["id"])
+
           expect(persona.commands).to eq([["search", { "base_query" => "test" }]])
+          expect(persona.top_p).to eq(0.1)
+          expect(persona.temperature).to eq(0.5)
         }.to change(AiPersona, :count).by(1)
       end
     end
@@ -143,6 +154,36 @@ RSpec.describe DiscourseAi::Admin::AiPersonasController do
   end
 
   describe "PUT #update" do
+    it "allows us to trivially clear top_p and temperature" do
+      persona = Fabricate(:ai_persona, name: "test_bot2", top_p: 0.5, temperature: 0.1)
+
+      put "/admin/plugins/discourse-ai/ai_personas/#{persona.id}.json",
+          params: {
+            ai_persona: {
+              top_p: "",
+              temperature: "",
+            },
+          }
+
+      expect(response).to have_http_status(:ok)
+      persona.reload
+
+      expect(persona.top_p).to eq(nil)
+      expect(persona.temperature).to eq(nil)
+    end
+
+    it "does not allow temperature and top p changes on stock personas" do
+      put "/admin/plugins/discourse-ai/ai_personas/#{DiscourseAi::AiBot::Personas::Persona.system_personas.values.first}.json",
+          params: {
+            ai_persona: {
+              top_p: 0.5,
+              temperature: 0.1,
+            },
+          }
+
+      expect(response).to have_http_status(:unprocessable_entity)
+    end
+
     context "with valid params" do
       it "updates the requested ai_persona" do
         put "/admin/plugins/discourse-ai/ai_personas/#{ai_persona.id}.json",
diff --git a/test/javascripts/unit/models/ai-persona-test.js b/test/javascripts/unit/models/ai-persona-test.js
index 0df91379..84ba05af 100644
--- a/test/javascripts/unit/models/ai-persona-test.js
+++ b/test/javascripts/unit/models/ai-persona-test.js
@@ -39,6 +39,8 @@ module("Discourse AI | Unit | Model | ai-persona", function () {
       system_prompt: "System Prompt",
       priority: false,
       description: "Description",
+      top_p: 0.8,
+      temperature: 0.7,
     };
 
     const aiPersona = AiPersona.create({ ...properties });
@@ -63,6 +65,8 @@ module("Discourse AI | Unit | Model | ai-persona", function () {
       system_prompt: "System Prompt",
       priority: false,
       description: "Description",
+      top_p: 0.8,
+      temperature: 0.7,
     };
 
     const aiPersona = AiPersona.create({ ...properties });