raif 1.4.0 → 1.5.0

data/app/controllers/raif/admin/prompt_studio/batch_runs_controller.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    module PromptStudio
+      class BatchRunsController < BaseController
+        def create
+          unless prompt_studio_runs_enabled?
+            redirect_to raif.admin_prompt_studio_tasks_path, alert: t("raif.admin.prompt_studio.common.runs_disabled")
+            return
+          end
+
+          source_tasks = resolve_source_tasks
+          if source_tasks.empty?
+            redirect_to raif.admin_prompt_studio_tasks_path(task_type: params[:task_type]),
+              alert: t("raif.admin.prompt_studio.batch_runs.create.no_tasks_selected")
+            return
+          end
+
+          available_keys = Raif.available_llm_keys.map(&:to_s)
+
+          unless params[:llm_model_key].present? && available_keys.include?(params[:llm_model_key])
+            redirect_to raif.admin_prompt_studio_tasks_path(task_type: params[:task_type]),
+              alert: t("raif.admin.prompt_studio.tasks.rerun.invalid_model")
+            return
+          end
+
+          if params[:judge_type].present? && params[:judge_llm_model_key].present? && !available_keys.include?(params[:judge_llm_model_key])
+            redirect_to raif.admin_prompt_studio_tasks_path(task_type: params[:task_type]),
+              alert: t("raif.admin.prompt_studio.tasks.rerun.invalid_model")
+            return
+          end
+
+          batch_run = Raif::PromptStudioBatchRun.new(
+            task_type: params[:task_type],
+            llm_model_key: params[:llm_model_key],
+            judge_type: params[:judge_type].presence,
+            judge_llm_model_key: params[:judge_llm_model_key].presence,
+            judge_config: build_judge_config,
+            total_count: source_tasks.size
+          )
+
+          batch_run.save!
+
+          source_tasks.each do |task|
+            batch_run.items.create!(source_task: task)
+          end
+
+          Raif::PromptStudioBatchRunJob.perform_later(batch_run: batch_run)
+
+          redirect_to raif.admin_prompt_studio_batch_run_path(batch_run)
+        rescue StandardError => e
+          redirect_to raif.admin_prompt_studio_tasks_path(task_type: params[:task_type]),
+            alert: t("raif.admin.prompt_studio.batch_runs.create.error", message: e.message)
+        end
+
+        def show
+          @batch_run = Raif::PromptStudioBatchRun.find(params[:id])
+          items = @batch_run.items.includes(:source_task, :result_task, :judge_task).order(:id)
+          @pagy, @items = pagy(items)
+        end
+
+      private
+
+        def resolve_source_tasks
+          ids = Array(params[:source_task_ids]).map(&:to_i).reject(&:zero?)
+          scope = Raif::Task.where(id: ids).completed
+          scope = scope.where(type: params[:task_type]) if params[:task_type].present?
+          scope
+        end
+
+        def build_judge_config
+          config = case params[:judge_type]
+          when "Raif::Evals::LlmJudges::Binary"
+            {
+              "criteria" => params[:judge_criteria].presence || "",
+              "strict_mode" => params[:judge_strict_mode] == "1"
+            }
+          when "Raif::Evals::LlmJudges::Scored"
+            {
+              "scoring_rubric" => params[:judge_scoring_rubric].presence || "accuracy"
+            }
+          when "Raif::Evals::LlmJudges::Comparative"
+            {
+              "comparison_criteria" => params[:judge_comparison_criteria].presence || ""
+            }
+          when "Raif::Evals::LlmJudges::Summarization"
+            {}
+          else
+            {}
+          end
+
+          if params[:judge_type].present?
+            config["include_original_prompt_as_context"] = params[:judge_include_original_prompt_as_context] == "1"
+          end
+
+          config
+        end
+      end
+    end
+  end
+end

data/app/controllers/raif/admin/prompt_studio/conversations_controller.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    module PromptStudio
+      class ConversationsController < BaseController
+        def index
+          @conversation_types = Raif::Conversation.distinct.pluck(:type).sort
+          @selected_type = params[:conversation_type] if params[:conversation_type].present?
+          @llm_model_keys = Raif::Conversation.where(type: @selected_type).distinct.pluck(:llm_model_key).compact.sort if @selected_type.present?
+
+          if @selected_type.present?
+            conversations = apply_filters(Raif::Conversation.where(type: @selected_type)).order(created_at: :desc)
+            @pagy, @conversations = pagy(conversations)
+          end
+        end
+
+        def show
+          @conversation = Raif::Conversation.find(params[:id])
+          @comparison = build_prompt_comparison(@conversation)
+        end
+      end
+    end
+  end
+end

data/app/controllers/raif/admin/prompt_studio/tasks_controller.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    module PromptStudio
+      class TasksController < BaseController
+        def index
+          @task_types = Raif::Task.distinct.pluck(:type).sort
+          @selected_type = params[:task_type] if params[:task_type].present?
+          @llm_model_keys = Raif::Task.where(type: @selected_type).distinct.pluck(:llm_model_key).compact.sort if @selected_type.present?
+
+          if @selected_type.present?
+            tasks = apply_filters(Raif::Task.where(type: @selected_type).completed).includes(:raif_model_completion).order(created_at: :desc)
+            @pagy, @tasks = pagy(tasks)
+          end
+
+          @show_batch_runs = prompt_studio_runs_enabled? && @selected_type.present? && @tasks.present?
+        end
+
+        def show
+          @task = Raif::Task.find(params[:id])
+          @comparison = build_prompt_comparison(@task)
+          @original_task = @task.source if @task.prompt_studio_run? && @task.source.is_a?(Raif::Task)
+          @available_llm_keys = Raif.available_llm_keys.map(&:to_s).sort
+        end
+
+        def create
+          original_task = Raif::Task.find(params[:source_task_id])
+
+          unless prompt_studio_runs_enabled?
+            redirect_to raif.admin_prompt_studio_task_path(original_task), alert: t("raif.admin.prompt_studio.common.runs_disabled")
+            return
+          end
+
+          llm_model_key = params[:llm_model_key]
+
+          unless llm_model_key.present? && Raif.available_llm_keys.map(&:to_s).include?(llm_model_key)
+            redirect_to raif.admin_prompt_studio_task_path(original_task), alert: t("raif.admin.prompt_studio.tasks.rerun.invalid_model")
+            return
+          end
+
+          new_task = original_task.class.new(
+            creator: original_task.creator,
+            source: original_task,
+            llm_model_key: llm_model_key,
+            available_model_tools: original_task.available_model_tools,
+            run_with: original_task.run_with,
+            prompt_studio_run: true,
+            started_at: Time.current
+          )
+          new_task.assign_attributes(original_task.prompt_studio_task_attributes)
+          new_task.save!
+          Raif::PromptStudioTaskRunJob.perform_later(task: new_task)
+
+          redirect_to raif.admin_prompt_studio_task_path(new_task)
+        rescue StandardError => e
+          new_task&.update(failed_at: Time.current) unless new_task&.failed_at?
+          redirect_to raif.admin_prompt_studio_task_path(original_task || params[:source_task_id]),
+            alert: t("raif.admin.prompt_studio.tasks.rerun.error", message: e.message)
+        end
+      end
+    end
+  end
+end

data/app/controllers/raif/admin/tasks_controller.rb

@@ -12,6 +12,9 @@ module Raif
         @task_statuses = [:all, :completed, :failed, :in_progress, :pending]
         @selected_statuses = params[:task_statuses].present? ? params[:task_statuses].to_sym : :all
 
+        @selected_llm_model_key = params[:llm_model_key].presence
+        @llm_model_keys = Raif::Task.distinct.order(:llm_model_key).pluck(:llm_model_key)
+
         tasks = Raif::Task.order(created_at: :desc)
         tasks = tasks.where(type: @selected_type) if @selected_type.present? && @selected_type != "all"
 
@@ -28,6 +31,8 @@ module Raif
           end
         end
 
+        tasks = tasks.where(llm_model_key: @selected_llm_model_key) if @selected_llm_model_key.present?
+
         @pagy, @tasks = pagy(tasks)
       end
 

data/app/helpers/raif/application_helper.rb

@@ -3,5 +3,45 @@
 module Raif
   module ApplicationHelper
     include Pagy::Frontend
+
+    def format_task_response(task)
+      if task.response_format_json? && task.raw_response.present?
+        JSON.pretty_generate(JSON.parse(task.raw_response))
+      else
+        task.raw_response
+      end
+    rescue JSON::ParserError
+      task.raw_response
+    end
+
+    def pretty_json(value)
+      JSON.pretty_generate(JSON.parse(value))
+    rescue StandardError
+      value
+    end
+
+    def llm_model_options(selected: nil)
+      options = Raif.available_llm_keys.map do |key|
+        label = I18n.t("raif.model_names.#{key}", default: key.to_s)
+        [label, key.to_s]
+      end.sort_by(&:first)
+
+      options_for_select(options, selected&.to_s)
+    end
+
+    def llm_pricing_json
+      pricing = {}
+      Raif.available_llm_keys.each do |key|
+        config = Raif.llm_config(key)
+        next unless config
+
+        pricing[key.to_s] = {
+          input: config[:input_token_cost] || 0,
+          output: config[:output_token_cost] || 0
+        }
+      end
+
+      pricing.to_json
+    end
   end
 end

data/app/jobs/raif/prompt_studio_batch_run_item_job.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Raif
+  class PromptStudioBatchRunItemJob < ApplicationJob
+
+    def perform(item:)
+      item.execute!
+    end
+
+  end
+end

data/app/jobs/raif/prompt_studio_batch_run_job.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Raif
+  class PromptStudioBatchRunJob < ApplicationJob
+
+    def perform(batch_run:)
+      batch_run.update!(started_at: Time.current)
+
+      batch_run.items.where(status: "pending").find_each do |item|
+        Raif::PromptStudioBatchRunItemJob.perform_later(item: item)
+      end
+    end
+
+  end
+end

data/app/jobs/raif/prompt_studio_task_run_job.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Raif
+  class PromptStudioTaskRunJob < ApplicationJob
+
+    def perform(task:)
+      task.run
+      broadcast_task_result(task)
+    rescue StandardError => e
+      logger.error "Error running prompt studio task: #{e.message}"
+      logger.error e.backtrace&.join("\n")
+
+      task.update(failed_at: Time.current) unless task.failed_at?
+      broadcast_task_result(task)
+    end
+
+  private
+
+    def broadcast_task_result(task)
+      comparison = Raif::PromptStudioComparisonBuilder.build(task)
+      original_task = task.prompt_studio_run? && task.source.is_a?(Raif::Task) ? task.source : nil
+
+      html = Raif::Admin::PromptStudio::TasksController.render(
+        partial: "raif/admin/prompt_studio/tasks/task_result",
+        locals: { task: task, comparison: comparison, original_task: original_task }
+      )
+
+      Turbo::StreamsChannel.broadcast_replace_to(
+        task,
+        target: ActionView::RecordIdentifier.dom_id(task, :result),
+        html: html
+      )
+    end
+
+  end
+end

data/app/models/raif/agent.rb

@@ -35,11 +35,15 @@
 #
 module Raif
   class Agent < ApplicationRecord
+    prepend Raif::Concerns::HasPromptTemplates
+
     include Raif::Concerns::HasLlm
     include Raif::Concerns::HasRequestedLanguage
     include Raif::Concerns::HasAvailableModelTools
+    include Raif::Concerns::HasRuntimeDuration
     include Raif::Concerns::InvokesModelTools
     include Raif::Concerns::AgentInferenceStats
+    include Raif::Concerns::LlmPromptCaching
     include Raif::Concerns::RunWith
 
     belongs_to :creator, polymorphic: true
@@ -122,7 +126,9 @@ module Raif
           source: self,
           system_prompt: system_prompt,
           available_model_tools: native_model_tools,
-          tool_choice: tool_choice_for_iteration
+          tool_choice: tool_choice_for_iteration,
+          anthropic_prompt_caching_enabled: self.class.anthropic_prompt_caching_enabled,
+          bedrock_prompt_caching_enabled: self.class.bedrock_prompt_caching_enabled
         )
 
         logger.debug <<~DEBUG
@@ -137,14 +143,14 @@ module Raif
         DEBUG
 
         process_iteration_model_completion(model_completion)
-        break if final_answer.present?
+        break if final_answer.present? || failed?
       end
 
-      completed!
+      finalize_run!
       final_answer
     rescue StandardError => e
-      self.failed_at = Time.current
-      self.failure_reason = e.message
+      self.failed_at ||= Time.current
+      self.failure_reason ||= e.message
       save!
 
       raise
@@ -160,6 +166,17 @@ module Raif
       # no-op by default. Can be overridden by subclasses to add default model tools
     end
 
+    def finalize_run!
+      validate_successful_completion
+      return if failed?
+
+      completed!
+    end
+
+    def validate_successful_completion
+      # no-op by default. Can be overridden by subclasses to enforce success criteria.
+    end
+
     def process_iteration_model_completion(model_completion)
       raise NotImplementedError, "#{self.class.name} must implement process_iteration_model_completion"
     end
@@ -181,6 +198,14 @@ module Raif
       nil
     end
 
+    # Hook for subclasses to require a specific tool on the current iteration.
+    # Override to align prompt warnings and provider-level tool_choice.
+    # Overrides should be deterministic and side-effect free for a given iteration.
+    # @return [Class, nil] A model tool class, or nil if no specific tool is required.
+    def required_tool_for_iteration
+      nil
+    end
+
     def add_conversation_history_entry(entry)
       entry_stringified = entry.stringify_keys
       conversation_history << entry_stringified
@@ -188,6 +213,12 @@ module Raif
       on_conversation_history_entry.call(entry_stringified) if on_conversation_history_entry.present?
     end
 
+    def fail_run!(reason)
+      self.failed_at ||= Time.current
+      self.failure_reason ||= reason
+      save!
+    end
+
     def build_system_prompt
       raise NotImplementedError, "Subclasses of Raif::Agent must implement build_system_prompt"
     end

data/app/models/raif/agents/native_tool_calling_agent.rb

@@ -88,24 +88,32 @@ module Raif
         available_model_tools_map["agent_final_answer"]
       end
 
-      # Warn the agent that it must provide a final answer on the next iteration
+      def required_tool_for_iteration
+        return final_answer_tool if final_iteration?
+
+        nil
+      end
+
       def before_iteration_llm_chat
-        return unless final_iteration?
+        required_tool = current_iteration_required_tool
+        return if required_tool.blank?
 
         warning_message = Raif::Messages::UserMessage.new(
-          content: I18n.t("raif.agents.native_tool_calling_agent.final_answer_warning")
+          content: required_tool_warning_message(required_tool)
         )
         add_conversation_history_entry(warning_message.to_h)
       end
 
-      # On the final iteration, force the agent to use the agent_final_answer tool
       def tool_choice_for_iteration
-        return unless final_iteration?
+        return current_iteration_required_tool if current_iteration_required_tool.present?
+        return :required if llm.supports_faithful_required_tool_choice?(native_model_tools)
 
-        final_answer_tool
+        log_required_tool_choice_fallback_once!
+        nil
       end
 
       def process_iteration_model_completion(model_completion)
+        required_tool = current_iteration_required_tool
         assistant_response_message = model_completion.parsed_response if model_completion.parsed_response.present?
 
         # The model made no tool call in this completion. Tell it to make a tool call.
@@ -115,36 +123,61 @@ module Raif
             add_conversation_history_entry(assistant_message.to_h)
           end
 
-          error_message = Raif::Messages::UserMessage.new(
-            content: "Error: Previous message contained no tool call. Make a tool call at each step. Available tools: #{available_model_tools_map.keys.join(", ")}" # rubocop:disable Layout/LineLength
-          )
-          add_conversation_history_entry(error_message.to_h)
+          error_content = if required_tool.present?
+            "Error: This iteration required the tool '#{required_tool.tool_name}', but the model response contained no tool call. Available tools: #{available_model_tools_map.keys.join(", ")}" # rubocop:disable Layout/LineLength
+          else
+            "Error: Previous message contained no tool call. Make a tool call at each step. Available tools: #{available_model_tools_map.keys.join(", ")}" # rubocop:disable Layout/LineLength
+          end
+          handle_iteration_error(error_content, required_tool:)
+
+          return
+        end
+
+        # The model returned multiple tool calls. We only allow one per step.
+        if model_completion.response_tool_calls.length > 1
+          if assistant_response_message.present?
+            assistant_message = Raif::Messages::AssistantMessage.new(content: assistant_response_message)
+            add_conversation_history_entry(assistant_message.to_h)
+          end
+
+          error_content = "Error: Multiple tool calls received. Only one tool call is allowed per step. " \
+            "Please call exactly one tool at a time."
+          handle_iteration_error(error_content, required_tool:)
 
           return
         end
 
         tool_call = model_completion.response_tool_calls.first
 
-        # Add the tool call to history
+        tool_name = tool_call["name"]
+        tool_arguments = tool_call["arguments"]
+        tool_klass = available_model_tools_map[tool_name]
+
+        # Prepare tool arguments before recording to history so the history
+        # accurately reflects what was actually invoked
+        tool_arguments = tool_klass.prepare_tool_arguments(tool_arguments) if tool_klass.present?
+
+        # Add the tool call to history (with prepared arguments if tool is known)
         tool_call_message = Raif::Messages::ToolCall.new(
           provider_tool_call_id: tool_call["provider_tool_call_id"],
           name: tool_call["name"],
-          arguments: tool_call["arguments"],
+          arguments: tool_arguments,
           assistant_message: assistant_response_message,
           provider_metadata: tool_call["provider_metadata"]
         )
         add_conversation_history_entry(tool_call_message.to_h)
 
-        tool_name = tool_call["name"]
-        tool_arguments = tool_call["arguments"]
-        tool_klass = available_model_tools_map[tool_name]
+        if required_tool.present? && tool_name != required_tool.tool_name
+          error_content = "Error: This iteration required the tool '#{required_tool.tool_name}', but the model called '#{tool_name}' instead."
+          handle_iteration_error(error_content, required_tool:)
+          return
+        end
 
         # The model tried to use a tool that doesn't exist
         if tool_klass.blank?
           error_content = "Error: Tool '#{tool_name}' is not a valid tool. " \
             "Available tools: #{available_model_tools_map.keys.join(", ")}"
-          error_message = Raif::Messages::UserMessage.new(content: error_content)
-          add_conversation_history_entry(error_message.to_h)
+          handle_iteration_error(error_content, required_tool:)
           return
         end
 
@@ -152,8 +185,7 @@ module Raif
         unless JSON::Validator.validate(tool_klass.tool_arguments_schema, tool_arguments)
           error_content = "Error: Invalid tool arguments for the tool '#{tool_name}'. " \
             "Tool arguments schema: #{tool_klass.tool_arguments_schema.to_json}"
-          error_message = Raif::Messages::UserMessage.new(content: error_content)
-          add_conversation_history_entry(error_message.to_h)
+          handle_iteration_error(error_content, required_tool:)
           return
         end
 
@@ -171,6 +203,56 @@ module Raif
         end
       end
 
+      def validate_successful_completion
+        return if failed? || final_answer.present?
+
+        fail_run!("Agent completed without calling agent_final_answer")
+      end
+
+      def required_tool_warning_message(required_tool)
+        if required_tool == final_answer_tool
+          if final_iteration?
+            I18n.t("raif.agents.native_tool_calling_agent.final_answer_warning")
+          else
+            "Warning: This iteration requires the agent_final_answer tool. If you do not use it now, the next iteration will be your final chance."
+          end
+        else
+          "Warning: This iteration requires the #{required_tool.tool_name} tool."
+        end
+      end
+
+      def current_iteration_required_tool
+        if @required_tool_iteration_count != iteration_count
+          @required_tool_iteration_count = iteration_count
+          @current_iteration_required_tool = required_tool_for_iteration
+        end
+
+        @current_iteration_required_tool
+      end
+
+      def handle_iteration_error(error_content, required_tool: nil)
+        error_message = Raif::Messages::UserMessage.new(content: error_content)
+        add_conversation_history_entry(error_message.to_h)
+
+        return if required_tool.blank? || retry_iteration_available?
+
+        fail_run!(error_content)
+      end
+
+      def retry_iteration_available?
+        iteration_count < max_iterations
+      end
+
+      def log_required_tool_choice_fallback_once!
+        return if @logged_required_tool_choice_fallback
+
+        @logged_required_tool_choice_fallback = true
+        Raif.logger.warn(
+          "NativeToolCallingAgent is falling back to runtime tool-call validation because #{llm.key} " \
+            "cannot faithfully enforce tool_choice: :required for tools: #{available_model_tools_map.keys.join(", ")}"
+        )
+      end
+
       def ensure_llm_supports_native_tool_use
         unless llm.supports_native_tool_use?
           errors.add(:base, "Raif::Agent#llm_model_key must use an LLM that supports native tool use")

data/app/models/raif/concerns/has_prompt_templates.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module Raif
+  module Concerns
+    module HasPromptTemplates
+      extend ActiveSupport::Concern
+
+      class TemplateContext < ActionView::Base.with_empty_template_cache
+        def initialize(lookup_context, instance)
+          super(lookup_context, {}, nil)
+          @_instance = instance
+        end
+
+        def method_missing(method_name, ...)
+          if @_instance.respond_to?(method_name)
+            @_instance.public_send(method_name, ...)
+          else
+            super
+          end
+        end
+
+        def respond_to_missing?(method_name, include_private = false)
+          @_instance.respond_to?(method_name, include_private) || super
+        end
+      end
+
+      class_methods do
+        # Returns the template prefix path derived from the class name.
+        # e.g. Raif::Tasks::SummarizeDocument -> "raif/tasks/summarize_document"
+        # e.g. Raif::Tasks::Docs::Summarize  -> "raif/tasks/docs/summarize"
+        def prompt_template_prefix
+          name.underscore
+        end
+
+        def prompt_template_view_paths
+          ActionController::Base.view_paths
+        end
+      end
+
+      def build_prompt
+        if prompt_template_exists?(:prompt)
+          render_prompt_template(:prompt)
+        else
+          super
+        end
+      end
+
+      def build_system_prompt
+        if prompt_template_exists?(:system_prompt)
+          render_prompt_template(:system_prompt)
+        else
+          super
+        end
+      end
+
+    private
+
+      def prompt_template_name
+        self.class.prompt_template_prefix.split("/").last
+      end
+
+      def prompt_template_dir
+        File.dirname(self.class.prompt_template_prefix)
+      end
+
+      def prompt_template_exists?(template_type)
+        prompt_lookup_context_for(template_type).exists?(prompt_template_name, prompt_template_dir)
+      end
+
+      def prompt_lookup_context_for(template_type)
+        lookup = ActionView::LookupContext.new(self.class.prompt_template_view_paths)
+        lookup.formats = [template_type]
+        lookup
+      end
+
+      def render_prompt_template(template_type)
+        lookup = prompt_lookup_context_for(template_type)
+        context = TemplateContext.new(lookup, self)
+        context.render(template: "#{prompt_template_dir}/#{prompt_template_name}").strip
+      rescue ActionView::Template::Error, ActionView::MissingTemplate => e
+        raise Raif::Errors::PromptTemplateError.new(
+          template_path: "#{self.class.prompt_template_prefix}.#{template_type}.erb",
+          original_error: e
+        )
+      end
+    end
+  end
+end