ruby_llm-agents 0.4.0 → 1.0.0.beta.1

data/lib/ruby_llm/agents/pipeline/middleware/base.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    module Pipeline
+      module Middleware
+        # Base class for all middleware in the pipeline.
+        #
+        # Middleware wraps the next handler in the chain and can:
+        # - Modify the context before passing it down
+        # - Short-circuit the chain (e.g., return cached result)
+        # - Handle errors from downstream
+        # - Modify the context after the response
+        #
+        # Each middleware receives:
+        # - @app: The next handler in the chain (another middleware or the executor)
+        # - @agent_class: The agent class, for reading DSL configuration
+        #
+        # @example Simple pass-through middleware
+        #   class Logger < Base
+        #     def call(context)
+        #       puts "Before: #{context.input}"
+        #       result = @app.call(context)
+        #       puts "After: #{context.output}"
+        #       result
+        #     end
+        #   end
+        #
+        # @example Short-circuiting middleware
+        #   class Cache < Base
+        #     def call(context)
+        #       if (cached = read_cache(context))
+        #         context.output = cached
+        #         context.cached = true
+        #         return context
+        #       end
+        #       @app.call(context)
+        #     end
+        #   end
+        #
+        # @abstract Subclass and implement {#call}
+        #
+        class Base
+          # @param app [#call] The next handler in the chain
+          # @param agent_class [Class] The agent class (for reading DSL config)
+          def initialize(app, agent_class)
+            @app = app
+            @agent_class = agent_class
+          end
+
+          # Process the context through this middleware
+          #
+          # Subclasses must implement this method. The typical pattern is:
+          # 1. Do pre-processing on context
+          # 2. Call @app.call(context) to continue the chain
+          # 3. Do post-processing on context
+          # 4. Return context
+          #
+          # @param context [Context] The execution context
+          # @return [Context] The (possibly modified) context
+          # @raise [NotImplementedError] If not implemented by subclass
+          def call(context)
+            raise NotImplementedError, "#{self.class} must implement #call"
+          end
+
+          private
+
+          # Read configuration from agent class DSL
+          #
+          # Safely reads a DSL method value from the agent class,
+          # returning a default if the method doesn't exist.
+          #
+          # @param method [Symbol] DSL method name
+          # @param default [Object] Default value if not set
+          # @return [Object] The configuration value
+          def config(method, default = nil)
+            return default unless @agent_class
+
+            if @agent_class.respond_to?(method)
+              @agent_class.send(method)
+            else
+              default
+            end
+          end
+
+          # Check if a DSL option is enabled
+          #
+          # Convenience method for boolean DSL options.
+          #
+          # @param method [Symbol] DSL method name (e.g., :cache_enabled?)
+          # @return [Boolean]
+          def enabled?(method)
+            config(method, false) == true
+          end
+
+          # Returns the global configuration
+          #
+          # @return [Configuration] The RubyLLM::Agents configuration
+          def global_config
+            RubyLLM::Agents.configuration
+          end
+
+          # Log a debug message if Rails logger is available
+          #
+          # @param message [String] The message to log
+          def debug(message)
+            return unless defined?(Rails) && Rails.logger
+
+            Rails.logger.debug("[RubyLLM::Agents::Pipeline] #{message}")
+          end
+
+          # Log an error message if Rails logger is available
+          #
+          # @param message [String] The message to log
+          def error(message)
+            return unless defined?(Rails) && Rails.logger
+
+            Rails.logger.error("[RubyLLM::Agents::Pipeline] #{message}")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/pipeline/middleware/budget.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    module Pipeline
+      module Middleware
+        # Checks budget limits before execution and records spend after.
+        #
+        # This middleware integrates with the BudgetTracker to:
+        # - Check if budget limits are exceeded before execution
+        # - Record spend after successful execution
+        #
+        # Budget checking is skipped if:
+        # - Budgets are disabled globally in configuration
+        # - The result was served from cache (no API call was made)
+        #
+        # @example With budget enforcement
+        #   # In config/initializers/ruby_llm_agents.rb
+        #   RubyLLM::Agents.configure do |config|
+        #     config.budgets_enabled = true
+        #   end
+        #
+        #   # Budget will be checked before execution
+        #   MyAgent.call(query: "test", tenant: { id: "org_123" })
+        #
+        class Budget < Base
+          # Process budget checking and spend recording
+          #
+          # @param context [Context] The execution context
+          # @return [Context] The context after budget processing
+          # @raise [BudgetExceededError] If budget is exceeded with hard enforcement
+          def call(context)
+            return @app.call(context) unless budgets_enabled?
+
+            # Check budget before execution
+            check_budget!(context)
+
+            # Execute the chain
+            @app.call(context)
+
+            # Record spend after successful execution (if not cached)
+            record_spend!(context) if context.success? && !context.cached?
+
+            context
+          end
+
+          private
+
+          # Returns whether budgets are enabled globally
+          #
+          # @return [Boolean]
+          def budgets_enabled?
+            global_config.budgets_enabled?
+          rescue StandardError
+            false
+          end
+
+          # Checks budget before execution
+          #
+          # @param context [Context] The execution context
+          # @raise [BudgetExceededError] If budget exceeded with hard enforcement
+          def check_budget!(context)
+            BudgetTracker.check!(
+              agent_type: context.agent_class&.name,
+              tenant_id: context.tenant_id,
+              execution_type: context.agent_type&.to_s
+            )
+          rescue BudgetExceededError
+            # Re-raise budget errors
+            raise
+          rescue StandardError => e
+            # Log but don't fail on budget check errors
+            error("Budget check failed: #{e.message}")
+          end
+
+          # Records spend after execution
+          #
+          # @param context [Context] The execution context
+          def record_spend!(context)
+            return unless context.total_cost&.positive?
+
+            BudgetTracker.record_spend!(
+              tenant_id: context.tenant_id,
+              cost: context.total_cost,
+              tokens: context.total_tokens
+            )
+          rescue StandardError => e
+            # Log but don't fail on spend recording errors
+            error("Failed to record spend: #{e.message}")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/pipeline/middleware/cache.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+require "digest"
+
+module RubyLLM
+  module Agents
+    module Pipeline
+      module Middleware
+        # Caches results to avoid redundant API calls.
+        #
+        # This middleware provides caching for agent executions:
+        # - Checks cache before execution
+        # - Stores successful results in cache
+        # - Respects TTL configuration from agent DSL
+        #
+        # Caching is skipped if:
+        # - Caching is not enabled on the agent class (no cache_for DSL)
+        # - The cache store is not configured
+        #
+        # @example Enable caching on an agent
+        #   class MyEmbedder < RubyLLM::Agents::Embedder
+        #     model "text-embedding-3-small"
+        #     cache_for 1.hour
+        #   end
+        #
+        # @example Cache versioning
+        #   class MyEmbedder < RubyLLM::Agents::Embedder
+        #     model "text-embedding-3-small"
+        #     version "2.0"  # Change to invalidate cache
+        #     cache_for 1.hour
+        #   end
+        #
+        class Cache < Base
+          # Process caching
+          #
+          # @param context [Context] The execution context
+          # @return [Context] The context (possibly from cache)
+          def call(context)
+            return @app.call(context) unless cache_enabled?
+
+            cache_key = generate_cache_key(context)
+
+            # Skip cache read if skip_cache is true
+            unless context.skip_cache
+              # Try to read from cache
+              if (cached = cache_read(cache_key))
+                context.output = cached
+                context.cached = true
+                context[:cache_key] = cache_key
+                debug("Cache hit for #{cache_key}")
+                return context
+              end
+            end
+
+            # Execute the chain
+            @app.call(context)
+
+            # Cache successful results
+            if context.success?
+              cache_write(cache_key, context.output)
+              debug("Cache write for #{cache_key}")
+            end
+
+            context
+          end
+
+          private
+
+          # Returns whether caching is enabled for this agent
+          #
+          # @return [Boolean]
+          def cache_enabled?
+            enabled?(:cache_enabled?) && cache_store.present?
+          end
+
+          # Returns the cache store
+          #
+          # @return [ActiveSupport::Cache::Store, nil]
+          def cache_store
+            global_config.cache_store
+          rescue StandardError
+            nil
+          end
+
+          # Returns the cache TTL
+          #
+          # @return [ActiveSupport::Duration, Integer, nil]
+          def cache_ttl
+            config(:cache_ttl)
+          end
+
+          # Generates a cache key for the context
+          #
+          # The cache key includes:
+          # - Namespace prefix
+          # - Agent type
+          # - Agent class name
+          # - Version (for cache invalidation)
+          # - Model
+          # - SHA256 hash of input
+          #
+          # @param context [Context] The execution context
+          # @return [String] The cache key
+          def generate_cache_key(context)
+            components = [
+              "ruby_llm_agents",
+              context.agent_type,
+              context.agent_class&.name,
+              config(:version, "1.0"),
+              context.model,
+              hash_input(context.input)
+            ].compact
+
+            components.join("/")
+          end
+
+          # Hashes the input for cache key
+          #
+          # @param input [Object] The input to hash
+          # @return [String] SHA256 hash
+          def hash_input(input)
+            Digest::SHA256.hexdigest(serialize_input(input))
+          end
+
+          # Serializes input for hashing
+          #
+          # @param input [Object] The input to serialize
+          # @return [String] Serialized representation
+          def serialize_input(input)
+            case input
+            when String
+              input
+            when Array
+              input.map { |i| serialize_input(i) }.join("|")
+            when Hash
+              input.sort.map { |k, v| "#{k}:#{serialize_input(v)}" }.join("|")
+            else
+              input.to_json
+            end
+          rescue StandardError
+            input.to_s
+          end
+
+          # Reads from cache
+          #
+          # @param key [String] Cache key
+          # @return [Object, nil] Cached value or nil
+          def cache_read(key)
+            cache_store.read(key)
+          rescue StandardError => e
+            error("Cache read failed: #{e.message}")
+            nil
+          end
+
+          # Writes to cache
+          #
+          # @param key [String] Cache key
+          # @param value [Object] Value to cache
+          def cache_write(key, value)
+            options = {}
+            options[:expires_in] = cache_ttl if cache_ttl
+
+            cache_store.write(key, value, **options)
+          rescue StandardError => e
+            error("Cache write failed: #{e.message}")
+          end
+        end
+      end
+    end
+  end
+end