ruby_llm-agents 1.3.4 → 2.1.0

data/lib/ruby_llm/agents/base_agent.rb

@@ -15,7 +15,6 @@ module RubyLLM
     # @example Creating an agent
     #   class SearchAgent < RubyLLM::Agents::BaseAgent
     #     model "gpt-4o"
-    #     version "1.0"
     #     description "Searches for relevant documents"
     #     timeout 30
     #
@@ -246,24 +245,38 @@ module RubyLLM
 
       # User prompt to send to the LLM
       #
-      # @abstract Subclasses must implement this method
+      # If a class-level `prompt` DSL is defined (string template or block),
+      # it will be used. Otherwise, subclasses must implement this method.
+      #
       # @return [String] The user prompt
       def user_prompt
-        raise NotImplementedError, "#{self.class} must implement #user_prompt"
+        prompt_config = self.class.prompt_config
+        return resolve_prompt_from_config(prompt_config) if prompt_config
+
+        raise NotImplementedError, "#{self.class} must implement #user_prompt or use the prompt DSL"
       end
 
       # System prompt for LLM instructions
       #
+      # If a class-level `system` DSL is defined, it will be used.
+      # Otherwise returns nil.
+      #
       # @return [String, nil] System instructions, or nil for none
       def system_prompt
+        system_config = self.class.system_config
+        return resolve_prompt_from_config(system_config) if system_config
+
         nil
       end
 
       # Response schema for structured output
       #
+      # Delegates to the class-level schema DSL by default.
+      # Override in subclass instances to customize per-instance.
+      #
       # @return [RubyLLM::Schema, nil] Schema definition, or nil for free-form
       def schema
-        nil
+        self.class.schema
       end
 
       # Conversation history for multi-turn conversations
@@ -288,9 +301,12 @@ module RubyLLM
 
       # Generates the cache key for this agent invocation
       #
-      # @return [String] Cache key in format "ruby_llm_agent/ClassName/version/hash"
+      # Cache keys are content-based, using a hash of the prompts and parameters.
+      # This automatically invalidates caches when prompts change.
+      #
+      # @return [String] Cache key in format "ruby_llm_agent/ClassName/hash"
       def agent_cache_key
-        ["ruby_llm_agent", self.class.name, self.class.version, cache_key_hash].join("/")
+        ["ruby_llm_agent", self.class.name, cache_key_hash].join("/")
       end
 
       # Generates a hash of the cache key data
@@ -462,6 +478,36 @@ module RubyLLM
         end
       end
 
+      # Resolves a prompt from DSL configuration (template string or block)
+      #
+      # For string templates, interpolates {placeholder} with parameter values.
+      # For blocks, evaluates in the instance context.
+      #
+      # @param config [String, Proc] The prompt configuration
+      # @return [String] The resolved prompt
+      def resolve_prompt_from_config(config)
+        case config
+        when String
+          interpolate_template(config)
+        when Proc
+          instance_eval(&config)
+        else
+          config.to_s
+        end
+      end
+
+      # Interpolates {placeholder} patterns in a template string
+      #
+      # @param template [String] Template with {placeholder} syntax
+      # @return [String] Interpolated string
+      def interpolate_template(template)
+        template.gsub(/\{(\w+)\}/) do
+          param_name = ::Regexp.last_match(1).to_sym
+          value = send(param_name) if respond_to?(param_name)
+          value.to_s
+        end
+      end
+
       # Execute the core LLM call
       #
       # This is called by the Pipeline::Executor after all middleware
@@ -567,29 +613,14 @@ module RubyLLM
         input_tokens = context.input_tokens || 0
         output_tokens = context.output_tokens || 0
 
-        input_price = extract_model_price(model_info, :input_price)
-        output_price = extract_model_price(model_info, :output_price)
+        input_price = model_info.pricing&.text_tokens&.input || 0
+        output_price = model_info.pricing&.text_tokens&.output || 0
 
         context.input_cost = (input_tokens / 1_000_000.0) * input_price
         context.output_cost = (output_tokens / 1_000_000.0) * output_price
         context.total_cost = (context.input_cost + context.output_cost).round(6)
       end
 
-      # Extracts price from model info (supports both hash and object access)
-      #
-      # @param model_info [Hash, Object] Model info
-      # @param key [Symbol] The price key
-      # @return [Float] The price, or 0 if not found
-      def extract_model_price(model_info, key)
-        if model_info.respond_to?(key)
-          model_info.send(key) || 0
-        elsif model_info.respond_to?(:[])
-          model_info[key] || 0
-        else
-          0
-        end
-      end
-
       # Finds model pricing info
       #
       # @param model_id [String] The model ID

data/lib/ruby_llm/agents/core/base/callbacks.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # DSL and execution support for before_call/after_call hooks
+    #
+    # Provides callbacks that run before and after the LLM call,
+    # allowing custom preprocessing (redaction, moderation, validation)
+    # and postprocessing (logging, transformation) logic.
+    #
+    # @example Using callbacks
+    #   class MyAgent < ApplicationAgent
+    #     before_call :sanitize_input
+    #     after_call :log_response
+    #
+    #     # Or with blocks
+    #     before_call { |context| context.params[:timestamp] = Time.current }
+    #     after_call { |context, response| notify(response) }
+    #
+    #     private
+    #
+    #     def sanitize_input(context)
+    #       # Mutate context as needed
+    #       # Raise to block execution
+    #     end
+    #
+    #     def log_response(context, response)
+    #       Rails.logger.info("Response: #{response}")
+    #     end
+    #   end
+    #
+    module CallbacksDSL
+      # Add a callback to run before the LLM call
+      #
+      # Callbacks receive the pipeline context and can:
+      # - Mutate the context (params, prompts, etc.)
+      # - Raise an exception to block execution
+      # - Return value is ignored
+      #
+      # @param method_name [Symbol, nil] Instance method to call
+      # @yield [context] Block to execute
+      # @yieldparam context [Pipeline::Context] The execution context
+      # @return [void]
+      #
+      # @example With method name
+      #   before_call :validate_input
+      #
+      # @example With block
+      #   before_call { |context| context.params[:sanitized] = true }
+      #
+      def before_call(method_name = nil, &block)
+        @callbacks ||= { before: [], after: [] }
+        @callbacks[:before] << (block || method_name)
+      end
+
+      # Add a callback to run after the LLM call
+      #
+      # Callbacks receive the pipeline context and the response.
+      # Return value is ignored.
+      #
+      # @param method_name [Symbol, nil] Instance method to call
+      # @yield [context, response] Block to execute
+      # @yieldparam context [Pipeline::Context] The execution context
+      # @yieldparam response [Object] The LLM response
+      # @return [void]
+      #
+      # @example With method name
+      #   after_call :log_response
+      #
+      # @example With block
+      #   after_call { |context, response| notify_completion(response) }
+      #
+      def after_call(method_name = nil, &block)
+        @callbacks ||= { before: [], after: [] }
+        @callbacks[:after] << (block || method_name)
+      end
+
+      # Simplified alias for before_call (block-only)
+      #
+      # This is the preferred method in the simplified DSL.
+      #
+      # @yield [context] Block to execute before the LLM call
+      # @yieldparam context [Pipeline::Context] The execution context
+      # @return [void]
+      #
+      # @example
+      #   before { |ctx| ctx.params[:timestamp] = Time.current }
+      #   before { |ctx| validate_input!(ctx.params[:query]) }
+      #
+      def before(&block)
+        before_call(&block)
+      end
+
+      # Simplified alias for after_call (block-only)
+      #
+      # This is the preferred method in the simplified DSL.
+      #
+      # @yield [context, response] Block to execute after the LLM call
+      # @yieldparam context [Pipeline::Context] The execution context
+      # @yieldparam response [Object] The LLM response
+      # @return [void]
+      #
+      # @example
+      #   after { |ctx, result| Rails.logger.info("Completed: #{result}") }
+      #   after { |ctx, result| notify_slack(result) if result.confidence < 0.5 }
+      #
+      def after(&block)
+        after_call(&block)
+      end
+
+      # Get all registered callbacks
+      #
+      # @return [Hash] Hash with :before and :after arrays
+      def callbacks
+        @callbacks ||= { before: [], after: [] }
+      end
+    end
+
+    # Instance methods for running callbacks
+    module CallbacksExecution
+      private
+
+      # Run callbacks of the specified type
+      #
+      # @param type [Symbol] :before or :after
+      # @param args [Array] Arguments to pass to callbacks
+      # @return [void]
+      def run_callbacks(type, *args)
+        callbacks = self.class.callbacks[type] || []
+
+        callbacks.each do |callback|
+          case callback
+          when Symbol
+            send(callback, *args)
+          when Proc
+            instance_exec(*args, &callback)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -1,20 +1,18 @@
 # frozen_string_literal: true
 
-require_relative "base/moderation_dsl"
-require_relative "base/moderation_execution"
+require_relative "base/callbacks"
 
 module RubyLLM
   module Agents
     # Base class for LLM-powered conversational agents
     #
     # Inherits from BaseAgent to use the middleware pipeline architecture
-    # while adding moderation capabilities for input/output content filtering.
+    # while adding callback hooks for custom preprocessing and postprocessing.
     #
     # @example Creating an agent
     #   class SearchAgent < ApplicationAgent
     #     model "gpt-4o"
     #     temperature 0.0
-    #     version "1.0"
     #     timeout 30
     #     cache_for 1.hour
     #
@@ -30,14 +28,19 @@ module RubyLLM
     #     end
     #   end
     #
-    # @example With moderation
+    # @example With callbacks
     #   class SafeAgent < ApplicationAgent
-    #     moderation :input, :output
-    #     # or
-    #     moderation :both
+    #     before_call :redact_pii
+    #     after_call :log_response
     #
-    #     def user_prompt
-    #       query
+    #     private
+    #
+    #     def redact_pii(context)
+    #       # Custom redaction logic
+    #     end
+    #
+    #     def log_response(context, response)
+    #       Rails.logger.info("Response received")
     #     end
     #   end
     #
@@ -49,8 +52,8 @@ module RubyLLM
     # @see RubyLLM::Agents::BaseAgent
     # @api public
     class Base < BaseAgent
-      extend ModerationDSL
-      include ModerationExecution
+      extend CallbacksDSL
+      include CallbacksExecution
 
       class << self
         # Returns the agent type for conversation agents
@@ -61,43 +64,27 @@ module RubyLLM
         end
       end
 
-      # Execute the core LLM call with moderation support
+      # Execute the core LLM call with callback support
       #
-      # This extends BaseAgent's execute method to add input and output
-      # moderation checks when configured via the moderation DSL.
+      # This extends BaseAgent's execute method to add before/after
+      # callback hooks for custom preprocessing and postprocessing.
       #
       # @param context [Pipeline::Context] The execution context
       # @return [void] Sets context.output with the result
       def execute(context)
         @execution_started_at = context.started_at || Time.current
 
-        # Input moderation check (before LLM call)
-        if self.class.moderation_enabled? && should_moderate?(:input)
-          input_text = build_moderation_input
-          moderate_input(input_text)
+        # Run before_call callbacks
+        run_callbacks(:before, context)
 
-          if moderation_blocked?
-            context.output = build_moderation_blocked_result(:input)
-            return
-          end
-        end
-
-        # Execute the LLM call via parent
+        # Execute the LLM call
         client = build_client(context)
         response = execute_llm_call(client, context)
         capture_response(response, context)
         processed_content = process_response(response)
 
-        # Output moderation check (after LLM call)
-        if self.class.moderation_enabled? && should_moderate?(:output)
-          output_text = processed_content.is_a?(String) ? processed_content : processed_content.to_s
-          moderate_output(output_text)
-
-          if moderation_blocked?
-            context.output = build_moderation_blocked_result(:output)
-            return
-          end
-        end
+        # Run after_call callbacks
+        run_callbacks(:after, context, response)
 
         context.output = build_result(processed_content, response, context)
       end
@@ -111,25 +98,6 @@ module RubyLLM
 
         tenant.is_a?(Hash) ? tenant[:id]&.to_s : nil
       end
-
-      private
-
-      # Check if execution model is available for moderation tracking
-      #
-      # @return [Boolean] true if Execution model can be used
-      def execution_model_available?
-        return @execution_model_available if defined?(@execution_model_available)
-
-        @execution_model_available = begin
-          RubyLLM::Agents::Execution.table_exists?
-        rescue StandardError
-          false
-        end
-      end
     end
   end
 end
-
-# Load moderation modules after class is defined (they reopen the class)
-require_relative "base/moderation_dsl"
-require_relative "base/moderation_execution"