ruby_llm-agents 1.1.0 → 1.2.0

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

@@ -2,262 +2,21 @@
 
 module RubyLLM
   module Agents
-    # Database-backed budget configuration for multi-tenant environments
+    # @deprecated Use {Tenant} instead. This class will be removed in a future major version.
     #
-    # Stores per-tenant budget limits that override the global configuration.
-    # Supports runtime updates without application restarts.
-    # Supports cost-based (USD), token-based, and execution-based limits.
+    # TenantBudget is now an alias to Tenant for backward compatibility.
+    # All functionality has been moved to the Tenant model with organized concerns.
     #
-    # @!attribute [rw] tenant_id
-    #   @return [String] Unique identifier for the tenant
-    # @!attribute [rw] name
-    #   @return [String, nil] Human-readable name for the tenant
-    # @!attribute [rw] daily_limit
-    #   @return [BigDecimal, nil] Daily budget limit in USD
-    # @!attribute [rw] monthly_limit
-    #   @return [BigDecimal, nil] Monthly budget limit in USD
-    # @!attribute [rw] daily_token_limit
-    #   @return [Integer, nil] Daily token limit (across all models)
-    # @!attribute [rw] monthly_token_limit
-    #   @return [Integer, nil] Monthly token limit (across all models)
-    # @!attribute [rw] daily_execution_limit
-    #   @return [Integer, nil] Daily execution/call limit
-    # @!attribute [rw] monthly_execution_limit
-    #   @return [Integer, nil] Monthly execution/call limit
-    # @!attribute [rw] per_agent_daily
-    #   @return [Hash] Per-agent daily cost limits: { "AgentName" => limit }
-    # @!attribute [rw] per_agent_monthly
-    #   @return [Hash] Per-agent monthly cost limits: { "AgentName" => limit }
-    # @!attribute [rw] enforcement
-    #   @return [String] Enforcement mode: "none", "soft", or "hard"
-    # @!attribute [rw] inherit_global_defaults
-    #   @return [Boolean] Whether to fall back to global config for unset limits
-    # @!attribute [rw] tenant_record
-    #   @return [ActiveRecord::Base, nil] Polymorphic association to tenant model
+    # @example Migration path
+    #   # Old usage (still works)
+    #   TenantBudget.for_tenant("acme_corp")
+    #   TenantBudget.create!(tenant_id: "acme", daily_limit: 100)
     #
-    # @example Creating a tenant budget with cost, token, and execution limits
-    #   TenantBudget.create!(
-    #     tenant_id: "acme_corp",
-    #     name: "Acme Corporation",
-    #     daily_limit: 50.0,              # USD
-    #     monthly_limit: 500.0,           # USD
-    #     daily_token_limit: 1_000_000,
-    #     monthly_token_limit: 10_000_000,
-    #     daily_execution_limit: 500,
-    #     monthly_execution_limit: 10_000,
-    #     enforcement: "hard"
-    #   )
+    #   # New usage (preferred)
+    #   Tenant.for("acme_corp")
+    #   Tenant.create!(tenant_id: "acme", daily_limit: 100)
     #
-    # @example Fetching budget for a tenant object
-    #   budget = TenantBudget.for_tenant(organization)
-    #   budget.effective_daily_limit           # => 50.0 (cost)
-    #   budget.effective_daily_token_limit     # => 1_000_000 (tokens)
-    #   budget.effective_daily_execution_limit # => 500 (executions)
-    #
-    # @see RubyLLM::Agents::BudgetTracker
-    # @see RubyLLM::Agents::LLMTenant
-    # @api public
-    class TenantBudget < ::ActiveRecord::Base
-      self.table_name = "ruby_llm_agents_tenant_budgets"
-
-      # Valid enforcement modes
-      ENFORCEMENT_MODES = %w[none soft hard].freeze
-
-      # Polymorphic association to the tenant model (e.g., Organization, Account)
-      belongs_to :tenant_record, polymorphic: true, optional: true
-
-      # Validations
-      validates :tenant_id, presence: true, uniqueness: true
-      validates :enforcement, inclusion: { in: ENFORCEMENT_MODES }, allow_nil: true
-      validates :daily_limit, :monthly_limit,
-                numericality: { greater_than_or_equal_to: 0 }, allow_nil: true
-      validates :daily_token_limit, :monthly_token_limit,
-                numericality: { greater_than_or_equal_to: 0, only_integer: true }, allow_nil: true
-      validates :daily_execution_limit, :monthly_execution_limit,
-                numericality: { greater_than_or_equal_to: 0, only_integer: true }, allow_nil: true
-
-      # Finds a budget for the given tenant
-      #
-      # @param tenant [String, Object] The tenant identifier string or object with llm_tenant_id
-      # @return [TenantBudget, nil] The budget record or nil if not found
-      def self.for_tenant(tenant)
-        return nil if tenant.blank?
-
-        if tenant.respond_to?(:llm_tenant_id)
-          # Object with llm_tenant DSL - try polymorphic first, then tenant_id
-          find_by(tenant_record: tenant) || find_by(tenant_id: tenant.llm_tenant_id)
-        else
-          # String tenant_id
-          find_by(tenant_id: tenant.to_s)
-        end
-      end
-
-      # Finds or creates a budget for the given tenant
-      #
-      # @param tenant_id [String] The tenant identifier
-      # @param name [String, nil] Optional human-readable name
-      # @return [TenantBudget] The budget record
-      def self.for_tenant!(tenant_id, name: nil)
-        find_or_create_by!(tenant_id: tenant_id) do |budget|
-          budget.name = name
-        end
-      end
-
-      # Returns the display name (name or tenant_id fallback)
-      #
-      # @return [String] The name to display
-      def display_name
-        name.presence || tenant_id
-      end
-
-      # Returns the effective daily limit, considering inheritance
-      #
-      # @return [Float, nil] The daily limit or nil if not set
-      def effective_daily_limit
-        return daily_limit if daily_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_daily)
-      end
-
-      # Returns the effective monthly limit, considering inheritance
-      #
-      # @return [Float, nil] The monthly limit or nil if not set
-      def effective_monthly_limit
-        return monthly_limit if monthly_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_monthly)
-      end
-
-      # Returns the effective per-agent daily limit
-      #
-      # @param agent_type [String] The agent class name
-      # @return [Float, nil] The limit or nil if not set
-      def effective_per_agent_daily(agent_type)
-        limit = per_agent_daily&.dig(agent_type)
-        return limit if limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:per_agent_daily, agent_type)
-      end
-
-      # Returns the effective per-agent monthly limit
-      #
-      # @param agent_type [String] The agent class name
-      # @return [Float, nil] The limit or nil if not set
-      def effective_per_agent_monthly(agent_type)
-        limit = per_agent_monthly&.dig(agent_type)
-        return limit if limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:per_agent_monthly, agent_type)
-      end
-
-      # Returns the effective daily token limit, considering inheritance
-      #
-      # @return [Integer, nil] The daily token limit or nil if not set
-      def effective_daily_token_limit
-        return daily_token_limit if daily_token_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_daily_tokens)
-      end
-
-      # Returns the effective monthly token limit, considering inheritance
-      #
-      # @return [Integer, nil] The monthly token limit or nil if not set
-      def effective_monthly_token_limit
-        return monthly_token_limit if monthly_token_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_monthly_tokens)
-      end
-
-      # Returns the effective daily execution limit, considering inheritance
-      #
-      # @return [Integer, nil] The daily execution limit or nil if not set
-      def effective_daily_execution_limit
-        return daily_execution_limit if daily_execution_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_daily_executions)
-      end
-
-      # Returns the effective monthly execution limit, considering inheritance
-      #
-      # @return [Integer, nil] The monthly execution limit or nil if not set
-      def effective_monthly_execution_limit
-        return monthly_execution_limit if monthly_execution_limit.present?
-        return nil unless inherit_global_defaults
-
-        global_config&.dig(:global_monthly_executions)
-      end
-
-      # Returns the effective enforcement mode
-      #
-      # @return [Symbol] :none, :soft, or :hard
-      def effective_enforcement
-        return enforcement.to_sym if enforcement.present?
-        return :soft unless inherit_global_defaults
-
-        RubyLLM::Agents.configuration.budget_enforcement
-      end
-
-      # Checks if budget enforcement is enabled for this tenant
-      #
-      # @return [Boolean] true if enforcement is :soft or :hard
-      def budgets_enabled?
-        effective_enforcement != :none
-      end
-
-      # Returns a hash suitable for BudgetTracker
-      #
-      # @return [Hash] Budget configuration hash
-      def to_budget_config
-        {
-          enabled: budgets_enabled?,
-          enforcement: effective_enforcement,
-          # Cost limits
-          global_daily: effective_daily_limit,
-          global_monthly: effective_monthly_limit,
-          per_agent_daily: merged_per_agent_daily,
-          per_agent_monthly: merged_per_agent_monthly,
-          # Token limits
-          global_daily_tokens: effective_daily_token_limit,
-          global_monthly_tokens: effective_monthly_token_limit,
-          # Execution limits
-          global_daily_executions: effective_daily_execution_limit,
-          global_monthly_executions: effective_monthly_execution_limit
-        }
-      end
-
-      private
-
-      # Returns the global budgets configuration
-      #
-      # @return [Hash, nil] Global budget config
-      def global_config
-        RubyLLM::Agents.configuration.budgets
-      end
-
-      # Merges per-agent daily limits with global defaults
-      #
-      # @return [Hash] Merged per-agent daily limits
-      def merged_per_agent_daily
-        return per_agent_daily || {} unless inherit_global_defaults
-
-        (global_config&.dig(:per_agent_daily) || {}).merge(per_agent_daily || {})
-      end
-
-      # Merges per-agent monthly limits with global defaults
-      #
-      # @return [Hash] Merged per-agent monthly limits
-      def merged_per_agent_monthly
-        return per_agent_monthly || {} unless inherit_global_defaults
-
-        (global_config&.dig(:per_agent_monthly) || {}).merge(per_agent_monthly || {})
-      end
-    end
+    # @see Tenant
+    TenantBudget = Tenant
   end
 end

data/lib/generators/ruby_llm_agents/multi_tenancy_generator.rb

@@ -10,8 +10,11 @@ module RubyLlmAgents
   #   rails generate ruby_llm_agents:multi_tenancy
   #
   # This will create migrations for:
-  #   - ruby_llm_agents_tenant_budgets table for per-tenant budget configuration
-  #   - Adding tenant_id column to ruby_llm_agents_executions
+  #   - ruby_llm_agents_tenants table for per-tenant configuration
+  #   - Adding tenant columns to ruby_llm_agents_executions
+  #
+  # For users upgrading from an older version:
+  #   - Renames ruby_llm_agents_tenant_budgets to ruby_llm_agents_tenants
   #
   class MultiTenancyGenerator < ::Rails::Generators::Base
     include ::ActiveRecord::Generators::Migration
@@ -20,21 +23,31 @@ module RubyLlmAgents
 
     desc "Adds multi-tenancy support to RubyLLM::Agents"
 
-    def create_tenant_budgets_migration
-      if table_exists?(:ruby_llm_agents_tenant_budgets)
-        say_status :skip, "ruby_llm_agents_tenant_budgets table already exists", :yellow
+    def create_tenants_migration
+      if table_exists?(:ruby_llm_agents_tenants)
+        say_status :skip, "ruby_llm_agents_tenants table already exists", :yellow
         return
       end
 
-      migration_template(
-        "create_tenant_budgets_migration.rb.tt",
-        File.join(db_migrate_path, "create_ruby_llm_agents_tenant_budgets.rb")
-      )
+      if table_exists?(:ruby_llm_agents_tenant_budgets)
+        # Upgrade path: rename existing table
+        say_status :upgrade, "Renaming tenant_budgets to tenants", :blue
+        migration_template(
+          "rename_tenant_budgets_to_tenants_migration.rb.tt",
+          File.join(db_migrate_path, "rename_tenant_budgets_to_tenants.rb")
+        )
+      else
+        # Fresh install: create new table
+        migration_template(
+          "create_tenants_migration.rb.tt",
+          File.join(db_migrate_path, "create_ruby_llm_agents_tenants.rb")
+        )
+      end
     end
 
     def create_add_tenant_to_executions_migration
       if column_exists?(:ruby_llm_agents_executions, :tenant_id)
-        say_status :skip, "tenant_id column already exists", :yellow
+        say_status :skip, "tenant_id column already exists on executions", :yellow
         return
       end
 
@@ -50,23 +63,30 @@ module RubyLlmAgents
       say ""
       say "Next steps:"
       say "  1. Run: rails db:migrate"
-      say "  2. Configure multi-tenancy in your initializer:"
+      say "  2. Add llm_tenant to your tenant model:"
       say ""
-      say "     RubyLLM::Agents.configure do |config|"
-      say "       config.multi_tenancy_enabled = true"
-      say "       config.tenant_resolver = -> { Current.tenant&.id }"
+      say "     class Organization < ApplicationRecord"
+      say "       include RubyLLM::Agents::LLMTenant"
+      say ""
+      say "       llm_tenant id: :id,            # Method for tenant_id"
+      say "                  budget: true,       # Auto-create budget on creation"
+      say "                  limits: {           # Optional default limits"
+      say "                    daily_cost: 100,"
+      say "                    monthly_cost: 1000"
+      say "                  },"
+      say "                  enforcement: :hard  # :none, :soft, or :hard"
       say "     end"
       say ""
-      say "  3. Set Current.tenant in your ApplicationController"
+      say "  3. Pass tenant to agents:"
+      say ""
+      say "     MyAgent.call(prompt, tenant: current_organization)"
       say ""
-      say "  4. Create tenant budgets:"
+      say "  4. Query usage:"
       say ""
-      say "     RubyLLM::Agents::TenantBudget.create!("
-      say "       tenant_id: 'acme_corp',"
-      say "       daily_limit: 50.0,"
-      say "       monthly_limit: 500.0,"
-      say "       enforcement: 'hard'"
-      say "     )"
+      say "     tenant = RubyLLM::Agents::Tenant.for(organization)"
+      say "     tenant.cost_today        # => 12.34"
+      say "     tenant.tokens_this_month # => 50000"
+      say "     tenant.usage_summary     # => { cost: ..., tokens: ..., ... }"
       say ""
     end
 

data/lib/generators/ruby_llm_agents/templates/add_tenant_to_executions_migration.rb.tt

@@ -1,12 +1,13 @@
 # frozen_string_literal: true
 
-# Migration to add tenant_id column to executions for multi-tenancy support
+# Migration to add tenant columns to executions for multi-tenancy support
 #
-# This migration adds a tenant_id column to track which tenant each execution
+# This migration adds tenant columns to track which tenant each execution
 # belongs to, enabling:
 # - Filtering executions by tenant
 # - Tenant-scoped analytics and reporting
 # - Per-tenant budget tracking and circuit breakers
+# - Polymorphic association with tenant models (e.g., Organization, Account)
 #
 # Run with: rails db:migrate
 class AddTenantIdToRubyLLMAgentsExecutions < ActiveRecord::Migration<%= migration_version %>
@@ -14,10 +15,20 @@ class AddTenantIdToRubyLLMAgentsExecutions < ActiveRecord::Migration<%= migratio
     # Add tenant_id column (nullable for backward compatibility)
     add_column :ruby_llm_agents_executions, :tenant_id, :string
 
+    # Polymorphic association to tenant model (for llm_tenant DSL)
+    # Uses string type for tenant_record_id to support both integer and UUID primary keys
+    add_column :ruby_llm_agents_executions, :tenant_record_type, :string
+    add_column :ruby_llm_agents_executions, :tenant_record_id, :string
+
     # Add indexes for efficient tenant-scoped queries
     add_index :ruby_llm_agents_executions, :tenant_id
     add_index :ruby_llm_agents_executions, [:tenant_id, :created_at]
     add_index :ruby_llm_agents_executions, [:tenant_id, :agent_type]
     add_index :ruby_llm_agents_executions, [:tenant_id, :status]
+
+    # Index for polymorphic tenant record lookup
+    add_index :ruby_llm_agents_executions,
+              [:tenant_record_type, :tenant_record_id],
+              name: "index_executions_on_tenant_record"
   end
 end

data/lib/generators/ruby_llm_agents/templates/create_tenant_budgets_migration.rb.tt

@@ -36,10 +36,21 @@ class CreateRubyLLMAgentsTenantBudgets < ActiveRecord::Migration<%= migration_ve
       # Whether to inherit from global config for unset limits
       t.boolean :inherit_global_defaults, default: true
 
+      # Polymorphic association to tenant model (e.g., Organization, Account)
+      # Allows direct association with ActiveRecord models using llm_tenant DSL
+      # Uses string type for tenant_record_id to support both integer and UUID primary keys
+      t.string :tenant_record_type
+      t.string :tenant_record_id
+
       t.timestamps
     end
 
     # Ensure unique tenant IDs
     add_index :ruby_llm_agents_tenant_budgets, :tenant_id, unique: true
+
+    # Index for polymorphic tenant record lookup
+    add_index :ruby_llm_agents_tenant_budgets,
+              [:tenant_record_type, :tenant_record_id],
+              name: "index_tenant_budgets_on_tenant_record"
   end
 end

data/lib/generators/ruby_llm_agents/templates/create_tenants_migration.rb.tt

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+# Migration to create the tenants table for multi-tenancy support
+#
+# This table stores tenant configuration including:
+# - Budget limits (cost, tokens, executions)
+# - Enforcement mode (none, soft, hard)
+# - Polymorphic link to user's tenant model
+#
+# Run with: rails db:migrate
+class CreateRubyLLMAgentsTenants < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :ruby_llm_agents_tenants do |t|
+      # Identity
+      t.string :tenant_id, null: false
+      t.string :name
+
+      # Polymorphic association to user's tenant model (e.g., Organization, Account)
+      # Uses string type for tenant_record_id to support both integer and UUID primary keys
+      t.string :tenant_record_type
+      t.string :tenant_record_id
+
+      # Cost limits (in USD, 6 decimal precision)
+      t.decimal :daily_limit, precision: 12, scale: 6
+      t.decimal :monthly_limit, precision: 12, scale: 6
+
+      # Token limits
+      t.bigint :daily_token_limit
+      t.bigint :monthly_token_limit
+
+      # Execution/call limits
+      t.bigint :daily_execution_limit
+      t.bigint :monthly_execution_limit
+
+      # Per-agent limits (JSON hash)
+      # Format: { "AgentName" => limit_value }
+      t.json :per_agent_daily, null: false, default: {}
+      t.json :per_agent_monthly, null: false, default: {}
+
+      # Enforcement mode: "none", "soft", or "hard"
+      # - none: no enforcement, only tracking
+      # - soft: log warnings when limits exceeded
+      # - hard: block execution when limits exceeded
+      t.string :enforcement, default: "soft"
+
+      # Whether to inherit from global config for unset limits
+      t.boolean :inherit_global_defaults, default: true
+
+      # Status (for soft-delete/disable)
+      t.boolean :active, default: true
+
+      # Extensible metadata (JSON)
+      t.json :metadata, null: false, default: {}
+
+      t.timestamps
+    end
+
+    # Ensure unique tenant IDs
+    add_index :ruby_llm_agents_tenants, :tenant_id, unique: true
+
+    # Index for name lookups
+    add_index :ruby_llm_agents_tenants, :name
+
+    # Index for active/inactive filtering
+    add_index :ruby_llm_agents_tenants, :active
+
+    # Index for polymorphic tenant record lookup
+    add_index :ruby_llm_agents_tenants,
+              [:tenant_record_type, :tenant_record_id],
+              name: "index_tenants_on_tenant_record"
+  end
+end

data/lib/generators/ruby_llm_agents/templates/rename_tenant_budgets_to_tenants_migration.rb.tt

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+# Migration to rename tenant_budgets table to tenants
+#
+# This migration is for users upgrading from an older version of the gem
+# where the table was called 'ruby_llm_agents_tenant_budgets'.
+#
+# If you're doing a fresh install, use the create_tenants_migration instead.
+#
+# Run with: rails db:migrate
+class RenameTenantBudgetsToTenants < ActiveRecord::Migration<%= migration_version %>
+  def change
+    # Rename the table
+    rename_table :ruby_llm_agents_tenant_budgets, :ruby_llm_agents_tenants
+
+    # Add new columns for Tenant model
+    add_column :ruby_llm_agents_tenants, :active, :boolean, default: true
+    add_column :ruby_llm_agents_tenants, :metadata, :json, null: false, default: {}
+
+    # Add index for active status
+    add_index :ruby_llm_agents_tenants, :active
+
+    # Rename indexes to match new table name
+    # Note: Some databases handle this automatically when renaming tables
+    # If you get an error about missing indexes, you may need to adjust this
+
+    # The polymorphic index needs to be renamed if it exists
+    if index_exists?(:ruby_llm_agents_tenants, [:tenant_record_type, :tenant_record_id], name: "index_tenant_budgets_on_tenant_record")
+      rename_index :ruby_llm_agents_tenants,
+                   "index_tenant_budgets_on_tenant_record",
+                   "index_tenants_on_tenant_record"
+    end
+  end
+end

data/lib/generators/ruby_llm_agents/upgrade_generator.rb

@@ -147,6 +147,26 @@ module RubyLlmAgents
       )
     end
 
+    def create_rename_tenant_budgets_migration
+      # Skip if already using new table name
+      if table_exists?(:ruby_llm_agents_tenants)
+        say_status :skip, "ruby_llm_agents_tenants table already exists", :yellow
+        return
+      end
+
+      # Only run if old table exists (needs upgrade)
+      unless table_exists?(:ruby_llm_agents_tenant_budgets)
+        say_status :skip, "No tenant_budgets table to upgrade", :yellow
+        return
+      end
+
+      say_status :upgrade, "Renaming tenant_budgets to tenants", :blue
+      migration_template(
+        "rename_tenant_budgets_to_tenants_migration.rb.tt",
+        File.join(db_migrate_path, "rename_tenant_budgets_to_tenants.rb")
+      )
+    end
+
     def migrate_agents_directory
       root_dir = RubyLLM::Agents.configuration.root_directory
       namespace = RubyLLM::Agents.configuration.root_namespace
@@ -211,6 +231,12 @@ module RubyLlmAgents
       false
     end
 
+    def table_exists?(table)
+      ActiveRecord::Base.connection.table_exists?(table)
+    rescue StandardError
+      false
+    end
+
     def migrate_directory(old_dir, new_dir, namespace)
       source = Rails.root.join("app", old_dir)
       destination = Rails.root.join("app", new_dir)