ruby_llm-agents 1.3.4 → 2.1.0

data/app/models/ruby_llm/agents/tenant/budgetable.rb

@@ -170,17 +170,30 @@ module RubyLLM
           effective_enforcement == :soft
         end
 
-        # Check if within budget for a specific type
+        # Check if within budget for a specific type using counter columns
         #
         # @param type [Symbol] :daily_cost, :monthly_cost, :daily_tokens,
         #   :monthly_tokens, :daily_executions, :monthly_executions
         # @return [Boolean]
         def within_budget?(type: :daily_cost)
-          status = budget_status
-          return true unless status[:enabled]
+          return true unless budgets_enabled?
 
-          key = budget_status_key(type)
-          (status.dig(key, :percentage_used) || 0) < 100
+          case type
+          when :daily_cost
+            within_daily_cost_budget?
+          when :monthly_cost
+            within_monthly_cost_budget?
+          when :daily_tokens
+            within_daily_token_budget?
+          when :monthly_tokens
+            within_monthly_token_budget?
+          when :daily_executions
+            within_daily_execution_budget?
+          when :monthly_executions
+            within_monthly_execution_budget?
+          else
+            true
+          end
         end
 
         # Get remaining budget for a specific type
@@ -188,9 +201,20 @@ module RubyLLM
         # @param type [Symbol] Budget type (see #within_budget?)
         # @return [Numeric, nil]
         def remaining_budget(type: :daily_cost)
-          status = budget_status
-          key = budget_status_key(type)
-          status.dig(key, :remaining)
+          case type
+          when :daily_cost
+            effective_daily_limit && (ensure_daily_reset!; effective_daily_limit - daily_cost_spent)
+          when :monthly_cost
+            effective_monthly_limit && (ensure_monthly_reset!; effective_monthly_limit - monthly_cost_spent)
+          when :daily_tokens
+            effective_daily_token_limit && (ensure_daily_reset!; effective_daily_token_limit - daily_tokens_used)
+          when :monthly_tokens
+            effective_monthly_token_limit && (ensure_monthly_reset!; effective_monthly_token_limit - monthly_tokens_used)
+          when :daily_executions
+            effective_daily_execution_limit && (ensure_daily_reset!; effective_daily_execution_limit - daily_executions_count)
+          when :monthly_executions
+            effective_monthly_execution_limit && (ensure_monthly_reset!; effective_monthly_execution_limit - monthly_executions_count)
+          end
         end
 
         # Check budget and raise if exceeded (for hard enforcement)
@@ -198,14 +222,98 @@ module RubyLLM
         # @param agent_type [String] The agent class name
         # @raise [BudgetExceededError] If hard enforcement and over budget
         def check_budget!(agent_type = nil)
-          BudgetTracker.check_budget!(agent_type || "Unknown", tenant_id: tenant_id)
+          return unless budgets_enabled?
+          return unless hard_enforcement?
+
+          ensure_daily_reset!
+          ensure_monthly_reset!
+
+          if effective_daily_limit && daily_cost_spent >= effective_daily_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_daily, effective_daily_limit, daily_cost_spent, tenant_id: tenant_id
+            )
+          end
+
+          if effective_monthly_limit && monthly_cost_spent >= effective_monthly_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_monthly, effective_monthly_limit, monthly_cost_spent, tenant_id: tenant_id
+            )
+          end
+
+          if effective_daily_token_limit && daily_tokens_used >= effective_daily_token_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_daily_tokens, effective_daily_token_limit, daily_tokens_used, tenant_id: tenant_id
+            )
+          end
+
+          if effective_monthly_token_limit && monthly_tokens_used >= effective_monthly_token_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_monthly_tokens, effective_monthly_token_limit, monthly_tokens_used, tenant_id: tenant_id
+            )
+          end
+
+          if effective_daily_execution_limit && daily_executions_count >= effective_daily_execution_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_daily_executions, effective_daily_execution_limit, daily_executions_count, tenant_id: tenant_id
+            )
+          end
+
+          if effective_monthly_execution_limit && monthly_executions_count >= effective_monthly_execution_limit
+            raise Reliability::BudgetExceededError.new(
+              :global_monthly_executions, effective_monthly_execution_limit, monthly_executions_count, tenant_id: tenant_id
+            )
+          end
         end
 
-        # Get full budget status from BudgetTracker
+        # Get full budget status using counter columns
         #
         # @return [Hash] Budget status with usage information
         def budget_status
-          BudgetTracker.status(tenant_id: tenant_id)
+          ensure_daily_reset!
+          ensure_monthly_reset!
+
+          {
+            enabled: budgets_enabled?,
+            enforcement: effective_enforcement,
+            global_daily: budget_status_for(effective_daily_limit, daily_cost_spent),
+            global_monthly: budget_status_for(effective_monthly_limit, monthly_cost_spent),
+            global_daily_tokens: budget_status_for(effective_daily_token_limit, daily_tokens_used),
+            global_monthly_tokens: budget_status_for(effective_monthly_token_limit, monthly_tokens_used),
+            global_daily_executions: budget_status_for(effective_daily_execution_limit, daily_executions_count),
+            global_monthly_executions: budget_status_for(effective_monthly_execution_limit, monthly_executions_count)
+          }
+        end
+
+        # Individual budget check methods
+
+        def within_daily_cost_budget?
+          ensure_daily_reset!
+          effective_daily_limit.nil? || daily_cost_spent < effective_daily_limit
+        end
+
+        def within_monthly_cost_budget?
+          ensure_monthly_reset!
+          effective_monthly_limit.nil? || monthly_cost_spent < effective_monthly_limit
+        end
+
+        def within_daily_token_budget?
+          ensure_daily_reset!
+          effective_daily_token_limit.nil? || daily_tokens_used < effective_daily_token_limit
+        end
+
+        def within_monthly_token_budget?
+          ensure_monthly_reset!
+          effective_monthly_token_limit.nil? || monthly_tokens_used < effective_monthly_token_limit
+        end
+
+        def within_daily_execution_budget?
+          ensure_daily_reset!
+          effective_daily_execution_limit.nil? || daily_executions_count < effective_daily_execution_limit
+        end
+
+        def within_monthly_execution_budget?
+          ensure_monthly_reset!
+          effective_monthly_execution_limit.nil? || monthly_executions_count < effective_monthly_execution_limit
         end
 
         # Convert to config hash for BudgetTracker
@@ -256,20 +364,20 @@ module RubyLLM
           (global_config&.dig(:per_agent_monthly) || {}).merge(per_agent_monthly || {})
         end
 
-        # Maps budget type to status key
+        # Builds a status hash for a single budget dimension
         #
-        # @param type [Symbol]
-        # @return [Symbol]
-        def budget_status_key(type)
-          case type
-          when :daily_cost then :global_daily
-          when :monthly_cost then :global_monthly
-          when :daily_tokens then :global_daily_tokens
-          when :monthly_tokens then :global_monthly_tokens
-          when :daily_executions then :global_daily_executions
-          when :monthly_executions then :global_monthly_executions
-          else :global_daily
-          end
+        # @param limit [Numeric, nil] The configured limit
+        # @param current [Numeric] The current usage
+        # @return [Hash, nil]
+        def budget_status_for(limit, current)
+          return nil unless limit
+
+          {
+            limit: limit,
+            current_spend: current,
+            remaining: [limit - current, 0].max,
+            percentage_used: (current.to_f / limit * 100).round(1)
+          }
         end
       end
     end

data/app/models/ruby_llm/agents/tenant/incrementable.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    class Tenant
+      # Provides atomic SQL increment of usage counters after each execution.
+      #
+      # @example Recording an execution
+      #   tenant.record_execution!(cost: 0.05, tokens: 1200)
+      #   tenant.record_execution!(cost: 0.01, tokens: 500, error: true)
+      #
+      # @api public
+      module Incrementable
+        extend ActiveSupport::Concern
+
+        # Records an execution by atomically incrementing all counter columns.
+        #
+        # @param cost [Numeric] The cost of the execution in USD
+        # @param tokens [Integer] The number of tokens used
+        # @param error [Boolean] Whether the execution was an error
+        # @return [void]
+        def record_execution!(cost:, tokens:, error: false)
+          ensure_daily_reset!
+          ensure_monthly_reset!
+
+          error_inc = error ? 1 : 0
+          status = error ? "error" : "success"
+
+          self.class.where(id: id).update_all(
+            self.class.sanitize_sql_array([
+              <<~SQL,
+                daily_cost_spent = daily_cost_spent + ?,
+                monthly_cost_spent = monthly_cost_spent + ?,
+                daily_tokens_used = daily_tokens_used + ?,
+                monthly_tokens_used = monthly_tokens_used + ?,
+                daily_executions_count = daily_executions_count + 1,
+                monthly_executions_count = monthly_executions_count + 1,
+                daily_error_count = daily_error_count + ?,
+                monthly_error_count = monthly_error_count + ?,
+                last_execution_at = ?,
+                last_execution_status = ?
+              SQL
+              cost.to_f, cost.to_f,
+              tokens.to_i, tokens.to_i,
+              error_inc, error_inc,
+              Time.current, status
+            ])
+          )
+
+          reload
+          check_soft_cap_alerts!
+        end
+
+        private
+
+        # Checks soft cap alerts after recording an execution.
+        #
+        # @return [void]
+        def check_soft_cap_alerts!
+          return unless soft_enforcement?
+
+          check_cost_alerts!
+          check_token_alerts!
+          check_execution_alerts!
+        end
+
+        # @return [void]
+        def check_cost_alerts!
+          if effective_daily_limit && daily_cost_spent >= effective_daily_limit
+            AlertManager.notify(:budget_soft_cap, {
+              tenant_id: tenant_id, type: :daily_cost,
+              limit: effective_daily_limit, total: daily_cost_spent
+            })
+          end
+          if effective_monthly_limit && monthly_cost_spent >= effective_monthly_limit
+            AlertManager.notify(:budget_soft_cap, {
+              tenant_id: tenant_id, type: :monthly_cost,
+              limit: effective_monthly_limit, total: monthly_cost_spent
+            })
+          end
+        end
+
+        # @return [void]
+        def check_token_alerts!
+          if effective_daily_token_limit && daily_tokens_used >= effective_daily_token_limit
+            AlertManager.notify(:token_soft_cap, {
+              tenant_id: tenant_id, type: :daily_tokens,
+              limit: effective_daily_token_limit, total: daily_tokens_used
+            })
+          end
+          if effective_monthly_token_limit && monthly_tokens_used >= effective_monthly_token_limit
+            AlertManager.notify(:token_soft_cap, {
+              tenant_id: tenant_id, type: :monthly_tokens,
+              limit: effective_monthly_token_limit, total: monthly_tokens_used
+            })
+          end
+        end
+
+        # @return [void]
+        def check_execution_alerts!
+          if effective_daily_execution_limit && daily_executions_count >= effective_daily_execution_limit
+            AlertManager.notify(:budget_soft_cap, {
+              tenant_id: tenant_id, type: :daily_executions,
+              limit: effective_daily_execution_limit, total: daily_executions_count
+            })
+          end
+          if effective_monthly_execution_limit && monthly_executions_count >= effective_monthly_execution_limit
+            AlertManager.notify(:budget_soft_cap, {
+              tenant_id: tenant_id, type: :monthly_executions,
+              limit: effective_monthly_execution_limit, total: monthly_executions_count
+            })
+          end
+        end
+      end
+    end
+  end
+end

data/app/models/ruby_llm/agents/tenant/resettable.rb

@@ -0,0 +1,128 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    class Tenant
+      # Handles lazy reset of usage counters when periods roll over,
+      # and provides refresh_counters! for reconciliation from the executions table.
+      #
+      # @api public
+      module Resettable
+        extend ActiveSupport::Concern
+
+        # Resets daily counters if the day has rolled over.
+        # Uses a WHERE guard to prevent race conditions with concurrent requests.
+        #
+        # @return [void]
+        def ensure_daily_reset!
+          return if daily_reset_date == Date.current
+
+          rows = self.class.where(id: id)
+            .where("daily_reset_date IS NULL OR daily_reset_date < ?", Date.current)
+            .update_all(
+              daily_cost_spent: 0,
+              daily_tokens_used: 0,
+              daily_executions_count: 0,
+              daily_error_count: 0,
+              daily_reset_date: Date.current
+            )
+
+          reload if rows > 0
+        end
+
+        # Resets monthly counters if the month has rolled over.
+        # Uses a WHERE guard to prevent race conditions with concurrent requests.
+        #
+        # @return [void]
+        def ensure_monthly_reset!
+          bom = Date.current.beginning_of_month
+          return if monthly_reset_date == bom
+
+          rows = self.class.where(id: id)
+            .where("monthly_reset_date IS NULL OR monthly_reset_date < ?", bom)
+            .update_all(
+              monthly_cost_spent: 0,
+              monthly_tokens_used: 0,
+              monthly_executions_count: 0,
+              monthly_error_count: 0,
+              monthly_reset_date: bom
+            )
+
+          reload if rows > 0
+        end
+
+        # Recalculates all counters from the source-of-truth executions table.
+        #
+        # Use when counters have drifted due to manual DB edits, failed writes,
+        # or after deleting/updating execution records.
+        #
+        # @return [void]
+        def refresh_counters!
+          today = Date.current
+          bom = today.beginning_of_month
+
+          daily_stats = aggregate_stats(
+            executions.where("created_at >= ?", today.beginning_of_day)
+          )
+          monthly_stats = aggregate_stats(
+            executions.where("created_at >= ?", bom.beginning_of_day)
+          )
+
+          last_exec = executions.order(created_at: :desc).pick(:created_at, :status)
+
+          update_columns(
+            daily_cost_spent:        daily_stats[:cost],
+            daily_tokens_used:       daily_stats[:tokens],
+            daily_executions_count:  daily_stats[:count],
+            daily_error_count:       daily_stats[:errors],
+            daily_reset_date:        today,
+
+            monthly_cost_spent:       monthly_stats[:cost],
+            monthly_tokens_used:      monthly_stats[:tokens],
+            monthly_executions_count: monthly_stats[:count],
+            monthly_error_count:      monthly_stats[:errors],
+            monthly_reset_date:       bom,
+
+            last_execution_at:     last_exec&.first,
+            last_execution_status: last_exec&.last
+          )
+
+          reload
+        end
+
+        class_methods do
+          # Refresh counters for all tenants
+          #
+          # @return [void]
+          def refresh_all_counters!
+            find_each(&:refresh_counters!)
+          end
+
+          # Refresh counters for active tenants only
+          #
+          # @return [void]
+          def refresh_active_counters!
+            active.find_each(&:refresh_counters!)
+          end
+        end
+
+        private
+
+        # Aggregates cost, tokens, count, and errors from an executions scope.
+        #
+        # @param scope [ActiveRecord::Relation]
+        # @return [Hash] { cost:, tokens:, count:, errors: }
+        def aggregate_stats(scope)
+          agg = scope.pick(
+            Arel.sql("COALESCE(SUM(total_cost), 0)"),
+            Arel.sql("COALESCE(SUM(total_tokens), 0)"),
+            Arel.sql("COUNT(*)"),
+            Arel.sql("COALESCE(SUM(CASE WHEN status = 'error' THEN 1 ELSE 0 END), 0)")
+          )
+
+          { cost: agg[0], tokens: agg[1], count: agg[2], errors: agg[3] }
+        end
+      end
+    end
+  end
+end

data/app/models/ruby_llm/agents/tenant/trackable.rb

@@ -50,11 +50,12 @@ module RubyLLM
           scope.sum(:total_cost) || 0
         end
 
-        # Returns today's cost
+        # Returns today's cost from counter columns
         #
         # @return [Float]
         def cost_today
-          cost(period: :today)
+          ensure_daily_reset!
+          daily_cost_spent
         end
 
         # Returns yesterday's cost
@@ -78,11 +79,12 @@ module RubyLLM
           cost(period: :last_week)
         end
 
-        # Returns this month's cost
+        # Returns this month's cost from counter columns
         #
         # @return [Float]
         def cost_this_month
-          cost(period: :this_month)
+          ensure_monthly_reset!
+          monthly_cost_spent
         end
 
         # Returns last month's cost
@@ -104,11 +106,12 @@ module RubyLLM
           scope.sum(:total_tokens) || 0
         end
 
-        # Returns today's token usage
+        # Returns today's token usage from counter columns
         #
         # @return [Integer]
         def tokens_today
-          tokens(period: :today)
+          ensure_daily_reset!
+          daily_tokens_used
         end
 
         # Returns yesterday's token usage
@@ -125,11 +128,12 @@ module RubyLLM
           tokens(period: :this_week)
         end
 
-        # Returns this month's token usage
+        # Returns this month's token usage from counter columns
         #
         # @return [Integer]
         def tokens_this_month
-          tokens(period: :this_month)
+          ensure_monthly_reset!
+          monthly_tokens_used
         end
 
         # Returns last month's token usage
@@ -151,11 +155,12 @@ module RubyLLM
           scope.count
         end
 
-        # Returns today's execution count
+        # Returns today's execution count from counter columns
         #
         # @return [Integer]
         def executions_today
-          execution_count(period: :today)
+          ensure_daily_reset!
+          daily_executions_count
         end
 
         # Returns yesterday's execution count
@@ -172,11 +177,12 @@ module RubyLLM
           execution_count(period: :this_week)
         end
 
-        # Returns this month's execution count
+        # Returns this month's execution count from counter columns
         #
         # @return [Integer]
         def executions_this_month
-          execution_count(period: :this_month)
+          ensure_monthly_reset!
+          monthly_executions_count
         end
 
         # Returns last month's execution count
@@ -186,6 +192,34 @@ module RubyLLM
           execution_count(period: :last_month)
         end
 
+        # Error count queries
+
+        # Returns today's error count from counter columns
+        #
+        # @return [Integer]
+        def errors_today
+          ensure_daily_reset!
+          daily_error_count
+        end
+
+        # Returns this month's error count from counter columns
+        #
+        # @return [Integer]
+        def errors_this_month
+          ensure_monthly_reset!
+          monthly_error_count
+        end
+
+        # Returns today's success rate from counter columns
+        #
+        # @return [Float] Percentage (0.0-100.0)
+        def success_rate_today
+          ensure_daily_reset!
+          return 100.0 if daily_executions_count.zero?
+
+          ((daily_executions_count - daily_error_count).to_f / daily_executions_count * 100).round(1)
+        end
+
         # Usage summaries
 
         # Returns a complete usage summary for the tenant

data/app/models/ruby_llm/agents/tenant.rb

@@ -7,7 +7,6 @@ module RubyLLM
     # Encapsulates all tenant-related functionality:
     # - Budget limits and enforcement (via Budgetable concern)
     # - Usage tracking: cost, tokens, executions (via Trackable concern)
-    # - API configuration per tenant (via Configurable concern)
     #
     # @example Creating a tenant
     #   Tenant.create!(
@@ -30,7 +29,6 @@ module RubyLLM
     #
     # @see Tenant::Budgetable
     # @see Tenant::Trackable
-    # @see Tenant::Configurable
     # @see LLMTenant
     # @api public
     class Tenant < ::ActiveRecord::Base
@@ -39,7 +37,8 @@ module RubyLLM
       # Include concerns for organized functionality
       include Tenant::Budgetable
       include Tenant::Trackable
-      include Tenant::Configurable
+      include Tenant::Resettable
+      include Tenant::Incrementable
 
       # Polymorphic association to user's tenant model (optional)
       # Allows linking to Organization, Account, or any ActiveRecord model

data/app/models/ruby_llm/agents/tenant_budget.rb

@@ -8,15 +8,18 @@ module RubyLLM
     # All functionality has been moved to the Tenant model with organized concerns.
     #
     # @example Migration path
-    #   # Old usage (still works)
+    #   # Old usage (still works but emits deprecation warning)
     #   TenantBudget.for_tenant("acme_corp")
-    #   TenantBudget.create!(tenant_id: "acme", daily_limit: 100)
     #
     #   # New usage (preferred)
     #   Tenant.for("acme_corp")
-    #   Tenant.create!(tenant_id: "acme", daily_limit: 100)
     #
     # @see Tenant
     TenantBudget = Tenant
+
+    ActiveSupport.deprecator.warn(
+      "RubyLLM::Agents::TenantBudget is deprecated. Use RubyLLM::Agents::Tenant instead. " \
+      "This alias will be removed in the next major version."
+    )
   end
 end