raif 1.3.0 → 1.4.0

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

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::Google::ToolFormatting
+  extend ActiveSupport::Concern
+
+  def build_tools_parameter(model_completion)
+    tools = []
+    function_declarations = []
+
+    # If we support native tool use and have tools available, add them to the request
+    if supports_native_tool_use? && model_completion.available_model_tools.any?
+      model_completion.available_model_tools_map.each do |_tool_name, tool|
+        if tool.provider_managed?
+          # Provider-managed tools are added as separate tool entries
+          tools << format_provider_managed_tool(tool)
+        else
+          function_declarations << {
+            name: tool.tool_name,
+            description: tool.tool_description,
+            parameters: sanitize_schema_for_google(tool.tool_arguments_schema)
+          }
+        end
+      end
+    end
+
+    # Add function declarations if any
+    if function_declarations.any?
+      tools << { functionDeclarations: function_declarations }
+    end
+
+    tools
+  end
+
+  def format_provider_managed_tool(tool)
+    validate_provider_managed_tool_support!(tool)
+
+    case tool.name
+    when "Raif::ModelTools::ProviderManaged::WebSearch"
+      { google_search: {} }
+    when "Raif::ModelTools::ProviderManaged::CodeExecution"
+      { code_execution: {} }
+    else
+      raise Raif::Errors::UnsupportedFeatureError,
+        "Invalid provider-managed tool: #{tool.name} for #{key}"
+    end
+  end
+
+  def build_forced_tool_choice(tool_name)
+    { mode: "ANY", allowedFunctionNames: [tool_name] }
+  end
+
+private
+
+  # Google's API doesn't support additionalProperties in JSON schemas
+  # This method recursively removes it from the schema
+  def sanitize_schema_for_google(schema)
+    return schema unless schema.is_a?(Hash)
+
+    sanitized = schema.except(:additionalProperties, "additionalProperties")
+
+    sanitized.transform_values do |value|
+      case value
+      when Hash
+        sanitize_schema_for_google(value)
+      when Array
+        value.map { |item| sanitize_schema_for_google(item) }
+      else
+        value
+      end
+    end
+  end
+end

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

@@ -5,11 +5,17 @@ module Raif::Concerns::Llms::MessageFormatting
 
   def format_messages(messages)
     messages.map do |message|
-      role = message["role"] || message[:role]
-      {
-        "role" => role,
-        "content" => format_message_content(message["content"] || message[:content], role: role)
-      }
+      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]
+        {
+          "role" => role,
+          "content" => format_message_content(message["content"] || message[:content], role: role)
+        }
+      end
     end
   end
 

data/app/models/raif/concerns/llms/open_ai/json_schema_validation.rb

@@ -23,7 +23,7 @@ module Raif::Concerns::Llms::OpenAi::JsonSchemaValidation
       # Check properties count (max 100 total)
       validate_properties_count(schema, errors)
 
-      # Check nesting depth (max 5 levels)
+      # Check nesting depth (max 10 levels)
       validate_nesting_depth(schema, errors)
 
       # Check for unsupported anyOf at root level
@@ -118,8 +118,8 @@ private
   def validate_nesting_depth(schema, errors, depth = 1)
     return unless schema.is_a?(Hash)
 
-    if depth > 5
-      errors << "Schema exceeds maximum nesting depth of 5 levels"
+    if depth > 10
+      errors << "Schema exceeds maximum nesting depth of 10 levels"
       return
     end
 

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

@@ -38,4 +38,26 @@ module Raif::Concerns::Llms::OpenAiCompletions::MessageFormatting
       raise Raif::Errors::InvalidModelFileInputError, "Invalid model image input source type: #{file_input.source_type}"
     end
   end
+
+  def format_tool_call_message(tool_call)
+    {
+      "role" => "assistant",
+      "tool_calls" => [{
+        "id" => tool_call["provider_tool_call_id"],
+        "type" => "function",
+        "function" => {
+          "name" => tool_call["name"],
+          "arguments" => JSON.generate(tool_call["arguments"])
+        }
+      }]
+    }
+  end
+
+  def format_tool_call_result_message(tool_call_result)
+    {
+      "role" => "tool",
+      "tool_call_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/open_ai_completions/response_tool_calls.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::OpenAiCompletions::ResponseToolCalls
+  extend ActiveSupport::Concern
+
+  def extract_response_tool_calls(resp)
+    tool_calls = resp.dig("choices", 0, "message", "tool_calls")
+    return if tool_calls.blank?
+
+    tool_calls.map do |tool_call|
+      {
+        "provider_tool_call_id" => tool_call["id"],
+        "name" => tool_call["function"]["name"],
+        "arguments" => begin
+          JSON.parse(tool_call["function"]["arguments"])
+        rescue JSON::ParserError
+          tool_call["function"]["arguments"]
+        end
+      }
+    end
+  end
+end

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

@@ -23,4 +23,8 @@ module Raif::Concerns::Llms::OpenAiCompletions::ToolFormatting
       end
     end
   end
+
+  def build_forced_tool_choice(tool_name)
+    { "type" => "function", "function" => { "name" => tool_name } }
+  end
 end

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

@@ -40,4 +40,21 @@ module Raif::Concerns::Llms::OpenAiResponses::MessageFormatting
       raise Raif::Errors::InvalidModelFileInputError, "Invalid model image input source type: #{file_input.source_type}"
     end
   end
+
+  def format_tool_call_message(tool_call)
+    {
+      "type" => "function_call",
+      "call_id" => tool_call["provider_tool_call_id"],
+      "name" => tool_call["name"],
+      "arguments" => JSON.generate(tool_call["arguments"])
+    }
+  end
+
+  def format_tool_call_result_message(tool_call_result)
+    {
+      "type" => "function_call_output",
+      "call_id" => tool_call_result["provider_tool_call_id"],
+      "output" => 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/open_ai_responses/response_tool_calls.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::Llms::OpenAiResponses::ResponseToolCalls
+  extend ActiveSupport::Concern
+
+  def extract_response_tool_calls(resp)
+    return if resp["output"].blank?
+
+    tool_calls = []
+    resp["output"].each do |output_item|
+      next unless output_item["type"] == "function_call"
+
+      tool_calls << {
+        "provider_tool_call_id" => output_item["call_id"],
+        "name" => output_item["name"],
+        "arguments" => begin
+          JSON.parse(output_item["arguments"])
+        rescue JSON::ParserError
+          output_item["arguments"]
+        end
+      }
+    end
+
+    tool_calls.any? ? tool_calls : nil
+  end
+end

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

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

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

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+module Raif::Concerns::RunWith
+  extend ActiveSupport::Concern
+
+  included do
+    class_attribute :_run_with_args, instance_writer: false, default: []
+
+    # Backward compatibility alias
+    class_attribute :_task_run_args, instance_writer: false, default: []
+
+    # Automatically serialize run_with args before validation on create
+    before_validation :serialize_run_with_to_column, on: :create
+  end
+
+  class_methods do
+    # Scope for querying records by run_with arguments
+    # @param args [Hash] Key-value pairs to match in the run_with column
+    # @example
+    #   Task.having_run_with(document: doc)
+    #   Task.having_run_with(user: user, options: { foo: "bar" })
+    def having_run_with(**args)
+      return all if args.empty?
+
+      # Serialize args the same way we do for storage (handles GID conversion)
+      serialized = serialize_run_with(args)
+
+      # Avoid matching all records if args didn't match declared run_with arguments
+      return none if args.any? && serialized.empty?
+
+      # Use database-specific JSON containment query
+      case connection.adapter_name.downcase
+      when "postgresql"
+        # PostgreSQL: Use JSONB containment operator
+        where("run_with @> ?", serialized.to_json)
+      when "mysql2", "trilogy"
+        # MySQL: Use JSON_CONTAINS function
+        where("JSON_CONTAINS(run_with, ?)", serialized.to_json)
+      else
+        raise "Unsupported database: #{connection.adapter_name}"
+      end
+    end
+
+    # DSL for declaring persistent run arguments that will be serialized to the database
+    # @param name [Symbol] The name of the argument
+    def run_with(name)
+      # Ensure each class has its own array copy
+      self._run_with_args = _run_with_args.dup
+      _run_with_args << name.to_sym
+
+      # Keep backward compatibility for _task_run_args class attribute
+      self._task_run_args = _task_run_args.dup
+      _task_run_args << name.to_sym
+
+      # Define getter that pulls from run_with JSON column
+      define_method(name) do
+        return instance_variable_get("@#{name}") if instance_variable_defined?("@#{name}")
+
+        value = run_with&.dig(name.to_s)
+        return unless value
+
+        # Deserialize GID if it's a string starting with gid://
+        deserialized = if value.is_a?(String) && value.start_with?("gid://")
+          begin
+            GlobalID::Locator.locate(value)
+          rescue ActiveRecord::RecordNotFound
+            nil
+          end
+        else
+          value
+        end
+
+        instance_variable_set("@#{name}", deserialized)
+      end
+
+      # Define setter that stores in memory (for use during run)
+      define_method("#{name}=") do |value|
+        instance_variable_set("@#{name}", value)
+      end
+    end
+
+    # Backward compatibility alias
+    alias_method :task_run_arg, :run_with
+
+    # Transform run args into a hash that can be stored in the run_with database column
+    def serialize_run_with(args)
+      serialized_args = {}
+      _run_with_args.each do |arg_name|
+        next unless args.key?(arg_name)
+
+        value = args[arg_name]
+        serialized_args[arg_name.to_s] = if value.respond_to?(:to_global_id)
+          value.to_global_id.to_s
+        else
+          value
+        end
+      end
+
+      serialized_args
+    end
+
+    # Backward compatibility alias
+    alias_method :serialize_task_run_args, :serialize_run_with
+  end
+
+private
+
+  # Automatically called before validation on create to serialize run_with args
+  # Collects all declared run_with arguments from instance variables and serializes them
+  # to the run_with JSON column
+  def serialize_run_with_to_column
+    args = {}
+
+    # Collect all run_with args that were set via instance variables
+    self.class._run_with_args.each do |arg_name|
+      if instance_variable_defined?("@#{arg_name}")
+        args[arg_name] = instance_variable_get("@#{arg_name}")
+      end
+    end
+
+    # Merge serialized args into run_with hash if any args were set
+    if args.any?
+      self.run_with ||= {}
+      self.run_with = self.run_with.merge(self.class.serialize_run_with(args))
+    end
+  end
+end

data/app/models/raif/conversation.rb

@@ -1,5 +1,33 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_conversations
+#
+#  id                         :bigint           not null, primary key
+#  available_model_tools      :jsonb            not null
+#  available_user_tools       :jsonb            not null
+#  conversation_entries_count :integer          default(0), not null
+#  creator_type               :string           not null
+#  generating_entry_response  :boolean          default(FALSE), not null
+#  llm_messages_max_length    :integer
+#  llm_model_key              :string           not null
+#  requested_language_key     :string
+#  response_format            :integer          default("text"), not null
+#  source_type                :string
+#  system_prompt              :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_conversations_on_created_at  (created_at)
+#  index_raif_conversations_on_creator     (creator_type,creator_id)
+#  index_raif_conversations_on_source      (source_type,source_id)
+#
 class Raif::Conversation < Raif::ApplicationRecord
   include Raif::Concerns::HasLlm
   include Raif::Concerns::HasRequestedLanguage
@@ -7,6 +35,27 @@ class Raif::Conversation < Raif::ApplicationRecord
   include Raif::Concerns::LlmResponseParsing
 
   belongs_to :creator, polymorphic: true
+  belongs_to :source, polymorphic: true, optional: true
+
+  class << self
+    def before_prompt_model_for_entry_response(&block)
+      @before_prompt_model_for_entry_response_blocks ||= []
+      @before_prompt_model_for_entry_response_blocks << block if block
+    end
+
+    def before_prompt_model_for_entry_response_blocks
+      blocks = []
+
+      # Collect blocks from ancestors (in reverse order so parent blocks run first)
+      ancestors.reverse_each do |klass|
+        if klass.instance_variable_defined?(:@before_prompt_model_for_entry_response_blocks)
+          blocks.concat(klass.instance_variable_get(:@before_prompt_model_for_entry_response_blocks))
+        end
+      end
+
+      blocks
+    end
+  end
 
   has_many :entries, class_name: "Raif::ConversationEntry", dependent: :destroy, foreign_key: :raif_conversation_id, inverse_of: :raif_conversation
 
@@ -14,6 +63,7 @@ class Raif::Conversation < Raif::ApplicationRecord
 
   after_initialize -> { self.available_model_tools ||= [] }
   after_initialize -> { self.available_user_tools ||= [] }
+  after_initialize -> { self.llm_messages_max_length ||= Raif.config.conversation_llm_messages_max_length_default }
 
   before_validation ->{ self.type ||= "Raif::Conversation" }, on: :create
 
@@ -39,9 +89,15 @@ class Raif::Conversation < Raif::ApplicationRecord
   end
 
   def prompt_model_for_entry_response(entry:, &block)
-    update(system_prompt: build_system_prompt)
+    self.class.before_prompt_model_for_entry_response_blocks.each do |callback_block|
+      instance_exec(entry, &callback_block)
+    end
+
+    self.system_prompt = build_system_prompt
+    self.generating_entry_response = true
+    save!
 
-    llm.chat(
+    model_completion = llm.chat(
       messages: llm_messages,
       source: entry,
       response_format: response_format.to_sym,
@@ -49,7 +105,15 @@ class Raif::Conversation < Raif::ApplicationRecord
       available_model_tools: available_model_tools,
       &block
     )
+
+    self.generating_entry_response = false
+    save!
+
+    model_completion
   rescue StandardError => e
+    self.generating_entry_response = false
+    save!
+
     Rails.logger.error("Error processing conversation entry ##{entry.id}. #{e.message}")
     Rails.logger.error(e.backtrace.join("\n"))
 
@@ -75,14 +139,33 @@ class Raif::Conversation < Raif::ApplicationRecord
   def llm_messages
     messages = []
 
-    entries.oldest_first.includes(:raif_model_tool_invocations).each do |entry|
-      messages << { "role" => "user", "content" => entry.user_message } unless entry.user_message.blank?
+    # Apply max length limit to entries if configured (nil means no limit)
+    included_entries = entries.oldest_first.includes(:raif_model_tool_invocations)
+    included_entries = included_entries.last(llm_messages_max_length) if llm_messages_max_length.present?
+
+    included_entries.each do |entry|
+      unless entry.user_message.blank?
+        messages << Raif::Messages::UserMessage.new(content: entry.user_message).to_h
+      end
+
       next unless entry.completed?
 
-      messages << { "role" => "assistant", "content" => entry.model_response_message } unless entry.model_response_message.blank?
-      entry.raif_model_tool_invocations.each do |tool_invocation|
-        messages << { "role" => "assistant", "content" => tool_invocation.as_llm_message }
-        messages << { "role" => "assistant", "content" => tool_invocation.result_llm_message } if tool_invocation.result_llm_message.present?
+      tool_invocations = entry.raif_model_tool_invocations.to_a
+
+      if tool_invocations.any?
+        # First tool call includes the assistant's message (if any)
+        first_invocation = tool_invocations.shift
+        messages << first_invocation.as_tool_call_message(assistant_message: entry.model_response_message.presence)
+        messages << first_invocation.as_tool_call_result_message
+
+        # Remaining tool calls (if multiple)
+        tool_invocations.each do |tool_invocation|
+          messages << tool_invocation.as_tool_call_message
+          messages << tool_invocation.as_tool_call_result_message
+        end
+      elsif entry.model_response_message.present?
+        # No tool calls, just a regular assistant response
+        messages << Raif::Messages::AssistantMessage.new(content: entry.model_response_message).to_h
       end
     end
 

data/app/models/raif/conversation_entry.rb

@@ -1,5 +1,32 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_conversation_entries
+#
+#  id                     :bigint           not null, primary key
+#  completed_at           :datetime
+#  creator_type           :string           not null
+#  failed_at              :datetime
+#  model_response_message :text
+#  raw_response           :text
+#  started_at             :datetime
+#  user_message           :text
+#  created_at             :datetime         not null
+#  updated_at             :datetime         not null
+#  creator_id             :bigint           not null
+#  raif_conversation_id   :bigint           not null
+#
+# Indexes
+#
+#  index_raif_conversation_entries_on_created_at            (created_at)
+#  index_raif_conversation_entries_on_creator               (creator_type,creator_id)
+#  index_raif_conversation_entries_on_raif_conversation_id  (raif_conversation_id)
+#
+# Foreign Keys
+#
+#  fk_rails_...  (raif_conversation_id => raif_conversations.id)
+#
 class Raif::ConversationEntry < Raif::ApplicationRecord
   include Raif::Concerns::InvokesModelTools
   include Raif::Concerns::HasAvailableModelTools
@@ -98,7 +125,11 @@ private
         tool_klass = available_model_tools_map[tool_call["name"]]
         next if tool_klass.nil?
 
-        tool_klass.invoke_tool(tool_arguments: tool_call["arguments"], source: self)
+        tool_klass.invoke_tool(
+          provider_tool_call_id: tool_call["provider_tool_call_id"],
+          tool_arguments: tool_call["arguments"],
+          source: self
+        )
       end
 
       completed!

data/app/models/raif/embedding_model.rb

@@ -5,6 +5,7 @@ class Raif::EmbeddingModel
 
   attr_accessor :key,
     :api_name,
+    :display_name,
     :input_token_cost,
     :default_output_vector_size
 
@@ -13,7 +14,7 @@ class Raif::EmbeddingModel
   validates :key, presence: true
 
   def name
-    I18n.t("raif.embedding_model_names.#{key}")
+    I18n.t("raif.embedding_model_names.#{key}", default: display_name || key.to_s.humanize)
   end
 
   def generate_embedding!(input, dimensions: nil)

data/app/models/raif/embedding_models/open_ai.rb

@@ -30,7 +30,7 @@ private
   end
 
   def connection
-    @connection ||= Faraday.new(url: "https://api.openai.com/v1") do |f|
+    @connection ||= Faraday.new(url: Raif.config.open_ai_embedding_base_url, request: Raif.default_request_options) do |f|
       f.headers["Authorization"] = "Bearer #{Raif.config.open_ai_api_key}"
       f.request :json
       f.response :json

data/app/models/raif/llm.rb

@@ -7,6 +7,7 @@ module Raif
 
     attr_accessor :key,
       :api_name,
+      :display_name,
       :default_temperature,
       :default_max_completion_tokens,
       :supports_native_tool_use,
@@ -25,6 +26,7 @@ module Raif
     def initialize(
       key:,
       api_name:,
+      display_name: nil,
       model_provider_settings: {},
       supported_provider_managed_tools: [],
       supports_native_tool_use: true,
@@ -35,6 +37,7 @@ module Raif
     )
       @key = key
       @api_name = api_name
+      @display_name = display_name
       @provider_settings = model_provider_settings
       @supports_native_tool_use = supports_native_tool_use
       @default_temperature = temperature || 0.7
@@ -45,11 +48,11 @@ module Raif
     end
 
     def name
-      I18n.t("raif.model_names.#{key}")
+      I18n.t("raif.model_names.#{key}", default: display_name || key.to_s.humanize)
     end
 
     def chat(message: nil, messages: nil, response_format: :text, available_model_tools: [], source: nil, system_prompt: nil, temperature: nil,
-      max_completion_tokens: nil, &block)
+      max_completion_tokens: nil, tool_choice: nil, &block)
       unless response_format.is_a?(Symbol)
         raise ArgumentError,
           "Raif::Llm#chat - Invalid response format: #{response_format}. Must be a symbol (you passed #{response_format.class}) and be one of: #{VALID_RESPONSE_FORMATS.join(", ")}" # rubocop:disable Layout/LineLength
@@ -67,6 +70,11 @@ module Raif
         raise ArgumentError, "Raif::Llm#chat - You must provide either a message: or messages: argument, not both"
       end
 
+      if tool_choice.present? && !available_model_tools.map(&:to_s).include?(tool_choice.to_s)
+        raise ArgumentError,
+          "Raif::Llm#chat - Invalid tool choice: #{tool_choice} is not included in the available model tools: #{available_model_tools.join(", ")}"
+      end
+
       unless Raif.config.llm_api_requests_enabled
         Raif.logger.warn("LLM API requests are disabled. Skipping request to #{api_name}.")
         return
@@ -87,20 +95,29 @@ module Raif
         temperature: temperature,
         max_completion_tokens: max_completion_tokens,
         available_model_tools: available_model_tools,
+        tool_choice: tool_choice&.to_s,
         stream_response: block_given?
       )
 
+      model_completion.started!
+
       retry_with_backoff(model_completion) do
         perform_model_completion!(model_completion, &block)
       end
 
+      model_completion.completed!
       model_completion
     rescue Raif::Errors::StreamingError => e
       Rails.logger.error("Raif streaming error -- code: #{e.code} -- type: #{e.type} -- message: #{e.message} -- event: #{e.event}")
+      model_completion&.record_failure!(e)
       raise e
     rescue Faraday::Error => e
       Raif.logger.error("LLM API request failed (status: #{e.response_status}): #{e.message}")
       Raif.logger.error(e.response_body)
+      model_completion&.record_failure!(e)
+      raise e
+    rescue StandardError => e
+      model_completion&.record_failure!(e)
       raise e
     end
 
@@ -116,6 +133,14 @@ module Raif
       supported_provider_managed_tools&.include?(tool_klass.to_s)
     end
 
+    # Build the tool_choice parameter to force a specific tool to be called.
+    # Each provider implements this to return the correct format.
+    # @param tool_name [String] The name of the tool to force
+    # @return [Hash] The tool_choice parameter for the provider's API
+    def build_forced_tool_choice(tool_name)
+      raise NotImplementedError, "#{self.class.name} must implement #build_forced_tool_choice"
+    end
+
     def validate_provider_managed_tool_support!(tool)
       unless supports_provider_managed_tool?(tool)
         raise Raif::Errors::UnsupportedFeatureError,