From 0d6d9a6ef59b00520b85b8451bd5bcc15899a1f5 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Wed, 19 Jun 2024 15:49:36 +1000
Subject: [PATCH] FEATURE: allow access to private topics if tool permits
 (#673)

Previously read tool only had access to public topics, this allows
access to all topics user has access to, if admin opts for the option
Also

- Fixes VLLM migration
- Display which llms have bot enabled
---
 .../discourse_ai/admin/ai_llms_controller.rb  |  2 +-
 app/models/llm_model.rb                       |  7 ++-
 .../components/ai-llms-list-editor.gjs        | 28 ++++++++++-
 config/locales/client.en.yml                  |  2 +-
 config/locales/server.en.yml                  |  4 ++
 .../20240603143158_seed_oss_models.rb         |  2 +-
 lib/ai_bot/tools/read.rb                      | 14 +++---
 lib/ai_bot/tools/tool.rb                      |  2 +-
 spec/lib/modules/ai_bot/tools/read_spec.rb    | 46 +++++++++++++++++++
 spec/system/ai_bot/persona_spec.rb            |  2 +-
 10 files changed, 93 insertions(+), 16 deletions(-)

diff --git a/app/controllers/discourse_ai/admin/ai_llms_controller.rb b/app/controllers/discourse_ai/admin/ai_llms_controller.rb
index 250dd873..6e3fc321 100644
--- a/app/controllers/discourse_ai/admin/ai_llms_controller.rb
+++ b/app/controllers/discourse_ai/admin/ai_llms_controller.rb
@@ -6,7 +6,7 @@ module DiscourseAi
       requires_plugin ::DiscourseAi::PLUGIN_NAME
 
       def index
-        llms = LlmModel.all
+        llms = LlmModel.all.order(:display_name)
 
         render json: {
                  ai_llms:
diff --git a/app/models/llm_model.rb b/app/models/llm_model.rb
index 4ac22fc1..07c60789 100644
--- a/app/models/llm_model.rb
+++ b/app/models/llm_model.rb
@@ -18,9 +18,7 @@ class LlmModel < ActiveRecord::Base
           provider: "vllm",
           tokenizer: "DiscourseAi::Tokenizer::MixtralTokenizer",
           url: RESERVED_VLLM_SRV_URL,
-          vllm_key: "",
-          user_id: nil,
-          enabled_chat_bot: false,
+          max_prompt_tokens: 8000,
         )
 
       record.save(validate: false) # Ignore reserved URL validation
@@ -55,7 +53,8 @@ class LlmModel < ActiveRecord::Base
         new_user.save!(validate: false)
         self.update!(user: new_user)
       else
-        user.update!(active: true)
+        user.active = true
+        user.save!(validate: false)
       end
     elsif user
       # will include deleted
diff --git a/assets/javascripts/discourse/components/ai-llms-list-editor.gjs b/assets/javascripts/discourse/components/ai-llms-list-editor.gjs
index 13b00bce..03f74b93 100644
--- a/assets/javascripts/discourse/components/ai-llms-list-editor.gjs
+++ b/assets/javascripts/discourse/components/ai-llms-list-editor.gjs
@@ -1,6 +1,10 @@
 import Component from "@glimmer/component";
-import { concat } from "@ember/helper";
+import { concat, fn } from "@ember/helper";
+import { on } from "@ember/modifier";
+import { action } from "@ember/object";
 import { LinkTo } from "@ember/routing";
+import DToggleSwitch from "discourse/components/d-toggle-switch";
+import { popupAjaxError } from "discourse/lib/ajax-error";
 import icon from "discourse-common/helpers/d-icon";
 import i18n from "discourse-common/helpers/i18n";
 import I18n from "discourse-i18n";
@@ -11,6 +15,21 @@ export default class AiLlmsListEditor extends Component {
     return this.args.llms.length !== 0;
   }
 
+  @action
+  async toggleEnabledChatBot(llm) {
+    const oldValue = llm.enabled_chat_bot;
+    const newValue = !oldValue;
+    try {
+      llm.set("enabled_chat_bot", newValue);
+      await llm.update({
+        enabled_chat_bot: newValue,
+      });
+    } catch (err) {
+      llm.set("enabled_chat_bot", oldValue);
+      popupAjaxError(err);
+    }
+  }
+
   <template>
     <section class="ai-llms-list-editor admin-detail pull-left">
       {{#if @currentLlm}}
@@ -35,6 +54,7 @@ export default class AiLlmsListEditor extends Component {
               <tr>
                 <th>{{i18n "discourse_ai.llms.display_name"}}</th>
                 <th>{{i18n "discourse_ai.llms.provider"}}</th>
+                <th>{{i18n "discourse_ai.llms.enabled_chat_bot"}}</th>
                 <th></th>
               </tr>
             </thead>
@@ -45,6 +65,12 @@ export default class AiLlmsListEditor extends Component {
                   <td>{{i18n
                       (concat "discourse_ai.llms.providers." llm.provider)
                     }}</td>
+                  <td>
+                    <DToggleSwitch
+                      @state={{llm.enabled_chat_bot}}
+                      {{on "click" (fn this.toggleEnabledChatBot llm)}}
+                    />
+                  </td>
                   <td>
                     <LinkTo
                       @route="adminPlugins.show.discourse-ai-llms.show"
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index fe8af127..3dc57650 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -210,7 +210,7 @@ en:
         max_prompt_tokens: "Number of tokens for the prompt"
         url: "URL of the service hosting the model"
         api_key: "API Key of the service hosting the model"
-        enabled_chat_bot: "Allow Companion user to act as an AI Bot"
+        enabled_chat_bot: "Allow AI Bot"
         save: "Save"
         edit: "Edit"
         saved: "LLM Model Saved"
diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
index 50b90e5a..d40f632f 100644
--- a/config/locales/server.en.yml
+++ b/config/locales/server.en.yml
@@ -207,6 +207,10 @@ en:
       summarizing: "Summarizing topic"
       searching: "Searching for: '%{query}'"
       tool_options:
+        read:
+          read_private:
+            name: "Read Private"
+            description: "Allow access to all topics user has access to (by default only public topics are included)"
         search:
           search_private:
             name: "Search Private"
diff --git a/db/post_migrate/20240603143158_seed_oss_models.rb b/db/post_migrate/20240603143158_seed_oss_models.rb
index 9eb86d29..a528fac9 100644
--- a/db/post_migrate/20240603143158_seed_oss_models.rb
+++ b/db/post_migrate/20240603143158_seed_oss_models.rb
@@ -40,7 +40,7 @@ class SeedOssModels < ActiveRecord::Migration[7.0]
         reserved: srv_reserved_url,
       ).first
 
-    if vllm_srv.present? && srv.record.nil?
+    if vllm_srv.present? && srv_record.nil?
       url = "https://vllm.shadowed-by-srv.invalid"
       name = "mistralai/Mixtral"
 
diff --git a/lib/ai_bot/tools/read.rb b/lib/ai_bot/tools/read.rb
index 3cc6e096..8bb3afc5 100644
--- a/lib/ai_bot/tools/read.rb
+++ b/lib/ai_bot/tools/read.rb
@@ -28,6 +28,10 @@ module DiscourseAi
           }
         end
 
+        def self.accepted_options
+          [option(:read_private, type: :boolean)]
+        end
+
         def self.name
           "read"
         end
@@ -44,20 +48,18 @@ module DiscourseAi
 
         def invoke
           not_found = { topic_id: topic_id, description: "Topic not found" }
+          guardian = Guardian.new(context[:user]) if options[:read_private] && context[:user]
+          guardian ||= Guardian.new
 
           @title = ""
 
           topic = Topic.find_by(id: topic_id.to_i)
-          return not_found if !topic || !Guardian.new.can_see?(topic)
+          return not_found if !topic || !guardian.can_see?(topic)
 
           @title = topic.title
 
           posts =
-            Post
-              .secured(Guardian.new)
-              .where(topic_id: topic_id)
-              .order(:post_number)
-              .limit(MAX_POSTS)
+            Post.secured(guardian).where(topic_id: topic_id).order(:post_number).limit(MAX_POSTS)
 
           post_number = 1
           post_number = post_numbers.first if post_numbers.present?
diff --git a/lib/ai_bot/tools/tool.rb b/lib/ai_bot/tools/tool.rb
index 8b8065f9..f5112c3e 100644
--- a/lib/ai_bot/tools/tool.rb
+++ b/lib/ai_bot/tools/tool.rb
@@ -72,7 +72,7 @@ module DiscourseAi
             if val
               case option.type
               when :boolean
-                val = val == "true"
+                val = (val.to_s == "true")
               when :integer
                 val = val.to_i
               end
diff --git a/spec/lib/modules/ai_bot/tools/read_spec.rb b/spec/lib/modules/ai_bot/tools/read_spec.rb
index 8f9a5ff4..10f105f1 100644
--- a/spec/lib/modules/ai_bot/tools/read_spec.rb
+++ b/spec/lib/modules/ai_bot/tools/read_spec.rb
@@ -30,6 +30,52 @@ RSpec.describe DiscourseAi::AiBot::Tools::Read do
   before { SiteSetting.ai_bot_enabled = true }
 
   describe "#process" do
+    it "can read private topics if allowed to" do
+      category = topic_with_tags.category
+      category.set_permissions(Group::AUTO_GROUPS[:staff] => :full)
+      category.save!
+
+      tool =
+        described_class.new(
+          { topic_id: topic_with_tags.id, post_numbers: [post1.post_number] },
+          bot_user: bot_user,
+          llm: llm,
+        )
+      results = tool.invoke
+
+      expect(results[:description]).to eq("Topic not found")
+
+      admin = Fabricate(:admin)
+
+      tool =
+        described_class.new(
+          { topic_id: topic_with_tags.id, post_numbers: [post1.post_number] },
+          bot_user: bot_user,
+          llm: llm,
+          persona_options: {
+            "read_private" => true,
+          },
+          context: {
+            user: admin,
+          },
+        )
+      results = tool.invoke
+      expect(results[:content]).to include("hello there")
+
+      tool =
+        described_class.new(
+          { topic_id: topic_with_tags.id, post_numbers: [post1.post_number] },
+          bot_user: bot_user,
+          llm: llm,
+          context: {
+            user: admin,
+          },
+        )
+
+      results = tool.invoke
+      expect(results[:description]).to eq("Topic not found")
+    end
+
     it "can read specific posts" do
       tool =
         described_class.new(
diff --git a/spec/system/ai_bot/persona_spec.rb b/spec/system/ai_bot/persona_spec.rb
index f8d9135c..256599a3 100644
--- a/spec/system/ai_bot/persona_spec.rb
+++ b/spec/system/ai_bot/persona_spec.rb
@@ -53,7 +53,7 @@ RSpec.describe "AI personas", type: :system, js: true do
     expect(persona.name).to eq("Test Persona")
     expect(persona.description).to eq("I am a test persona")
     expect(persona.system_prompt).to eq("You are a helpful bot")
-    expect(persona.tools).to eq(["Read"])
+    expect(persona.tools).to eq([["Read", { "read_private" => nil }]])
   end
 
   it "will not allow deletion or editing of system personas" do