raif 1.3.0 → 1.5.0

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/admin/task_stat.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Raif
+  module Admin
+    TaskStat = Data.define(:type, :llm_model_key, :count, :input_cost, :output_cost, :total_cost)
+  end
+end

data/app/models/raif/agent.rb

@@ -1,18 +1,59 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_agents
+#
+#  id                     :bigint           not null, primary key
+#  available_model_tools  :jsonb            not null
+#  completed_at           :datetime
+#  conversation_history   :jsonb            not null
+#  creator_type           :string           not null
+#  failed_at              :datetime
+#  failure_reason         :text
+#  final_answer           :text
+#  iteration_count        :integer          default(0), not null
+#  llm_model_key          :string           not null
+#  max_iterations         :integer          default(10), not null
+#  requested_language_key :string
+#  run_with               :jsonb
+#  source_type            :string
+#  started_at             :datetime
+#  system_prompt          :text
+#  task                   :text
+#  type                   :string           not null
+#  created_at             :datetime         not null
+#  updated_at             :datetime         not null
+#  creator_id             :bigint           not null
+#  source_id              :bigint
+#
+# Indexes
+#
+#  index_raif_agents_on_created_at  (created_at)
+#  index_raif_agents_on_creator     (creator_type,creator_id)
+#  index_raif_agents_on_source      (source_type,source_id)
+#
 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
+    belongs_to :source, polymorphic: true, optional: true
 
     has_many :raif_model_completions, as: :source, dependent: :destroy, class_name: "Raif::ModelCompletion"
 
     after_initialize -> { self.available_model_tools ||= [] }
     after_initialize -> { self.conversation_history ||= [] }
+    after_initialize -> { self.run_with ||= {} }
 
     boolean_timestamp :started_at
     boolean_timestamp :completed_at
@@ -69,16 +110,25 @@ module Raif
         Task: #{task}
       DEBUG
 
-      add_conversation_history_entry({ role: "user", content: task })
+      add_conversation_history_entry(Raif::Messages::UserMessage.new(content: task).to_h)
 
       while iteration_count < max_iterations
         update_columns(iteration_count: iteration_count + 1)
 
+        # Update the system prompt on each iteration in case it has changed since the last iteration
+        self.system_prompt = build_system_prompt
+
+        # Hook for subclasses to perform actions before the LLM chat (e.g., add warnings)
+        before_iteration_llm_chat
+
         model_completion = llm.chat(
           messages: conversation_history,
           source: self,
           system_prompt: system_prompt,
-          available_model_tools: native_model_tools
+          available_model_tools: native_model_tools,
+          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
@@ -93,25 +143,40 @@ 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
     end
 
+    def final_iteration?
+      iteration_count == max_iterations
+    end
+
   private
 
     def populate_default_model_tools
       # 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
@@ -120,6 +185,27 @@ module Raif
       # no-op by default
     end
 
+    # Hook for subclasses to perform actions before the LLM chat on each iteration
+    # Override in subclasses to add warnings, context, etc.
+    def before_iteration_llm_chat
+      # no-op by default
+    end
+
+    # Hook for subclasses to specify tool_choice for the current iteration
+    # Override in subclasses to force specific tools (e.g., on final iteration)
+    # @return [Class, nil] A model tool class (e.g., Raif::ModelTools::AgentFinalAnswer), or nil for default behavior
+    def tool_choice_for_iteration
+      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
@@ -127,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

@@ -1,5 +1,38 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_agents
+#
+#  id                     :bigint           not null, primary key
+#  available_model_tools  :jsonb            not null
+#  completed_at           :datetime
+#  conversation_history   :jsonb            not null
+#  creator_type           :string           not null
+#  failed_at              :datetime
+#  failure_reason         :text
+#  final_answer           :text
+#  iteration_count        :integer          default(0), not null
+#  llm_model_key          :string           not null
+#  max_iterations         :integer          default(10), not null
+#  requested_language_key :string
+#  run_with               :jsonb
+#  source_type            :string
+#  started_at             :datetime
+#  system_prompt          :text
+#  task                   :text
+#  type                   :string           not null
+#  created_at             :datetime         not null
+#  updated_at             :datetime         not null
+#  creator_id             :bigint           not null
+#  source_id              :bigint
+#
+# Indexes
+#
+#  index_raif_agents_on_created_at  (created_at)
+#  index_raif_agents_on_creator     (creator_type,creator_id)
+#  index_raif_agents_on_source      (source_type,source_id)
+#
 module Raif
   module Agents
     class NativeToolCallingAgent < Raif::Agent
@@ -12,7 +45,10 @@ module Raif
       }
 
       before_validation -> {
-        available_model_tools << "Raif::ModelTools::AgentFinalAnswer" unless available_model_tools.include?("Raif::ModelTools::AgentFinalAnswer")
+        # If there is no final answer tool added, add it
+        unless available_model_tools_map.key?("agent_final_answer")
+          available_model_tools << "Raif::ModelTools::AgentFinalAnswer"
+        end
       }
 
       def build_system_prompt
@@ -28,7 +64,7 @@ module Raif
           6. Provide a final answer to the user's request.
 
           For your final answer:
-          - Use the agent_final_answer tool/function with your complete answer as the "final_answer" parameter.
+          - You **MUST** use the agent_final_answer tool/function to provide your final answer.
           - Your answer should be comprehensive and directly address the user's request.
 
           Guidelines
@@ -48,82 +84,173 @@ module Raif
         available_model_tools
       end
 
+      def final_answer_tool
+        available_model_tools_map["agent_final_answer"]
+      end
+
+      def required_tool_for_iteration
+        return final_answer_tool if final_iteration?
+
+        nil
+      end
+
+      def before_iteration_llm_chat
+        required_tool = current_iteration_required_tool
+        return if required_tool.blank?
+
+        warning_message = Raif::Messages::UserMessage.new(
+          content: required_tool_warning_message(required_tool)
+        )
+        add_conversation_history_entry(warning_message.to_h)
+      end
+
+      def tool_choice_for_iteration
+        return current_iteration_required_tool if current_iteration_required_tool.present?
+        return :required if llm.supports_faithful_required_tool_choice?(native_model_tools)
+
+        log_required_tool_choice_fallback_once!
+        nil
+      end
+
       def process_iteration_model_completion(model_completion)
-        if model_completion.parsed_response.present?
-          add_conversation_history_entry({
-            role: "assistant",
-            content: model_completion.parsed_response
-          })
-        end
+        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.
         if model_completion.response_tool_calls.blank?
-          add_conversation_history_entry({
-            role: "assistant",
-            content: "<observation>Error: No tool call found. I need to make a tool call at each step. Available tools: #{available_model_tools_map.keys.join(", ")}</observation>" # rubocop:disable Layout/LineLength
-          })
+          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 = 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
 
-        tool_call = model_completion.response_tool_calls.first
+        # 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:)
 
-        unless tool_call["name"] && tool_call["arguments"]
-          add_conversation_history_entry({
-            role: "assistant",
-            content: "<observation>Error: Invalid action specified. Please provide a valid action, formatted as a JSON object with 'tool' and 'arguments' keys.</observation>" # rubocop:disable Layout/LineLength
-          })
           return
         end
 
+        tool_call = model_completion.response_tool_calls.first
+
         tool_name = tool_call["name"]
         tool_arguments = tool_call["arguments"]
+        tool_klass = available_model_tools_map[tool_name]
 
-        # Add assistant's response to conversation history (without the actual tool calls)
-        # add_conversation_history_entry({
-        #   role: "assistant",
-        #   content: "<thought>I need to use the #{tool_name} tool to help with this task.</thought>"
-        # })
+        # 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?
 
-        # Check if we have a final answer. If yes, we're done.
-        if tool_name == "agent_final_answer"
-          self.final_answer = tool_arguments["final_answer"]
-          add_conversation_history_entry({ role: "assistant", content: "<answer>#{final_answer}</answer>" })
+        # 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_arguments,
+          assistant_message: assistant_response_message,
+          provider_metadata: tool_call["provider_metadata"]
+        )
+        add_conversation_history_entry(tool_call_message.to_h)
+
+        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
 
-        # Add the tool call to conversation history
-        add_conversation_history_entry({
-          role: "assistant",
-          content: "<action>#{JSON.pretty_generate(tool_call)}</action>"
-        })
-
-        # Find the tool class and process it
-        tool_klass = available_model_tools_map[tool_name]
-
         # The model tried to use a tool that doesn't exist
-        unless tool_klass
-          add_conversation_history_entry({
-            role: "assistant",
-            content: "<observation>Error: Tool '#{tool_name}' not found. Available tools: #{available_model_tools_map.keys.join(", ")}</observation>"
-          })
+        if tool_klass.blank?
+          error_content = "Error: Tool '#{tool_name}' is not a valid tool. " \
+            "Available tools: #{available_model_tools_map.keys.join(", ")}"
+          handle_iteration_error(error_content, required_tool:)
           return
         end
 
+        # Make sure the tool arguments match the tool's schema
         unless JSON::Validator.validate(tool_klass.tool_arguments_schema, tool_arguments)
-          add_conversation_history_entry({
-            role: "assistant",
-            content: "<observation>Error: Invalid tool arguments. Please provide valid arguments for the tool '#{tool_name}'. Tool arguments schema: #{tool_klass.tool_arguments_schema.to_json}</observation>" # rubocop:disable Layout/LineLength
-          })
+          error_content = "Error: Invalid tool arguments for the tool '#{tool_name}'. " \
+            "Tool arguments schema: #{tool_klass.tool_arguments_schema.to_json}"
+          handle_iteration_error(error_content, required_tool:)
           return
         end
 
-        # Process the tool and add observation to history
-        tool_invocation = tool_klass.invoke_tool(tool_arguments: tool_arguments, source: self)
-        observation = tool_klass.observation_for_invocation(tool_invocation)
+        # Process the tool invocation and add observation/result to history
+        tool_invocation = tool_klass.invoke_tool(
+          provider_tool_call_id: tool_call["provider_tool_call_id"],
+          tool_arguments: tool_arguments,
+          source: self
+        )
+
+        if tool_name == "agent_final_answer"
+          self.final_answer = tool_invocation.result
+        else
+          add_conversation_history_entry(tool_invocation.as_tool_call_result_message)
+        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
 
-        add_conversation_history_entry({
-          role: "assistant",
-          content: "<observation>#{observation}</observation>"
-        })
+        @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

data/app/models/raif/application_record.rb

@@ -8,6 +8,24 @@ class Raif::ApplicationRecord < Raif.config.model_superclass.constantize
   scope :newest_first, -> { order(created_at: :desc) }
   scope :oldest_first, -> { order(created_at: :asc) }
 
+  # Returns a scope that checks if a JSON column is not blank (not null and not empty array)
+  # @param column_name [Symbol, String] the name of the JSON column
+  # @return [ActiveRecord::Relation]
+  def self.where_json_not_blank(column_name)
+    quoted_column = connection.quote_column_name(column_name.to_s)
+
+    case connection.adapter_name.downcase
+    when "postgresql"
+      where.not(column_name => nil)
+        .where("jsonb_array_length(#{quoted_column}) > 0")
+    when "mysql2", "trilogy"
+      where.not(column_name => nil)
+        .where("JSON_LENGTH(#{quoted_column}) > 0")
+    else
+      raise "Unsupported database: #{connection.adapter_name}"
+    end
+  end
+
   def self.table_name_prefix
     "raif_"
   end

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

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::AgentInferenceStats
+  extend ActiveSupport::Concern
+
+  # Returns the total number of prompt tokens across all model completions
+  def total_prompt_tokens
+    @total_prompt_tokens ||= raif_model_completions.sum(:prompt_tokens)
+  end
+
+  # Returns the total number of completion tokens across all model completions
+  def total_completion_tokens
+    @total_completion_tokens ||= raif_model_completions.sum(:completion_tokens)
+  end
+
+  # Returns the total number of tokens across all model completions
+  def total_tokens_sum
+    @total_tokens_sum ||= raif_model_completions.sum(:total_tokens)
+  end
+
+  # Returns the total cost of prompt tokens across all model completions
+  def total_prompt_token_cost
+    @total_prompt_token_cost ||= raif_model_completions.sum(:prompt_token_cost)
+  end
+
+  # Returns the total cost of output tokens across all model completions
+  def total_output_token_cost
+    @total_output_token_cost ||= raif_model_completions.sum(:output_token_cost)
+  end
+
+  # Returns the total cost across all model completions
+  def total_cost
+    @total_cost ||= raif_model_completions.sum(:total_cost)
+  end
+end

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