mirror of
https://github.com/discourse/discourse-ai.git
synced 2025-03-03 15:59:59 +00:00
Adds a comprehensive quota management system for LLM models that allows: - Setting per-group (applied per user in the group) token and usage limits with configurable durations - Tracking and enforcing token/usage limits across user groups - Quota reset periods (hourly, daily, weekly, or custom) - Admin UI for managing quotas with real-time updates This system provides granular control over LLM API usage by allowing admins to define limits on both total tokens and number of requests per group. Supports multiple concurrent quotas per model and automatically handles quota resets. Co-authored-by: Keegan George <kgeorge13@gmail.com>
170 lines
4.4 KiB
Ruby
170 lines
4.4 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
class LlmModel < ActiveRecord::Base
|
|
FIRST_BOT_USER_ID = -1200
|
|
BEDROCK_PROVIDER_NAME = "aws_bedrock"
|
|
|
|
has_many :llm_quotas, dependent: :destroy
|
|
belongs_to :user
|
|
|
|
validates :display_name, presence: true, length: { maximum: 100 }
|
|
validates :tokenizer, presence: true, inclusion: DiscourseAi::Completions::Llm.tokenizer_names
|
|
validates :provider, presence: true, inclusion: DiscourseAi::Completions::Llm.provider_names
|
|
validates :url, presence: true, unless: -> { provider == BEDROCK_PROVIDER_NAME }
|
|
validates_presence_of :name, :api_key
|
|
validates :max_prompt_tokens, numericality: { greater_than: 0 }
|
|
validate :required_provider_params
|
|
scope :in_use,
|
|
-> do
|
|
model_ids = DiscourseAi::Configuration::LlmEnumerator.global_usage.keys
|
|
where(id: model_ids)
|
|
end
|
|
|
|
def self.provider_params
|
|
{
|
|
aws_bedrock: {
|
|
access_key_id: :text,
|
|
region: :text,
|
|
disable_native_tools: :checkbox,
|
|
},
|
|
anthropic: {
|
|
disable_native_tools: :checkbox,
|
|
},
|
|
open_ai: {
|
|
organization: :text,
|
|
disable_native_tools: :checkbox,
|
|
disable_streaming: :checkbox,
|
|
},
|
|
mistral: {
|
|
disable_native_tools: :checkbox,
|
|
},
|
|
google: {
|
|
disable_native_tools: :checkbox,
|
|
},
|
|
azure: {
|
|
disable_native_tools: :checkbox,
|
|
},
|
|
hugging_face: {
|
|
disable_system_prompt: :checkbox,
|
|
},
|
|
vllm: {
|
|
disable_system_prompt: :checkbox,
|
|
},
|
|
ollama: {
|
|
disable_system_prompt: :checkbox,
|
|
enable_native_tool: :checkbox,
|
|
disable_streaming: :checkbox,
|
|
},
|
|
open_router: {
|
|
disable_native_tools: :checkbox,
|
|
provider_order: :text,
|
|
provider_quantizations: :text,
|
|
disable_streaming: :checkbox,
|
|
},
|
|
}
|
|
end
|
|
|
|
def to_llm
|
|
DiscourseAi::Completions::Llm.proxy(identifier)
|
|
end
|
|
|
|
def identifier
|
|
"custom:#{id}"
|
|
end
|
|
|
|
def toggle_companion_user
|
|
return if name == "fake" && Rails.env.production?
|
|
|
|
enable_check = SiteSetting.ai_bot_enabled && enabled_chat_bot
|
|
|
|
if enable_check
|
|
if !user
|
|
next_id = DB.query_single(<<~SQL).first
|
|
SELECT min(id) - 1 FROM users
|
|
SQL
|
|
|
|
new_user =
|
|
User.new(
|
|
id: [FIRST_BOT_USER_ID, next_id].min,
|
|
email: "no_email_#{SecureRandom.hex}",
|
|
name: name.titleize,
|
|
username: UserNameSuggester.suggest(name),
|
|
active: true,
|
|
approved: true,
|
|
admin: true,
|
|
moderator: true,
|
|
trust_level: TrustLevel[4],
|
|
)
|
|
new_user.save!(validate: false)
|
|
self.update!(user: new_user)
|
|
else
|
|
user.active = true
|
|
user.save!(validate: false)
|
|
end
|
|
elsif user
|
|
# will include deleted
|
|
has_posts = DB.query_single("SELECT 1 FROM posts WHERE user_id = #{user.id} LIMIT 1").present?
|
|
|
|
if has_posts
|
|
user.update!(active: false) if user.active
|
|
else
|
|
user.destroy!
|
|
self.update!(user: nil)
|
|
end
|
|
end
|
|
end
|
|
|
|
def tokenizer_class
|
|
tokenizer.constantize
|
|
end
|
|
|
|
def lookup_custom_param(key)
|
|
provider_params&.dig(key)
|
|
end
|
|
|
|
def seeded?
|
|
id.present? && id < 0
|
|
end
|
|
|
|
def api_key
|
|
if seeded?
|
|
env_key = "DISCOURSE_AI_SEEDED_LLM_API_KEY_#{id.abs}"
|
|
ENV[env_key] || self[:api_key]
|
|
else
|
|
self[:api_key]
|
|
end
|
|
end
|
|
|
|
private
|
|
|
|
def required_provider_params
|
|
return if provider != BEDROCK_PROVIDER_NAME
|
|
|
|
%w[access_key_id region].each do |field|
|
|
if lookup_custom_param(field).blank?
|
|
errors.add(:base, I18n.t("discourse_ai.llm_models.missing_provider_param", param: field))
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
# == Schema Information
|
|
#
|
|
# Table name: llm_models
|
|
#
|
|
# id :bigint not null, primary key
|
|
# display_name :string
|
|
# name :string not null
|
|
# provider :string not null
|
|
# tokenizer :string not null
|
|
# max_prompt_tokens :integer not null
|
|
# created_at :datetime not null
|
|
# updated_at :datetime not null
|
|
# url :string
|
|
# api_key :string
|
|
# user_id :integer
|
|
# enabled_chat_bot :boolean default(FALSE), not null
|
|
# provider_params :jsonb
|
|
# vision_enabled :boolean default(FALSE), not null
|
|
#
|