Merge pull request #124 from merefield/token_based_quota

FEATURE: add option to manage quotas by token

app/jobs/scheduled/chatbot_quota_reset.rb

@@ -5,20 +5,6 @@ class ::Jobs::ChatbotQuotaReset < ::Jobs::Scheduled
   every 1.week
 
   def execute(args)
-
-    ::User.all.each do |u|
-      if current_record = UserCustomField.find_by(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_QUERIES_CUSTOM_FIELD)
-        current_record.value = "0"
-        current_record.save!
-      else
-        current_queries = "0"
-        UserCustomField.create!(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_QUERIES_CUSTOM_FIELD, value: current_queries)
-      end
-
-      if current_record = UserCustomField.find_by(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_QUERIES_QUOTA_REACH_ESCALATION_DATE_CUSTOM_FIELD)
-        current_record.delete
-      end
-    end
-
+    ::DiscourseChatbot::Bot.new.reset_all_quotas
   end
 end

config/locales/server.en.yml

@@ -41,6 +41,7 @@ en:
     chatbot_quick_access_talk_button: "Display floating button which links to direct discussion one-to-one with the bot on specified channel type"
     chatbot_quick_access_talk_button_bot_icon: "Name of icon used for floating quick bot access button.  If blank it will use bot user's avatar instead. (if icon make sure it's permitted and added to svg_icon_subset if necessary)"
     chatbot_quick_access_bot_kicks_off: "Quick access launches a convesation in which the bot speaks first based on provided text at <a target='_blank' rel='noopener' href='/admin/customize/site_texts?q=chatbot.quick_access_kick_off.announcement'>Customize Text</a>"
+    chatbot_quota_basis: "Choose number of interactions ('queries') or tokens.  Tokens allows a more fine grained system that takes into account tool call cost, e.g. generating images."
     chatbot_quota_high_trust: "The allowed number of bot responses allowed by prompting user per week - high trust groups"
     chatbot_quota_medium_trust: "The allowed number of bot responses allowed by prompting user per week - medium trust groups"
     chatbot_quota_low_trust: "The allowed number of bot responses allowed by prompting user per week - low trust groups"
@@ -79,11 +80,16 @@ en:
     chatbot_escalate_to_staff_max_history: "(Chat only) number of chat messages included in transcript added to escalation PM"
     chatbot_user_fields_collection: "(EXPERIMENTAL) Collect empty user fields from the user as part of the conversation"
     chatbot_news_api_token: "News API token for news (if left blank, news will never be searched)<a target='_blank' rel='noopener' href='https://newsapi.org/'>Get one at NewsAPI.org</a>"
+    chatbot_news_api_call_token_cost: "The notional token cost used for each call to the News API against users' overall quota."
     chatbot_firecrawl_api_token: "Firecrawl API token for crawling remote websites.  If left blank, crawling will not be available. <a target='_blank' rel='noopener' href='https://www.firecrawl.dev/'>Get one at https://www.firecrawl.dev/</a>"
+    chatbot_firecrawl_api_call_token_cost: "The notional token cost used for each call to the Firecrawl API against users' overall quota."
     chatbot_jina_api_token: "Jina API token for web crawl and search.  <a target='_blank' rel='noopener' href='https://jina.ai'>Get one at https://jina.ai</a>. Alternative to Firecrawl and Serp.  If Firecrawl API token is populated, it will be used in preference to Jina for crawling. Ditto Serp API for Searching."
+    chatbot_jina_api_token_cost_multiplier: "The notional token cost multiplier used for each call to the Jina API against users' overall quota.  A Users's cost is the amount of text returned multiplied by this number."
     chatbot_function_response_char_limit: "The maximum number of characters taken from the response from web crawl in order to mitigate the risk of a token breach when later including result in LLM call"
     chatbot_serp_api_key: "Serp API token for google search (if left blank, google will never be searched). <a target='_blank' rel='noopener' href='https://serpapi.com/'>Get one at SerpAPI.com</a>"
+    chatbot_serp_api_call_token_cost: "The notional token cost used for each call to the Serp API against users' overall quota."
     chatbot_marketstack_key: "Marketstack API key for stock price information (if left blank, Marketstack will never be queried).<a target='_blank' rel='noopener' href='https://marketstack.com/'>Get one at MarketStack.com</a>"
+    chatbot_marketstack_api_call_token_cost: "The notional token cost used for each call to the Marketstack API against users' overall quota."
     chatbot_enable_verbose_console_logging: "Enable response retrieval progress logging to console to help debug issues"
     chatbot_enable_verbose_rails_logging: "Enable response retrieval progress logging to rails logs to help debug issues. 'api_calls_only' restricts this to just API calls, 'all' logs all progress"
     chatbot_verbose_rails_logging_destination_level: "Choose which category of logs to send verbose logs to.  'warn' is useful in Production as 'info' logs are not exposed at /logs."
@@ -113,6 +119,11 @@ en:
       quick_access_kick_off:
         announcement: "I must greet @%{username} warmly and kick off any questions I want answered."
       function:
+        remaining_bot_quota:
+          description: |
+            Get the current number of tokens or queries you have left for use of the bot this week
+          answer: "The user has %{quota} bot %{units} left for use this week.  There are %{days_remaining} before it resets!\n\n"
+          error:  "There was an error when trying to retrieve your remaining bot quota: '%{error}'"
         user_information:
           name: "user_information_for_%{user_field}"
           system_message:

config/settings.yml

@@ -197,6 +197,13 @@ plugins:
     default: ''
     type: category_list
     list_type: "compact"
+  chatbot_quota_basis:
+    client: false
+    default: 'queries'
+    type: enum
+    choices:
+      - 'queries'
+      - 'tokens'
   chatbot_quota_high_trust:
     default: 100
     client: false
@@ -373,21 +380,36 @@ plugins:
   chatbot_news_api_token:
     client: false
     default: ''
+  chatbot_news_api_call_token_cost:
+    client: false
+    default: 10000
   chatbot_firecrawl_api_token:
     client: false
     default: ''
+  chatbot_firecrawl_api_call_token_cost:
+    client: false
+    default: 10000
   chatbot_jina_api_token:
     client: false
     default: ''
+  chatbot_jina_api_token_cost_multiplier:
+    client: false
+    default: 10000
   chatbot_function_response_char_limit:
     client: false
     default: 350000
   chatbot_serp_api_key:
     client: false
     default: ''
+  chatbot_serp_api_call_token_cost:
+    client: false
+    default: 10000
   chatbot_marketstack_key:
     client: false
     default: ''
+  chatbot_marketstack_api_call_token_cost:
+    client: false
+    default: 10000
   chatbot_enable_verbose_console_logging:
     client: false
     default: false

lib/discourse_chatbot/bot.rb

@@ -8,9 +8,61 @@ def get_response(prompt, opts)
     end
 
     def ask(opts)
+      user_id = opts[:user_id]
       content = opts[:type] == POST ? PostPromptUtils.create_prompt(opts) : MessagePromptUtils.create_prompt(opts)
 
       response = get_response(content, opts)
+
+      consume_quota(opts[:user_id], response[:total_tokens])
+      response
+    end
+
+    def consume_quota(user_id, token_usage)
+      return if token_usage == 0
+
+      remaining_quota_field_name = SiteSetting.chatbot_quota_basis == "queries" ? CHATBOT_REMAINING_QUOTA_QUERIES_CUSTOM_FIELD : CHATBOT_REMAINING_QUOTA_TOKENS_CUSTOM_FIELD
+      deduction = SiteSetting.chatbot_quota_basis == "queries" ? 1 : token_usage
+
+      current_record = UserCustomField.find_by(user_id: user_id, name: remaining_quota_field_name)
+
+      if current_record.present?
+        remaining_quota = current_record.value.to_i - deduction
+        current_record.value = remaining_quota.to_s
+      else
+        max_quota = ::DiscourseChatbot::EventEvaluation.new.get_max_quota(user_id)
+        current_record = UserCustomField.create!(user_id: user_id, name: remaining_quota_field_name, value: max_quota.to_s)
+        remaining_quota = current_record.value.to_i - deduction
+        current_record.value = remaining_quota.to_s
+      end
+      current_record.save!
+    end
+
+    def reset_all_quotas
+      ::User.all.each do |u|
+        max_quota = ::DiscourseChatbot::EventEvaluation.new.get_max_quota(u.id)
+
+        current_record = UserCustomField.find_by(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_REMAINING_QUOTA_QUERIES_CUSTOM_FIELD)
+
+        if current_record.present?
+          current_record.value = max_quota.to_s
+          current_record.save!
+        else
+          UserCustomField.create!(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_REMAINING_QUOTA_QUERIES_CUSTOM_FIELD, value: max_quota.to_s)
+        end
+
+        current_record = UserCustomField.find_by(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_REMAINING_QUOTA_TOKENS_CUSTOM_FIELD)
+
+        if current_record.present?
+          current_record.value = max_quota.to_s
+          current_record.save!
+        else
+          UserCustomField.create!(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_REMAINING_QUOTA_TOKENS_CUSTOM_FIELD, value: max_quota.to_s)
+        end
+
+        if current_record = UserCustomField.find_by(user_id: u.id, name: ::DiscourseChatbot::CHATBOT_QUERIES_QUOTA_REACH_ESCALATION_DATE_CUSTOM_FIELD)
+          current_record.delete
+        end
+      end
     end
   end
 end

lib/discourse_chatbot/bots/open_ai_bot_base.rb

@@ -41,6 +41,7 @@ def initialize(opts)
       end
 
       @model_name = get_model(opts)
+      @total_tokens = 0
     end
 
     def get_response(prompt, opts)

lib/discourse_chatbot/bots/open_ai_bot_basic.rb

@@ -43,6 +43,9 @@ def get_response(prompt, opts)
           parameters: parameters
         )
 
+        token_usage = res.dig("usage", "total_tokens")
+        @total_tokens += token_usage
+
         {
           reply: response.dig("choices", 0, "message", "content"),
           inner_thoughts: nil

lib/discourse_chatbot/bots/open_ai_bot_rag.rb

@@ -21,7 +21,8 @@
 "DiscourseChatbot::NewsFunction",
 "DiscourseChatbot::UserFieldFunction",
 "DiscourseChatbot::EscalateToStaffFunction",
-"DiscourseChatbot::CalculatorFunction"]
+"DiscourseChatbot::CalculatorFunction",
+"DiscourseChatbot::RemainingQuotaFunction"]
 
 module ::DiscourseChatbot
 
@@ -66,7 +67,8 @@ def get_response(prompt, opts)
 
       {
         reply: res["choices"][0]["message"]["content"],
-        inner_thoughts: @inner_thoughts
+        inner_thoughts: @inner_thoughts,
+        total_tokens: @total_tokens
       }
     end
 
@@ -100,6 +102,7 @@ def get_system_message_suffix(opts)
     end
 
     def merge_functions(opts)
+      quota_function = ::DiscourseChatbot::RemainingQuotaFunction.new
       calculator_function = ::DiscourseChatbot::CalculatorFunction.new
       wikipedia_function = ::DiscourseChatbot::WikipediaFunction.new
       news_function = ::DiscourseChatbot::NewsFunction.new
@@ -147,6 +150,7 @@ def merge_functions(opts)
         end
       end
 
+      functions << quota_function
       functions << forum_search_function if forum_search_function
       functions << vision_function if vision_function
       functions << paint_function if SiteSetting.chatbot_support_picture_creation
@@ -218,6 +222,9 @@ def create_chat_completion(messages, use_functions = true, iteration)
           parameters: parameters
         )
 
+        token_usage = res.dig("usage", "total_tokens")
+        @total_tokens += token_usage
+
         ::DiscourseChatbot.progress_debug_message <<~EOS
           +++++++++++++++++++++++++++++++++++++++
           The llm responded with
@@ -345,15 +352,17 @@ def call_function(func_name, args_str, opts)
         +++++++++++++++++++++++++++++++++++++++
       EOS
       begin
+        token_usage = 0
         args = JSON.parse(args_str)
         func = @func_mapping[func_name]
-        if ["escalate_to_staff"].include?(func_name)
-          res = func.process(args, opts)
+        if ["escalate_to_staff", "remaining_bot_quota"].include?(func_name)
+          res, token_usage = func.process(args, opts).values_at(:answer, :token_usage)
         elsif ["vision"].include?(func_name)
-          res = func.process(args, opts, @client)
+          res, token_usage = func.process(args, opts, @client).values_at(:answer, :token_usage)
         else
-          res = func.process(args)
+          res, token_usage = func.process(args).values_at(:answer, :token_usage)
         end
+        @total_tokens += token_usage
         res
        rescue
          I18n.t("chatbot.prompt.rag.call_function.error")

lib/discourse_chatbot/event_evaluation.rb

@@ -26,37 +26,44 @@ def trust_level(user_id)
     end
 
     def over_quota(user_id)
-      max_quota = 0
+      max_quota = get_max_quota(user_id)
+      remaining_quota_field_name =  SiteSetting.chatbot_quota_basis == "queries" ? CHATBOT_REMAINING_QUOTA_QUERIES_CUSTOM_FIELD : CHATBOT_REMAINING_QUOTA_TOKENS_CUSTOM_FIELD
+      remaining_quota = get_remaining_quota(user_id, remaining_quota_field_name)
+
+      if remaining_quota.nil?
+        UserCustomField.create!(user_id: user_id, name: remaining_quota_field_name, value: max_quota.to_s)
+        remaining_quota = max_quota
+      end
+
+      breach = remaining_quota < 0
+      escalate_as_required(user_id) if breach
+      breach
+    end
 
+    def get_remaining_quota(user_id, remaining_quota_field_name)
+      UserCustomField.find_by(user_id: user_id, name: remaining_quota_field_name)&.value.to_i
+    end
+
+    def get_max_quota(user_id)
+      max_quota = 0
       GroupUser.where(user_id: user_id).each do |gu|
-        if SiteSetting.chatbot_high_trust_groups.split('|').include? gu.group_id.to_s
-          max_quota = SiteSetting.chatbot_quota_high_trust if max_quota < SiteSetting.chatbot_quota_high_trust
+        if SiteSetting.chatbot_low_trust_groups.split('|').include? gu.group_id.to_s
+          max_quota = SiteSetting.chatbot_quota_low_trust if max_quota < SiteSetting.chatbot_quota_low_trust
         end
         if SiteSetting.chatbot_medium_trust_groups.split('|').include? gu.group_id.to_s
           max_quota = SiteSetting.chatbot_quota_medium_trust if max_quota < SiteSetting.chatbot_quota_medium_trust
         end
-        if SiteSetting.chatbot_low_trust_groups.split('|').include? gu.group_id.to_s
-          max_quota = SiteSetting.chatbot_quota_low_trust if max_quota < SiteSetting.chatbot_quota_low_trust
+        if SiteSetting.chatbot_high_trust_groups.split('|').include? gu.group_id.to_s
+          max_quota = SiteSetting.chatbot_quota_high_trust if max_quota < SiteSetting.chatbot_quota_high_trust
         end
       end
 
       # deal with 'everyone' group
-      max_quota = SiteSetting.chatbot_quota_high_trust if SiteSetting.chatbot_high_trust_groups.split('|').include?("0") && max_quota < SiteSetting.chatbot_quota_high_trust
-      max_quota = SiteSetting.chatbot_quota_medium_trust if SiteSetting.chatbot_medium_trust_groups.split('|').include?("0") && max_quota < SiteSetting.chatbot_quota_medium_trust
       max_quota = SiteSetting.chatbot_quota_low_trust if SiteSetting.chatbot_low_trust_groups.split('|').include?("0") && max_quota < SiteSetting.chatbot_quota_low_trust
+      max_quota = SiteSetting.chatbot_quota_medium_trust if SiteSetting.chatbot_medium_trust_groups.split('|').include?("0") && max_quota < SiteSetting.chatbot_quota_medium_trust
+      max_quota = SiteSetting.chatbot_quota_high_trust if SiteSetting.chatbot_high_trust_groups.split('|').include?("0") && max_quota < SiteSetting.chatbot_quota_high_trust
 
-      if current_record = UserCustomField.find_by(user_id: user_id, name: CHATBOT_QUERIES_CUSTOM_FIELD)
-        current_queries = current_record.value.to_i + 1
-        current_record.value = current_queries.to_s
-        current_record.save!
-      else
-        current_queries = 1
-        UserCustomField.create!(user_id: user_id, name: CHATBOT_QUERIES_CUSTOM_FIELD, value: current_queries)
-      end
-
-      breach = current_queries > max_quota
-      escalate_as_required(user_id) if breach
-      breach
+      max_quota
     end
 
     def escalate_as_required(user_id)

lib/discourse_chatbot/functions/calculator_function.rb

@@ -27,9 +27,15 @@ def process(args)
       begin
         super(args)
 
-        ::SafeRuby.eval(args[parameters[0][:name]], timeout: 5)
+        {
+          answer: ::SafeRuby.eval(args[parameters[0][:name]], timeout: 5),
+          token_usage: 0
+        }
       rescue
-        I18n.t("chatbot.prompt.function.calculator.error", parameter: args[parameters[0][:name]])
+        {
+          answer: I18n.t("chatbot.prompt.function.calculator.error", parameter: args[parameters[0][:name]]),
+          token_usage: 0
+        }
       end
     end
   end

lib/discourse_chatbot/functions/coords_from_location_description_search.rb

@@ -35,9 +35,15 @@ def process(args)
 
         response = I18n.t("chatbot.prompt.function.return_coords_from_location_description.answer_summary", query: query, coords: coords)
 
-        response
+        {
+          answer: response,
+          token_usage: 0
+        }
       rescue
-        I18n.t("chatbot.prompt.function.return_coords_from_location_description.error", query: args[parameters[0][:name]])
+        {
+          answer: I18n.t("chatbot.prompt.function.return_coords_from_location_description.error", query: args[parameters[0][:name]]),
+          token_usage: 0
+        }
       end
     end
   end

lib/discourse_chatbot/functions/escalate_to_staff_function.rb

@@ -69,8 +69,15 @@ def process(args, opts)
         else
           response = I18n.t("chatbot.prompt.function.escalate_to_staff.no_escalation_groups")
         end
+        {
+          answer: response,
+          token_usage: 0
+        }
       rescue
-        I18n.t("chatbot.prompt.function.escalate_to_staff.error", parameter: args[parameters[0][:name]])
+        {
+          answer: I18n.t("chatbot.prompt.function.escalate_to_staff.error", parameter: args[parameters[0][:name]]),
+          token_usage: 0
+        }
       end
     end
 

lib/discourse_chatbot/functions/forum_get_user_address_function.rb

@@ -34,9 +34,15 @@ def process(args)
 
         response = I18n.t("chatbot.prompt.function.forum_get_user_address.answer_summary", username: username, address: result.address, latitude: result.latitude, longitude: result.longitude)
 
-        response
+        {
+          answer: response, 
+          token_usage: 0
+        }
       rescue
-        I18n.t("chatbot.prompt.function.forum_get_user_address.error")
+        {
+          answer: I18n.t("chatbot.prompt.function.forum_get_user_address.error"),
+          token_usage: 0
+        }
       end
     end
   end