discourse-ai/app/models/llm_quota.rb

# frozen_string_literal: true

class LlmQuota < ActiveRecord::Base
  self.table_name = "llm_quotas"

  belongs_to :group
  belongs_to :llm_model
  has_many :llm_quota_usages

  validates :group_id, presence: true
  # we can not validate on create cause it breaks build
  validates :llm_model_id, presence: true, on: :update
  validates :duration_seconds, presence: true, numericality: { greater_than: 0 }
  validates :max_tokens, numericality: { only_integer: true, greater_than: 0, allow_nil: true }
  validates :max_usages, numericality: { greater_than: 0, allow_nil: true }

  validate :at_least_one_limit

  def self.check_quotas!(llm, user)
    return true if user.blank?
    quotas = joins(:group).where(llm_model: llm).where(group: user.groups)

    return true if quotas.empty?
    errors =
      quotas.map do |quota|
        usage = LlmQuotaUsage.find_or_create_for(user: user, llm_quota: quota)
        begin
          usage.check_quota!
          nil
        rescue LlmQuotaUsage::QuotaExceededError => e
          e
        end
      end

    return if errors.include?(nil)

    raise errors.first
  end

  def self.log_usage(llm, user, input_tokens, output_tokens)
    return if user.blank?

    quotas = joins(:group).where(llm_model: llm).where(group: user.groups)

    quotas.each do |quota|
      usage = LlmQuotaUsage.find_or_create_for(user: user, llm_quota: quota)
      usage.increment_usage!(input_tokens: input_tokens, output_tokens: output_tokens)
    end
  end

  def available_tokens
    max_tokens
  end

  def available_usages
    max_usages
  end

  private

  def at_least_one_limit
    if max_tokens.nil? && max_usages.nil?
      errors.add(:base, I18n.t("discourse_ai.errors.quota_required"))
    end
  end
end

# == Schema Information
#
# Table name: llm_quotas
#
#  id               :bigint           not null, primary key
#  group_id         :bigint           not null
#  llm_model_id     :bigint           not null
#  max_tokens       :integer
#  max_usages       :integer
#  duration_seconds :integer          not null
#  created_at       :datetime         not null
#  updated_at       :datetime         not null
#
# Indexes
#
#  index_llm_quotas_on_group_id_and_llm_model_id  (group_id,llm_model_id) UNIQUE
#  index_llm_quotas_on_llm_model_id               (llm_model_id)
#