ruby_llm-agents 3.5.4 → 3.6.0

data/lib/ruby_llm/agents/dsl/queryable.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    module DSL
+      # Adds execution querying capabilities to agent classes.
+      #
+      # Mixed into BaseAgent via `extend DSL::Queryable`, making all methods
+      # available as class methods on agent classes.
+      #
+      # @example Basic queries
+      #   SupportAgent.executions.successful.recent
+      #   SupportAgent.executions.today.expensive(0.50)
+      #
+      # @example Convenience methods
+      #   SupportAgent.last_run
+      #   SupportAgent.stats
+      #   SupportAgent.total_spent(since: 1.week)
+      #
+      module Queryable
+        # Returns an ActiveRecord::Relation scoped to this agent's executions.
+        #
+        # @return [ActiveRecord::Relation]
+        #
+        # @example
+        #   SupportAgent.executions.successful.last(5)
+        #   SupportAgent.executions.where("total_cost > ?", 0.01)
+        #
+        def executions
+          RubyLLM::Agents::Execution.by_agent(name)
+        end
+
+        # Returns the most recent execution for this agent.
+        #
+        # @return [RubyLLM::Agents::Execution, nil]
+        #
+        def last_run
+          executions.order(created_at: :desc).first
+        end
+
+        # Returns recent failed executions.
+        #
+        # @param since [ActiveSupport::Duration] Time window (default: 24.hours)
+        # @return [ActiveRecord::Relation]
+        #
+        def failures(since: 24.hours)
+          executions.failed.where("created_at > ?", since.ago)
+        end
+
+        # Returns total cost spent by this agent.
+        #
+        # @param since [ActiveSupport::Duration, nil] Optional time window
+        # @return [BigDecimal] Total cost in USD
+        #
+        def total_spent(since: nil)
+          scope = executions
+          scope = scope.where("created_at > ?", since.ago) if since
+          scope.sum(:total_cost)
+        end
+
+        # Returns a stats summary hash for this agent.
+        #
+        # @param since [ActiveSupport::Duration, nil] Time window
+        # @return [Hash] Stats summary
+        #
+        def stats(since: nil)
+          scope = executions
+          scope = scope.where("created_at > ?", since.ago) if since
+
+          total = scope.count
+          successful = scope.successful.count
+
+          {
+            total: total,
+            successful: successful,
+            failed: scope.failed.count,
+            success_rate: total.zero? ? 0.0 : (successful.to_f / total * 100).round(1),
+            avg_duration_ms: scope.average(:duration_ms)&.round,
+            avg_cost: total.zero? ? 0 : (scope.sum(:total_cost).to_f / total).round(6),
+            total_cost: scope.sum(:total_cost),
+            total_tokens: scope.sum(:total_tokens),
+            avg_tokens: scope.average(:total_tokens)&.round
+          }
+        end
+
+        # Returns cost breakdown by model for this agent.
+        #
+        # @param since [ActiveSupport::Duration, nil] Time window
+        # @return [Hash{String => Hash}] Costs per model
+        #
+        def cost_by_model(since: nil)
+          scope = executions
+          scope = scope.where("created_at > ?", since.ago) if since
+
+          scope.group(:model_id).pluck(
+            :model_id,
+            Arel.sql("COUNT(*)"),
+            Arel.sql("SUM(total_cost)"),
+            Arel.sql("AVG(total_cost)")
+          ).each_with_object({}) do |(model, count, total, avg), hash|
+            hash[model] = {
+              count: count,
+              total_cost: total&.to_f&.round(6) || 0,
+              avg_cost: avg&.to_f&.round(6) || 0
+            }
+          end
+        end
+
+        # Returns executions matching specific parameter values.
+        #
+        # @param params [Hash] Parameter key-value pairs to match
+        # @return [ActiveRecord::Relation]
+        #
+        def with_params(**params)
+          scope = executions
+          params.each do |key, value|
+            scope = scope.with_parameter(key, value)
+          end
+          scope
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/dsl.rb

@@ -3,6 +3,7 @@
 require_relative "dsl/base"
 require_relative "dsl/reliability"
 require_relative "dsl/caching"
+require_relative "dsl/queryable"
 
 module RubyLLM
   module Agents

data/lib/ruby_llm/agents/image/concerns/image_operation_execution.rb

@@ -85,7 +85,8 @@ module RubyLLM
           if config.async_logging && defined?(ExecutionLoggerJob)
             ExecutionLoggerJob.perform_later(execution_data)
           else
-            RubyLLM::Agents::Execution.create!(execution_data)
+            execution = RubyLLM::Agents::Execution.create!(execution_data)
+            result.execution_id = execution.id if result.respond_to?(:execution_id=)
           end
         rescue => e
           Rails.logger.error("[RubyLLM::Agents] Failed to record #{execution_type} execution: #{e.message}") if defined?(Rails)

data/lib/ruby_llm/agents/image/generator/pricing.rb

@@ -1,17 +1,20 @@
 # frozen_string_literal: true
 
-require "net/http"
-require "json"
+require_relative "../../pricing/data_store"
+require_relative "../../pricing/ruby_llm_adapter"
+require_relative "../../pricing/litellm_adapter"
 
 module RubyLLM
   module Agents
     class ImageGenerator
-      # Dynamic pricing resolution for image generation models
+      # Dynamic pricing resolution for image generation models.
       #
-      # Uses a three-tier strategy:
-      # 1. LiteLLM JSON (primary) - comprehensive, community-maintained
-      # 2. Configurable pricing table - user overrides
-      # 3. Hardcoded fallbacks - last resort
+      # Uses a three-tier strategy (no hardcoded prices):
+      # 1. Configurable pricing table - user overrides
+      # 2. RubyLLM gem (local, no HTTP) - model registry pricing
+      # 3. LiteLLM (via shared DataStore) - comprehensive, community-maintained
+      #
+      # When no pricing is found, returns 0 to signal unknown cost.
       #
       # @example Get price for a model
       #   Pricing.cost_per_image("gpt-image-1", size: "1024x1024", quality: "hd")
@@ -24,9 +27,6 @@ module RubyLLM
       module Pricing
         extend self
 
-        LITELLM_PRICING_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
-        DEFAULT_CACHE_TTL = 24 * 60 * 60 # 24 hours in seconds
-
         # Calculate total cost for image generation
         #
         # @param model_id [String] The model identifier
@@ -44,170 +44,50 @@ module RubyLLM
         # @param model_id [String] The model identifier
         # @param size [String] Image size
         # @param quality [String] Quality setting
-        # @return [Float] Cost per image in USD
+        # @return [Float] Cost per image in USD (0 if unknown)
         def cost_per_image(model_id, size: nil, quality: nil)
-          # 1. Try LiteLLM pricing data
-          if (litellm_price = from_litellm(model_id, size, quality))
-            return litellm_price
-          end
-
-          # 2. Try configurable pricing table
+          # Tier 1: User-configurable pricing table
           if (config_price = from_config(model_id, size, quality))
             return config_price
           end
 
-          # 3. Fall back to hardcoded estimates
-          fallback_price(model_id, size, quality)
+          # Tier 2: RubyLLM gem (local, no HTTP)
+          if (ruby_llm_price = from_ruby_llm(model_id))
+            return ruby_llm_price
+          end
+
+          # Tier 3: LiteLLM (via shared DataStore + adapter)
+          if (litellm_price = from_litellm(model_id, size, quality))
+            return litellm_price
+          end
+
+          # No pricing found — return user-configured default or 0
+          config.default_image_cost || 0
         end
 
-        # Refresh pricing data from LiteLLM
+        # Refresh pricing data
         #
-        # @return [Hash] The fetched pricing data
+        # @return [void]
         def refresh!
-          @litellm_data = nil
-          @litellm_fetched_at = nil
-          litellm_data
+          Agents::Pricing::DataStore.refresh!
         end
 
-        # Get all known pricing for debugging/display
+        # Expose all known pricing for debugging/console inspection
         #
         # @return [Hash] Merged pricing from all sources
         def all_pricing
           {
             litellm: litellm_image_models,
-            configured: config.image_model_pricing || {},
-            fallbacks: fallback_pricing_table
+            configured: config.image_model_pricing || {}
           }
         end
 
         private
 
-        # Fetch from LiteLLM JSON
-        def from_litellm(model_id, size, quality)
-          data = litellm_data
-          return nil unless data
-
-          # Try exact match first
-          model_data = find_litellm_model(data, model_id, size, quality)
-          return nil unless model_data
-
-          extract_litellm_price(model_data, size, quality)
-        end
-
-        def find_litellm_model(data, model_id, size, quality)
-          normalized = normalize_model_id(model_id)
-
-          # Try various key formats LiteLLM uses
-          candidates = [
-            model_id,
-            normalized,
-            "#{size}/#{model_id}",
-            "#{size}/#{quality}/#{model_id}",
-            "aiml/#{normalized}",
-            "together_ai/#{normalized}"
-          ]
-
-          candidates.each do |key|
-            return data[key] if data[key]
-          end
-
-          # Fuzzy match by model name pattern
-          data.find do |key, _value|
-            key_lower = key.to_s.downcase
-            normalized_lower = normalized.downcase
-
-            key_lower.include?(normalized_lower) ||
-              normalized_lower.include?(key_lower.split("/").last.to_s)
-          end&.last
-        end
-
-        def extract_litellm_price(model_data, size, quality)
-          # LiteLLM uses different pricing fields for images
-          if model_data["input_cost_per_image"]
-            return model_data["input_cost_per_image"]
-          end
-
-          if model_data["input_cost_per_pixel"] && size
-            width, height = size.split("x").map(&:to_i)
-            pixels = width * height
-            return (model_data["input_cost_per_pixel"] * pixels).round(6)
-          end
-
-          # Some models have quality-based pricing
-          if quality == "hd" && model_data["input_cost_per_image_hd"]
-            return model_data["input_cost_per_image_hd"]
-          end
-
-          nil
-        end
-
-        def litellm_data
-          return @litellm_data if @litellm_data && !cache_expired?
-
-          @litellm_data = fetch_litellm_data
-          @litellm_fetched_at = Time.now
-          @litellm_data
-        end
-
-        def fetch_litellm_data
-          # Use Rails cache if available
-          if defined?(Rails) && Rails.cache
-            Rails.cache.fetch("litellm_pricing_data", expires_in: cache_ttl) do
-              fetch_from_url
-            end
-          else
-            fetch_from_url
-          end
-        rescue => e
-          warn "[RubyLLM::Agents] Failed to fetch LiteLLM pricing: #{e.message}"
-          {}
-        end
-
-        def fetch_from_url
-          uri = URI(config.litellm_pricing_url || LITELLM_PRICING_URL)
-          http = Net::HTTP.new(uri.host, uri.port)
-          http.use_ssl = uri.scheme == "https"
-          http.open_timeout = 5
-          http.read_timeout = 10
-
-          request = Net::HTTP::Get.new(uri)
-          response = http.request(request)
+        # ============================================================
+        # Tier 1: User-configurable pricing table
+        # ============================================================
 
-          if response.is_a?(Net::HTTPSuccess)
-            JSON.parse(response.body)
-          else
-            {}
-          end
-        rescue => e
-          warn "[RubyLLM::Agents] HTTP error fetching LiteLLM pricing: #{e.message}"
-          {}
-        end
-
-        def cache_expired?
-          return true unless @litellm_fetched_at
-          Time.now - @litellm_fetched_at > cache_ttl
-        end
-
-        def cache_ttl
-          ttl = config.litellm_pricing_cache_ttl
-          return DEFAULT_CACHE_TTL unless ttl
-
-          # Handle ActiveSupport::Duration
-          ttl.respond_to?(:to_i) ? ttl.to_i : ttl
-        end
-
-        # Get image-specific models from LiteLLM data
-        def litellm_image_models
-          litellm_data.select do |key, value|
-            value.is_a?(Hash) && (
-              value["input_cost_per_image"] ||
-              value["input_cost_per_pixel"] ||
-              key.to_s.match?(/dall-e|flux|sdxl|stable|imagen|image/i)
-            )
-          end
-        end
-
-        # Fetch from configurable pricing table
         def from_config(model_id, size, quality)
           table = config.image_model_pricing
           return nil unless table
@@ -251,80 +131,58 @@ module RubyLLM
           pricing[:base] || pricing["base"] || pricing[:default] || pricing["default"] || pricing[:standard] || pricing["standard"]
         end
 
-        # Hardcoded fallback prices
-        def fallback_price(model_id, size, quality)
-          normalized = normalize_model_id(model_id)
+        # ============================================================
+        # Tier 2: RubyLLM gem (local, no HTTP)
+        # ============================================================
 
-          case normalized
-          when /gpt-image-1|dall-e-3/i
-            dalle3_price(size, quality)
-          when /dall-e-2/i
-            dalle2_price(size)
-          when /imagen/i
-            0.02
-          when /flux.*pro.*ultra/i
-            0.063
-          when /flux.*pro/i
-            0.05
-          when /flux.*dev/i
-            0.025
-          when /flux.*schnell/i
-            0.003
-          when /sdxl.*lightning/i
-            0.002
-          when /sdxl|stable-diffusion-xl/i
-            0.04
-          when /stable-diffusion/i
-            0.02
-          when /ideogram/i
-            0.04
-          when /recraft/i
-            0.04
-          when /real-esrgan|upscal/i
-            0.01
-          when /blip|caption|analyz/i
-            0.001
-          when /segment|background|rembg/i
-            0.01
-          else
-            config.default_image_cost || 0.04
-          end
+        def from_ruby_llm(model_id)
+          data = Agents::Pricing::RubyLLMAdapter.find_model(model_id)
+          return nil unless data
+
+          data[:input_cost_per_image]
         end
 
-        def dalle3_price(size, quality)
-          pixels = parse_pixels(size)
-          is_large = pixels && pixels >= 1_000_000
+        # ============================================================
+        # Tier 3: LiteLLM (via shared DataStore + adapter)
+        # ============================================================
 
-          case quality
-          when "hd"
-            is_large ? 0.12 : 0.08
-          else
-            is_large ? 0.08 : 0.04
-          end
+        def from_litellm(model_id, size, quality)
+          data = Agents::Pricing::LiteLLMAdapter.find_model(model_id)
+          return nil unless data
+
+          extract_image_price(data, size, quality)
         end
 
-        def dalle2_price(size)
-          case size
-          when "1024x1024" then 0.02
-          when "512x512" then 0.018
-          when "256x256" then 0.016
-          else 0.02
+        def extract_image_price(data, size, quality)
+          # Check quality-specific pricing first when HD requested
+          if quality == "hd" && data[:input_cost_per_image_hd]
+            return data[:input_cost_per_image_hd]
           end
+
+          if data[:input_cost_per_image]
+            return data[:input_cost_per_image]
+          end
+
+          if data[:input_cost_per_pixel] && size
+            width, height = size.split("x").map(&:to_i)
+            pixels = width * height
+            return (data[:input_cost_per_pixel] * pixels).round(6)
+          end
+
+          nil
         end
 
-        def fallback_pricing_table
-          {
-            "gpt-image-1" => {standard: 0.04, hd: 0.08, large_hd: 0.12},
-            "dall-e-3" => {standard: 0.04, hd: 0.08, large_hd: 0.12},
-            "dall-e-2" => {"1024x1024" => 0.02, "512x512" => 0.018, "256x256" => 0.016},
-            "flux-pro" => 0.05,
-            "flux-dev" => 0.025,
-            "flux-schnell" => 0.003,
-            "sdxl" => 0.04,
-            "stable-diffusion-3.5" => 0.03,
-            "imagen-3" => 0.02,
-            "ideogram-2" => 0.04
-          }
+        def litellm_image_models
+          data = Agents::Pricing::DataStore.litellm_data
+          return {} unless data.is_a?(Hash)
+
+          data.select do |key, value|
+            value.is_a?(Hash) && (
+              value["input_cost_per_image"] ||
+              value["input_cost_per_pixel"] ||
+              key.to_s.match?(/dall-e|flux|sdxl|stable|imagen|image/i)
+            )
+          end
         end
 
         def parse_pixels(size)

data/lib/ruby_llm/agents/image/generator.rb

@@ -248,7 +248,8 @@ module RubyLLM
           images: images,
           started_at: context.started_at || execution_started_at,
           completed_at: execution_completed_at,
-          tenant_id: context.tenant_id
+          tenant_id: context.tenant_id,
+          execution_id: context.execution_id
         )
 
         # Update context with cost info
@@ -348,7 +349,7 @@ module RubyLLM
       end
 
       # Build successful result
-      def build_result(images:, started_at:, completed_at:, tenant_id:)
+      def build_result(images:, started_at:, completed_at:, tenant_id:, execution_id: nil)
         ImageGenerationResult.new(
           images: images,
           prompt: prompt,
@@ -359,7 +360,8 @@ module RubyLLM
           started_at: started_at,
           completed_at: completed_at,
           tenant_id: tenant_id,
-          generator_class: self.class.name
+          generator_class: self.class.name,
+          execution_id: execution_id
         )
       end
 

data/lib/ruby_llm/agents/infrastructure/attempt_tracker.rb

@@ -121,6 +121,8 @@ module RubyLLM
 
       # Finds the last failed attempt
       #
+      # Useful for debugging in rails console.
+      #
       # @return [Hash, nil] The last failed attempt or nil
       def last_failed_attempt
         @attempts.reverse.find { |a| a[:error_class].present? }
@@ -189,6 +191,8 @@ module RubyLLM
 
       # Calculates total duration across all attempts
       #
+      # Useful for debugging in rails console.
+      #
       # @return [Integer] Total duration in milliseconds
       def total_duration_ms
         @attempts.sum { |a| a[:duration_ms] || 0 }
@@ -203,6 +207,8 @@ module RubyLLM
 
       # Returns the number of failed attempts
       #
+      # Useful for debugging in rails console.
+      #
       # @return [Integer] Number of failed attempts
       def failed_attempts_count
         @attempts.count { |a| a[:error_class].present? }
@@ -210,6 +216,8 @@ module RubyLLM
 
       # Returns the number of short-circuited attempts
       #
+      # Useful for debugging in rails console.
+      #
       # @return [Integer] Number of short-circuited attempts
       def short_circuited_count
         @attempts.count { |a| a[:short_circuited] }

data/lib/ruby_llm/agents/infrastructure/circuit_breaker.rb

@@ -123,17 +123,19 @@ module RubyLLM
 
       # Returns the time remaining until the breaker closes
       #
+      # Useful for debugging circuit breaker state in rails console.
+      #
       # @return [Integer, nil] Seconds until cooldown expires, or nil if not open
       def time_until_close
         return nil unless open?
 
-        # We can't easily get TTL from Rails.cache, so this is an approximation
-        # In a real implementation, you might store the open time as well
         cooldown_seconds
       end
 
       # Returns status information for the circuit breaker
       #
+      # Useful for debugging circuit breaker state in rails console.
+      #
       # @return [Hash] Status information including open state and failure count
       def status
         {

data/lib/ruby_llm/agents/pipeline/builder.rb

@@ -144,6 +144,12 @@ module RubyLLM
 
               # Reliability (if agent has retries or fallbacks configured)
               builder.use(Middleware::Reliability) if reliability_enabled?(agent_class)
+
+              # Global custom middleware
+              apply_custom_middleware!(builder, global_middleware_entries)
+
+              # Per-agent custom middleware
+              apply_custom_middleware!(builder, agent_middleware_entries(agent_class))
             end
           end
 
@@ -159,6 +165,43 @@ module RubyLLM
 
           private
 
+          # Returns global custom middleware entries from configuration
+          #
+          # @return [Array<Hash>] Middleware entries
+          def global_middleware_entries
+            RubyLLM::Agents.configuration.middleware_stack
+          rescue
+            []
+          end
+
+          # Returns per-agent custom middleware entries
+          #
+          # @param agent_class [Class] The agent class
+          # @return [Array<Hash>] Middleware entries
+          def agent_middleware_entries(agent_class)
+            return [] unless agent_class&.respond_to?(:agent_middleware)
+
+            agent_class.agent_middleware
+          rescue
+            []
+          end
+
+          # Applies custom middleware entries to a builder
+          #
+          # @param builder [Builder] The builder to modify
+          # @param entries [Array<Hash>] Middleware entries with :klass, :before, :after
+          def apply_custom_middleware!(builder, entries)
+            entries.each do |entry|
+              if entry[:before]
+                builder.insert_before(entry[:before], entry[:klass])
+              elsif entry[:after]
+                builder.insert_after(entry[:after], entry[:klass])
+              else
+                builder.use(entry[:klass])
+              end
+            end
+          end
+
           # Check if budgets are enabled globally
           #
           # @return [Boolean]

data/lib/ruby_llm/agents/pipeline/context.rb

@@ -34,6 +34,9 @@ module RubyLLM
         # Execution tracking (set by Instrumentation middleware)
         attr_accessor :started_at, :completed_at, :attempt, :attempts_made, :execution_id
 
+        # Execution hierarchy (agent-as-tool)
+        attr_accessor :parent_execution_id, :root_execution_id
+
         # Result data (set by core execute method)
         attr_accessor :output, :error, :cached
 
@@ -59,13 +62,17 @@ module RubyLLM
         # @param skip_cache [Boolean] Whether to skip caching
         # @param stream_block [Proc, nil] Block for streaming
         # @param options [Hash] Additional options passed to the agent
-        def initialize(input:, agent_class:, agent_instance: nil, model: nil, tenant: nil, skip_cache: false, stream_block: nil, **options)
+        def initialize(input:, agent_class:, agent_instance: nil, model: nil, tenant: nil, skip_cache: false, stream_block: nil, parent_execution_id: nil, root_execution_id: nil, **options)
           @input = input
           @agent_class = agent_class
           @agent_instance = agent_instance
           @agent_type = extract_agent_type(agent_class)
           @model = model || extract_model(agent_class)
 
+          # Execution hierarchy
+          @parent_execution_id = parent_execution_id
+          @root_execution_id = root_execution_id
+
           # Store tenant in options for middleware to resolve
           @options = options.merge(tenant: tenant).compact
 
@@ -202,6 +209,9 @@ module RubyLLM
           new_ctx.tenant_config = @tenant_config
           new_ctx.started_at = @started_at
           new_ctx.attempts_made = @attempts_made
+          # Preserve execution hierarchy
+          new_ctx.parent_execution_id = @parent_execution_id
+          new_ctx.root_execution_id = @root_execution_id
           new_ctx
         end