raif 1.3.0 → 1.4.0

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,91 @@ module Raif
         available_model_tools
       end
 
+      def final_answer_tool
+        available_model_tools_map["agent_final_answer"]
+      end
+
+      # Warn the agent that it must provide a final answer on the next iteration
+      def before_iteration_llm_chat
+        return unless final_iteration?
+
+        warning_message = Raif::Messages::UserMessage.new(
+          content: I18n.t("raif.agents.native_tool_calling_agent.final_answer_warning")
+        )
+        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?
+
+        final_answer_tool
+      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
+        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_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)
+
           return
         end
 
         tool_call = model_completion.response_tool_calls.first
 
-        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
+        # Add the tool call to history
+        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"],
+          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"]
-
-        # 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>"
-        # })
-
-        # 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>" })
-          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(", ")}"
+          error_message = Raif::Messages::UserMessage.new(content: error_content)
+          add_conversation_history_entry(error_message.to_h)
           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}"
+          error_message = Raif::Messages::UserMessage.new(content: error_content)
+          add_conversation_history_entry(error_message.to_h)
           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
+        )
 
-        add_conversation_history_entry({
-          role: "assistant",
-          content: "<observation>#{observation}</observation>"
-        })
+        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 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/json_schema_definition.rb

@@ -9,19 +9,54 @@ module Raif
         def json_schema_definition(schema_name, &block)
           raise ArgumentError, "A block must be provided to define the JSON schema" unless block_given?
 
-          @schemas ||= {}
-          @schemas[schema_name] = Raif::JsonSchemaBuilder.new
-          @schemas[schema_name].instance_eval(&block)
-          @schemas[schema_name]
+          # Check if block expects an instance parameter (arity == 1)
+          # arity == 0: no parameters (class-level schema)
+          # arity == 1: one parameter (instance-dependent schema)
+          if block.arity == 1
+            # Store block for instance-dependent schema building
+            @schema_blocks ||= {}
+            @schema_blocks[schema_name] = block
+          else
+            # Build schema immediately for class-level (backward compatible)
+            @schemas ||= {}
+            @schemas[schema_name] = Raif::JsonSchemaBuilder.new
+            @schemas[schema_name].instance_eval(&block)
+          end
         end
 
         def schema_defined?(schema_name)
-          @schemas&.dig(schema_name).present?
+          @schemas&.dig(schema_name).present? || @schema_blocks&.dig(schema_name).present?
         end
 
         def schema_for(schema_name)
+          # Check if this is an instance-dependent schema
+          if @schema_blocks&.dig(schema_name).present?
+            raise Raif::Errors::InstanceDependentSchemaError,
+              "The schema '#{schema_name}' is instance-dependent and cannot be accessed at the class level. " \
+                "Call this method on an instance instead."
+          end
+
           @schemas[schema_name].to_schema
         end
+
+        def instance_dependent_schema?(schema_name)
+          @schema_blocks&.dig(schema_name).present?
+        end
+      end
+
+      # Instance method to build schema with instance context
+      def schema_for_instance(schema_name)
+        block = self.class.instance_variable_get(:@schema_blocks)&.[](schema_name)
+
+        if block
+          # Build schema with instance context
+          builder = Raif::JsonSchemaBuilder.new
+          builder.build_with_instance(self, &block)
+          builder.to_schema
+        elsif self.class.schema_defined?(schema_name)
+          # Fall back to class-level schema
+          self.class.schema_for(schema_name)
+        end
       end
     end
   end

data/app/models/raif/concerns/llms/anthropic/message_formatting.rb

@@ -48,4 +48,32 @@ module Raif::Concerns::Llms::Anthropic::MessageFormatting
       raise Raif::Errors::InvalidModelFileInputError, "Invalid model file input source type: #{file_input.source_type}"
     end
   end
+
+  def format_tool_call_message(tool_call)
+    content_array = []
+    content_array << format_string_message(tool_call["assistant_message"]) if tool_call["assistant_message"].present?
+
+    content_array << {
+      "type" => "tool_use",
+      "id" => tool_call["provider_tool_call_id"],
+      "name" => tool_call["name"],
+      "input" => tool_call["arguments"]
+    }
+
+    {
+      "role" => "assistant",
+      "content" => content_array
+    }
+  end
+
+  def format_tool_call_result_message(tool_call_result)
+    {
+      "role" => "user",
+      "content" => [{
+        "type" => "tool_result",
+        "tool_use_id" => tool_call_result["provider_tool_call_id"],
+        "content" => tool_call_result["result"].is_a?(String) ? tool_call_result["result"] : JSON.generate(tool_call_result["result"])
+      }]
+    }
+  end
 end

data/app/models/raif/concerns/llms/anthropic/response_tool_calls.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::Anthropic::ResponseToolCalls
+  extend ActiveSupport::Concern
+
+  def extract_response_tool_calls(resp)
+    return if resp&.dig("content").nil?
+
+    # Find any tool_use content blocks
+    tool_uses = resp&.dig("content")&.select do |content|
+      content["type"] == "tool_use"
+    end
+
+    return if tool_uses.blank?
+
+    tool_uses.map do |tool_use|
+      {
+        "provider_tool_call_id" => tool_use["id"],
+        "name" => tool_use["name"],
+        "arguments" => tool_use["input"],
+      }
+    end
+  end
+end

data/app/models/raif/concerns/llms/anthropic/tool_formatting.rb

@@ -53,4 +53,8 @@ module Raif::Concerns::Llms::Anthropic::ToolFormatting
         "Invalid provider-managed tool: #{tool.name} for #{key}"
     end
   end
+
+  def build_forced_tool_choice(tool_name)
+    { "type" => "tool", "name" => tool_name }
+  end
 end

data/app/models/raif/concerns/llms/bedrock/message_formatting.rb

@@ -67,4 +67,40 @@ module Raif::Concerns::Llms::Bedrock::MessageFormatting
       "text/markdown" => "md"
     }[content_type]
   end
+
+  def format_tool_call_message(tool_call)
+    content_array = []
+    content_array << format_string_message(tool_call["assistant_message"]) if tool_call["assistant_message"].present?
+
+    content_array << {
+      "tool_use" => {
+        "tool_use_id" => tool_call["provider_tool_call_id"],
+        "name" => tool_call["name"],
+        "input" => tool_call["arguments"]
+      }
+    }
+
+    {
+      "role" => "assistant",
+      "content" => content_array
+    }
+  end
+
+  def format_tool_call_result_message(tool_call_result)
+    tool_result_content = if tool_call_result["result"].is_a?(String)
+      { "text" => tool_call_result["result"] }
+    else
+      { "json" => tool_call_result["result"] }
+    end
+
+    {
+      "role" => "user",
+      "content" => [{
+        "tool_result" => {
+          "tool_use_id" => tool_call_result["provider_tool_call_id"],
+          "content" => [tool_result_content]
+        }
+      }]
+    }
+  end
 end

data/app/models/raif/concerns/llms/bedrock/response_tool_calls.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::Bedrock::ResponseToolCalls
+  extend ActiveSupport::Concern
+
+  def extract_response_tool_calls(resp)
+    # Get the message from the response object
+    message = resp.output.message
+    return if message.content.nil?
+
+    # Find any tool_use blocks in the content array
+    tool_uses = message.content.select do |content|
+      content.respond_to?(:tool_use) && content.tool_use.present?
+    end
+
+    return if tool_uses.blank?
+
+    tool_uses.map do |content|
+      {
+        "provider_tool_call_id" => content.tool_use.tool_use_id,
+        "name" => content.tool_use.name,
+        "arguments" => content.tool_use.input
+      }
+    end
+  end
+end

data/app/models/raif/concerns/llms/bedrock/tool_formatting.rb

@@ -34,4 +34,8 @@ module Raif::Concerns::Llms::Bedrock::ToolFormatting
       tools: tools.map{|tool| { tool_spec: tool } }
     }
   end
+
+  def build_forced_tool_choice(tool_name)
+    { tool: { name: tool_name } }
+  end
 end

data/app/models/raif/concerns/llms/google/message_formatting.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::Google::MessageFormatting
+  extend ActiveSupport::Concern
+
+  # Override the base format_messages to use Google's message format
+  def format_messages(messages)
+    messages.map do |message|
+      if message.is_a?(Hash) && message["type"] == "tool_call"
+        format_tool_call_message(message)
+      elsif message.is_a?(Hash) && message["type"] == "tool_call_result"
+        format_tool_call_result_message(message)
+      else
+        role = message["role"] || message[:role]
+        # Google uses "model" instead of "assistant"
+        google_role = role == "assistant" ? "model" : role
+        {
+          "role" => google_role,
+          "parts" => format_message_content(message["content"] || message[:content], role: role)
+        }
+      end
+    end
+  end
+
+  def format_string_message(content, role: nil)
+    { "text" => content }
+  end
+
+  def format_model_image_input_message(image_input)
+    if image_input.source_type == :url
+      {
+        "fileData" => {
+          "mimeType" => image_input.content_type,
+          "fileUri" => image_input.url
+        }
+      }
+    elsif image_input.source_type == :file_content
+      {
+        "inlineData" => {
+          "mimeType" => image_input.content_type,
+          "data" => image_input.base64_data
+        }
+      }
+    else
+      raise Raif::Errors::InvalidModelImageInputError, "Invalid model image input source type: #{image_input.source_type}"
+    end
+  end
+
+  def format_model_file_input_message(file_input)
+    if file_input.source_type == :url
+      {
+        "fileData" => {
+          "mimeType" => file_input.content_type,
+          "fileUri" => file_input.url
+        }
+      }
+    elsif file_input.source_type == :file_content
+      {
+        "inlineData" => {
+          "mimeType" => file_input.content_type,
+          "data" => file_input.base64_data
+        }
+      }
+    else
+      raise Raif::Errors::InvalidModelFileInputError, "Invalid model file input source type: #{file_input.source_type}"
+    end
+  end
+
+  def format_tool_call_message(tool_call)
+    parts = []
+
+    if tool_call["assistant_message"].present?
+      parts << format_string_message(tool_call["assistant_message"])
+    end
+
+    function_call_part = {
+      "functionCall" => {
+        "name" => tool_call["name"],
+        "args" => tool_call["arguments"]
+      }
+    }
+
+    # Include thoughtSignature if present (required for Gemini 2.5+ thinking models)
+    thought_signature = tool_call.dig("provider_metadata", "thought_signature")
+    function_call_part["thoughtSignature"] = thought_signature if thought_signature.present?
+
+    parts << function_call_part
+
+    {
+      "role" => "model",
+      "parts" => parts
+    }
+  end
+
+  def format_tool_call_result_message(tool_call_result)
+    result = tool_call_result["result"]
+    response_content = result.is_a?(String) ? { "output" => result } : result
+
+    {
+      "role" => "user",
+      "parts" => [{
+        "functionResponse" => {
+          "name" => tool_call_result["name"],
+          "response" => response_content
+        }
+      }]
+    }
+  end
+end

data/app/models/raif/concerns/llms/google/response_tool_calls.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::Google::ResponseToolCalls
+  extend ActiveSupport::Concern
+
+  def extract_response_tool_calls(resp)
+    parts = resp&.dig("candidates", 0, "content", "parts")
+    return if parts.blank?
+
+    # Find any functionCall parts
+    function_calls = parts.select { |part| part.key?("functionCall") }
+
+    return if function_calls.blank?
+
+    function_calls.map do |part|
+      function_call = part["functionCall"]
+      tool_call = {
+        # Google doesn't provide a unique ID for function calls, so we generate one
+        "provider_tool_call_id" => SecureRandom.uuid,
+        "name" => function_call["name"],
+        "arguments" => function_call["args"]
+      }
+
+      # Capture thoughtSignature if present (required for Gemini 2.5+ thinking models)
+      if part["thoughtSignature"].present?
+        tool_call["provider_metadata"] = { "thought_signature" => part["thoughtSignature"] }
+      end
+
+      tool_call
+    end
+  end
+end