ruby_llm-agents 0.4.0 → 1.0.0.beta.1

data/lib/ruby_llm/agents/core/errors.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # Base error class for RubyLLM::Agents
+    class Error < StandardError; end
+
+    # ============================================================
+    # Pipeline Errors
+    # ============================================================
+
+    # Base class for pipeline-related errors
+    class PipelineError < Error; end
+
+    # ============================================================
+    # Reliability Errors
+    # ============================================================
+
+    # Base class for reliability-related errors
+    class ReliabilityError < Error; end
+
+    # Raised when an error is retryable (transient)
+    class RetryableError < ReliabilityError; end
+
+    # Raised when a circuit breaker is open
+    class CircuitOpenError < ReliabilityError
+      # @return [String] The model that has an open circuit
+      attr_reader :model
+
+      def initialize(message = nil, model: nil)
+        @model = model
+        super(message || "Circuit breaker is open#{model ? " for #{model}" : ""}")
+      end
+    end
+
+    # Raised when total timeout is exceeded across all attempts
+    class TotalTimeoutError < ReliabilityError
+      # @return [Float] The timeout that was exceeded
+      attr_reader :timeout
+
+      # @return [Float] The elapsed time
+      attr_reader :elapsed
+
+      def initialize(message = nil, timeout: nil, elapsed: nil)
+        @timeout = timeout
+        @elapsed = elapsed
+        super(message || "Total timeout of #{timeout}s exceeded (elapsed: #{elapsed&.round(2)}s)")
+      end
+    end
+
+    # Raised when all models (primary + fallbacks) fail
+    class AllModelsFailedError < ReliabilityError
+      # @return [Array<Hash>] Details of each failed attempt
+      attr_reader :attempts
+
+      def initialize(message = nil, attempts: [])
+        @attempts = attempts
+        models = attempts.map { |a| a[:model] }.compact.join(", ")
+        super(message || "All models failed: #{models}")
+      end
+    end
+
+    # ============================================================
+    # Budget Errors
+    # ============================================================
+
+    # Base class for budget-related errors
+    class BudgetError < Error; end
+
+    # Raised when budget is exceeded
+    class BudgetExceededError < BudgetError
+      # @return [String, nil] The tenant ID
+      attr_reader :tenant_id
+
+      # @return [String, nil] The budget type (daily, monthly, etc.)
+      attr_reader :budget_type
+
+      def initialize(message = nil, tenant_id: nil, budget_type: nil)
+        @tenant_id = tenant_id
+        @budget_type = budget_type
+        super(message || "Budget exceeded#{tenant_id ? " for tenant #{tenant_id}" : ""}")
+      end
+    end
+
+    # ============================================================
+    # Configuration Errors
+    # ============================================================
+
+    # Raised for configuration issues
+    class ConfigurationError < Error; end
+
+    # Raised when content is flagged during moderation
+    #
+    # Contains the full moderation result and the phase where
+    # the content was flagged.
+    #
+    # @example Handling moderation errors
+    #   begin
+    #     result = MyAgent.call(message: user_input)
+    #   rescue RubyLLM::Agents::ModerationError => e
+    #     puts "Content blocked: #{e.flagged_categories.join(', ')}"
+    #     puts "Phase: #{e.phase}"
+    #     puts "Scores: #{e.category_scores}"
+    #   end
+    #
+    # @api public
+    class ModerationError < Error
+      # @return [Object] The raw moderation result from RubyLLM
+      attr_reader :moderation_result
+
+      # @return [Symbol] The phase where content was flagged (:input or :output)
+      attr_reader :phase
+
+      # Creates a new ModerationError
+      #
+      # @param moderation_result [Object] The moderation result from RubyLLM
+      # @param phase [Symbol] The phase where content was flagged
+      def initialize(moderation_result, phase)
+        @moderation_result = moderation_result
+        @phase = phase
+
+        categories = moderation_result.flagged_categories
+        category_list = categories.respond_to?(:join) ? categories.join(", ") : categories.to_s
+
+        super("Content flagged during #{phase} moderation: #{category_list}")
+      end
+
+      # Returns the flagged categories from the moderation result
+      #
+      # @return [Array<String, Symbol>] List of flagged categories
+      def flagged_categories
+        moderation_result.flagged_categories
+      end
+
+      # Returns the category scores from the moderation result
+      #
+      # @return [Hash{String, Symbol => Float}] Category to score mapping
+      def category_scores
+        moderation_result.category_scores
+      end
+
+      # Returns whether the moderation result was flagged
+      #
+      # @return [Boolean] Always true for ModerationError
+      def flagged?
+        true
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/{instrumentation.rb → core/instrumentation.rb}

@@ -245,7 +245,9 @@ module RubyLLM
           metadata: metadata,
           system_prompt: config.persist_prompts ? redacted_system_prompt : nil,
           user_prompt: config.persist_prompts ? redacted_user_prompt : nil,
-          streaming: self.class.streaming
+          streaming: self.class.streaming,
+          messages_count: resolved_messages.size,
+          messages_summary: config.persist_messages_summary ? messages_summary : {}
         }
 
         # Extract tracing fields from metadata if present
@@ -326,6 +328,9 @@ module RubyLLM
             Rails.logger.warn("[RubyLLM::Agents] Cost calculation failed: #{cost_error.message}")
           end
         end
+
+        # Record token usage for budget tracking
+        record_token_usage(execution)
       rescue ActiveRecord::RecordInvalid => e
         Rails.logger.error("[RubyLLM::Agents] Validation failed for execution #{execution&.id}: #{e.record.errors.full_messages.join(', ')}")
         if Rails.env.development? || Rails.env.test?
@@ -415,6 +420,9 @@ module RubyLLM
             Rails.logger.warn("[RubyLLM::Agents] Cost calculation failed: #{cost_error.message}")
           end
         end
+
+        # Record token usage for budget tracking
+        record_token_usage(execution)
       rescue ActiveRecord::RecordInvalid => e
         Rails.logger.error("[RubyLLM::Agents] Validation failed for execution #{execution&.id}: #{e.record.errors.full_messages.join(', ')}")
         if Rails.env.development? || Rails.env.test?
@@ -440,6 +448,8 @@ module RubyLLM
       # @param error [Exception, nil] The exception if failed
       # @return [void]
       def legacy_log_execution(completed_at:, status:, response: nil, error: nil)
+        config = RubyLLM::Agents.configuration
+
         execution_data = {
           agent_type: self.class.name,
           agent_version: self.class.version,
@@ -452,7 +462,9 @@ module RubyLLM
           parameters: sanitized_parameters,
           metadata: execution_metadata,
           system_prompt: safe_system_prompt,
-          user_prompt: safe_user_prompt
+          user_prompt: safe_user_prompt,
+          messages_count: resolved_messages.size,
+          messages_summary: config.persist_messages_summary ? messages_summary : {}
         }
 
         # Add response data if available (using safe extraction)
@@ -516,6 +528,38 @@ module RubyLLM
         Redactor.redact_string(prompt)
       end
 
+      # Returns a summary of messages (first and last, truncated)
+      #
+      # Creates a summary of the conversation messages containing the first
+      # and last messages (if different) with content truncated for storage.
+      #
+      # @return [Hash] Summary with :first and :last message hashes, or empty hash
+      def messages_summary
+        msgs = resolved_messages
+        return {} if msgs.blank?
+
+        max_len = RubyLLM::Agents.configuration.messages_summary_max_length || 500
+
+        summary = {}
+
+        if msgs.first
+          summary[:first] = {
+            role: msgs.first[:role].to_s,
+            content: Redactor.redact_string(msgs.first[:content].to_s).truncate(max_len)
+          }
+        end
+
+        # Only add last if there are multiple messages and last is different from first
+        if msgs.size > 1 && msgs.last
+          summary[:last] = {
+            role: msgs.last[:role].to_s,
+            content: Redactor.redact_string(msgs.last[:content].to_s).truncate(max_len)
+          }
+        end
+
+        summary
+      end
+
       # Returns the response with redaction applied
       #
       # @param response [RubyLLM::Message] The LLM response
@@ -583,6 +627,9 @@ module RubyLLM
         # during multi-turn conversations (when tools are used)
         tool_calls_data = respond_to?(:accumulated_tool_calls) ? accumulated_tool_calls : []
 
+        # Extract thinking data if present
+        thinking_data = safe_extract_thinking_data(response)
+
         {
           input_tokens: safe_response_value(response, :input_tokens),
           output_tokens: safe_response_value(response, :output_tokens),
@@ -593,7 +640,7 @@ module RubyLLM
           response: safe_serialize_response(response),
           tool_calls: tool_calls_data || [],
           tool_calls_count: tool_calls_data&.size || 0
-        }.compact
+        }.merge(thinking_data).compact
       end
 
       # Extracts finish reason from response, normalizing to standard values
@@ -621,6 +668,24 @@ module RubyLLM
         end
       end
 
+      # Extracts thinking data from response
+      #
+      # Handles different response structures from various providers.
+      # The thinking object typically has text, signature, and tokens.
+      #
+      # @param response [RubyLLM::Message] The LLM response
+      # @return [Hash] Thinking data (empty if none present)
+      def safe_extract_thinking_data(response)
+        thinking = safe_response_value(response, :thinking)
+        return {} unless thinking
+
+        {
+          thinking_text: thinking.respond_to?(:text) ? thinking.text : thinking[:text],
+          thinking_signature: thinking.respond_to?(:signature) ? thinking.signature : thinking[:signature],
+          thinking_tokens: thinking.respond_to?(:tokens) ? thinking.tokens : thinking[:tokens]
+        }.compact
+      end
+
       # Extracts routing/retry tracking data from attempt tracker
       #
       # Analyzes the execution attempts to determine:
@@ -773,7 +838,9 @@ module RubyLLM
           total_cost: 0,
           parameters: redacted_parameters,
           metadata: execution_metadata,
-          streaming: self.class.streaming
+          streaming: self.class.streaming,
+          messages_count: resolved_messages.size,
+          messages_summary: config.persist_messages_summary ? messages_summary : {}
         }
 
         # Add tracing fields from metadata if present
@@ -798,6 +865,28 @@ module RubyLLM
         Rails.logger.error("[RubyLLM::Agents] Failed to record cache hit execution: #{e.message}")
       end
 
+      # Records token usage to the BudgetTracker
+      #
+      # @param execution [Execution] The completed execution record
+      # @return [void]
+      def record_token_usage(execution)
+        return unless execution&.total_tokens && execution.total_tokens > 0
+
+        begin
+          tenant_id = respond_to?(:resolved_tenant_id) ? resolved_tenant_id : nil
+          tenant_config = respond_to?(:runtime_tenant_config) ? runtime_tenant_config : nil
+
+          BudgetTracker.record_tokens!(
+            self.class.name,
+            execution.total_tokens,
+            tenant_id: tenant_id,
+            tenant_config: tenant_config
+          )
+        rescue StandardError => e
+          Rails.logger.warn("[RubyLLM::Agents] Failed to record token usage: #{e.message}")
+        end
+      end
+
       # Emergency fallback to mark execution as failed
       #
       # Uses update_all to bypass ActiveRecord callbacks and validations,

data/lib/ruby_llm/agents/core/llm_tenant.rb

@@ -0,0 +1,358 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+
+module RubyLLM
+  module Agents
+    # DSL for declaring Rails models as LLM tenants
+    #
+    # Provides automatic budget management and usage tracking when included
+    # in ActiveRecord models. Models using this concern can be passed as
+    # the `tenant:` parameter to agents.
+    #
+    # @example Basic usage
+    #   class Organization < ApplicationRecord
+    #     llm_tenant
+    #   end
+    #
+    # @example With custom ID method
+    #   class Organization < ApplicationRecord
+    #     llm_tenant id: :slug
+    #   end
+    #
+    # @example With auto-created budget
+    #   class Organization < ApplicationRecord
+    #     llm_tenant id: :slug, budget: true
+    #   end
+    #
+    # @example With limits (auto-creates budget)
+    #   class Organization < ApplicationRecord
+    #     llm_tenant(
+    #       id: :slug,
+    #       name: :company_name,
+    #       limits: {
+    #         daily_cost: 100,
+    #         monthly_cost: 1000,
+    #         daily_executions: 500
+    #       },
+    #       enforcement: :hard
+    #     )
+    #   end
+    #
+    # @example With API keys from model columns/methods
+    #   class Organization < ApplicationRecord
+    #     encrypts :openai_api_key, :anthropic_api_key  # Rails 7+ encryption
+    #
+    #     llm_tenant(
+    #       id: :slug,
+    #       api_keys: {
+    #         openai: :openai_api_key,        # column name
+    #         anthropic: :anthropic_api_key,  # column name
+    #         gemini: :fetch_gemini_key       # custom method
+    #       }
+    #     )
+    #
+    #     def fetch_gemini_key
+    #       Vault.read("secret/#{slug}/gemini")
+    #     end
+    #   end
+    #
+    # @see RubyLLM::Agents::TenantBudget
+    # @api public
+    module LLMTenant
+      extend ActiveSupport::Concern
+
+      included do
+        # Executions tracked for this tenant
+        has_many :llm_executions,
+                 class_name: "RubyLLM::Agents::Execution",
+                 as: :tenant_record,
+                 dependent: :nullify
+
+        # Budget association (optional)
+        has_one :llm_budget,
+                class_name: "RubyLLM::Agents::TenantBudget",
+                as: :tenant_record,
+                dependent: :destroy
+
+        # Store options at class level
+        class_attribute :llm_tenant_options, default: {}
+      end
+
+      class_methods do
+        # Declares this model as an LLM tenant
+        #
+        # @param id [Symbol] Method to call for tenant_id string (default: :id)
+        # @param name [Symbol] Method for budget display name (default: :to_s)
+        # @param budget [Boolean] Auto-create TenantBudget on model creation (default: false)
+        # @param limits [Hash] Default budget limits (implies budget: true)
+        # @param enforcement [Symbol] Budget enforcement mode (:none, :soft, :hard)
+        # @param inherit_global [Boolean] Inherit from global config (default: true)
+        # @param api_keys [Hash] Provider API keys mapping (e.g., { openai: :openai_api_key })
+        # @return [void]
+        def llm_tenant(id: :id, name: :to_s, budget: false, limits: nil, enforcement: nil, inherit_global: true, api_keys: nil)
+          self.llm_tenant_options = {
+            id: id,
+            name: name,
+            budget: budget || limits.present?,
+            limits: normalize_limits(limits),
+            enforcement: enforcement,
+            inherit_global: inherit_global,
+            api_keys: api_keys
+          }
+
+          # Auto-create budget callback
+          after_create :create_default_llm_budget if llm_tenant_options[:budget]
+        end
+
+        private
+
+        # Normalizes the limits hash to internal column names
+        #
+        # @param limits [Hash, nil] User-provided limits
+        # @return [Hash] Normalized limits
+        def normalize_limits(limits)
+          return {} if limits.blank?
+
+          {
+            daily_cost: limits[:daily_cost],
+            monthly_cost: limits[:monthly_cost],
+            daily_tokens: limits[:daily_tokens],
+            monthly_tokens: limits[:monthly_tokens],
+            daily_executions: limits[:daily_executions],
+            monthly_executions: limits[:monthly_executions]
+          }.compact
+        end
+      end
+
+      # Returns the tenant_id string for this model
+      #
+      # @return [String] The tenant identifier
+      def llm_tenant_id
+        id_method = self.class.llm_tenant_options[:id] || :id
+        send(id_method).to_s
+      end
+
+      # Returns API keys resolved from the DSL configuration
+      #
+      # Maps provider names (e.g., :openai, :anthropic) to their resolved values
+      # by calling the configured method/column on this model instance.
+      #
+      # @return [Hash] Provider to API key mapping (e.g., { openai: "sk-..." })
+      # @example
+      #   org.llm_api_keys
+      #   # => { openai: "sk-abc123", anthropic: "sk-ant-xyz789" }
+      def llm_api_keys
+        api_keys_config = self.class.llm_tenant_options[:api_keys]
+        return {} if api_keys_config.blank?
+
+        api_keys_config.transform_values do |method_name|
+          value = send(method_name)
+          value.presence
+        end.compact
+      end
+
+      # Returns cost for a given period
+      #
+      # @param period [Symbol, Range, nil] Time period (:today, :this_month, etc.)
+      # @return [BigDecimal] Total cost
+      def llm_cost(period: nil)
+        scope = llm_executions
+        scope = apply_period_scope(scope, period) if period
+        scope.sum(:total_cost) || 0
+      end
+
+      # Returns cost for today
+      #
+      # @return [BigDecimal] Today's cost
+      def llm_cost_today
+        llm_cost(period: :today)
+      end
+
+      # Returns cost for this month
+      #
+      # @return [BigDecimal] This month's cost
+      def llm_cost_this_month
+        llm_cost(period: :this_month)
+      end
+
+      # Returns token count for a given period
+      #
+      # @param period [Symbol, Range, nil] Time period
+      # @return [Integer] Total tokens
+      def llm_tokens(period: nil)
+        scope = llm_executions
+        scope = apply_period_scope(scope, period) if period
+        scope.sum(:total_tokens) || 0
+      end
+
+      # Returns tokens for today
+      #
+      # @return [Integer] Today's tokens
+      def llm_tokens_today
+        llm_tokens(period: :today)
+      end
+
+      # Returns tokens for this month
+      #
+      # @return [Integer] This month's tokens
+      def llm_tokens_this_month
+        llm_tokens(period: :this_month)
+      end
+
+      # Returns execution count for a given period
+      #
+      # @param period [Symbol, Range, nil] Time period
+      # @return [Integer] Execution count
+      def llm_execution_count(period: nil)
+        scope = llm_executions
+        scope = apply_period_scope(scope, period) if period
+        scope.count
+      end
+
+      # Returns executions for today
+      #
+      # @return [Integer] Today's execution count
+      def llm_executions_today
+        llm_execution_count(period: :today)
+      end
+
+      # Returns executions for this month
+      #
+      # @return [Integer] This month's execution count
+      def llm_executions_this_month
+        llm_execution_count(period: :this_month)
+      end
+
+      # Returns a usage summary for a given period
+      #
+      # @param period [Symbol] Time period (default: :this_month)
+      # @return [Hash] Usage summary with cost, tokens, and executions
+      def llm_usage_summary(period: :this_month)
+        {
+          cost: llm_cost(period: period),
+          tokens: llm_tokens(period: period),
+          executions: llm_execution_count(period: period),
+          period: period
+        }
+      end
+
+      # Returns or builds the associated TenantBudget
+      #
+      # @return [TenantBudget] The budget record
+      def llm_budget
+        super || build_llm_budget(tenant_id: llm_tenant_id)
+      end
+
+      # Configure budget with a block
+      #
+      # @yield [budget] The budget to configure
+      # @return [TenantBudget] The saved budget
+      def llm_configure_budget
+        budget = llm_budget
+        yield(budget) if block_given?
+        budget.save!
+        budget
+      end
+
+      # Returns the budget status from BudgetTracker
+      #
+      # @return [Hash] Budget status
+      def llm_budget_status
+        BudgetTracker.status(tenant_id: llm_tenant_id)
+      end
+
+      # Checks if within budget for a given limit type
+      #
+      # @param type [Symbol] Limit type (:daily_cost, :monthly_cost, :daily_tokens, etc.)
+      # @return [Boolean] true if within budget
+      def llm_within_budget?(type: :daily_cost)
+        status = llm_budget_status
+        return true unless status[:enabled]
+
+        key = budget_status_key(type)
+        status.dig(key, :percentage_used).to_f < 100
+      end
+
+      # Returns remaining budget for a given limit type
+      #
+      # @param type [Symbol] Limit type
+      # @return [Numeric, nil] Remaining amount
+      def llm_remaining_budget(type: :daily_cost)
+        status = llm_budget_status
+        key = budget_status_key(type)
+        status.dig(key, :remaining)
+      end
+
+      # Raises an error if over budget
+      #
+      # @raise [BudgetExceededError] if budget is exceeded
+      # @return [void]
+      def llm_check_budget!
+        BudgetTracker.check_budget!(self.class.name, tenant_id: llm_tenant_id)
+      end
+
+      private
+
+      # Applies a period scope to an execution query
+      #
+      # @param scope [ActiveRecord::Relation] The query scope
+      # @param period [Symbol, Range] The period to filter by
+      # @return [ActiveRecord::Relation] Filtered scope
+      def apply_period_scope(scope, period)
+        case period
+        when :today then scope.where(created_at: Time.current.all_day)
+        when :yesterday then scope.where(created_at: 1.day.ago.all_day)
+        when :this_week then scope.where(created_at: Time.current.all_week)
+        when :this_month then scope.where(created_at: Time.current.all_month)
+        when Range then scope.where(created_at: period)
+        else scope
+        end
+      end
+
+      # Maps user-friendly type to budget status key
+      #
+      # @param type [Symbol] User-friendly type
+      # @return [Symbol] Status key
+      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
+      end
+
+      # Creates the default budget on model creation
+      #
+      # @return [void]
+      def create_default_llm_budget
+        return if self.class.llm_tenant_options.blank?
+        return if llm_budget&.persisted?
+
+        options = self.class.llm_tenant_options
+        limits = options[:limits] || {}
+        name_method = options[:name] || :to_s
+
+        budget = build_llm_budget(
+          tenant_id: llm_tenant_id,
+          name: send(name_method).to_s,
+          daily_limit: limits[:daily_cost],
+          monthly_limit: limits[:monthly_cost],
+          daily_token_limit: limits[:daily_tokens],
+          monthly_token_limit: limits[:monthly_tokens],
+          daily_execution_limit: limits[:daily_executions],
+          monthly_execution_limit: limits[:monthly_executions],
+          enforcement: options[:enforcement]&.to_s || "soft",
+          inherit_global_defaults: options.fetch(:inherit_global, true)
+        )
+
+        budget.tenant_record = self
+        budget.save!
+      end
+    end
+  end
+end