raif 1.2.2 → 1.4.0

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
@@ -77,7 +85,7 @@ module Raif
       temperature ||= default_temperature
       max_completion_tokens ||= default_max_completion_tokens
 
-      model_completion = Raif::ModelCompletion.new(
+      model_completion = Raif::ModelCompletion.create!(
         messages: format_messages(messages),
         system_prompt: system_prompt,
         response_format: response_format,
@@ -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,

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

@@ -3,6 +3,7 @@
 class Raif::Llms::Anthropic < Raif::Llm
   include Raif::Concerns::Llms::Anthropic::MessageFormatting
   include Raif::Concerns::Llms::Anthropic::ToolFormatting
+  include Raif::Concerns::Llms::Anthropic::ResponseToolCalls
 
   def perform_model_completion!(model_completion, &block)
     params = build_request_parameters(model_completion)
@@ -21,7 +22,7 @@ class Raif::Llms::Anthropic < Raif::Llm
 private
 
   def connection
-    @connection ||= Faraday.new(url: "https://api.anthropic.com/v1") do |f|
+    @connection ||= Faraday.new(url: "https://api.anthropic.com/v1", request: Raif.default_request_options) do |f|
       f.headers["x-api-key"] = Raif.config.anthropic_api_key
       f.headers["anthropic-version"] = "2023-06-01"
       f.request :json
@@ -64,6 +65,11 @@ private
     if supports_native_tool_use?
       tools = build_tools_parameter(model_completion)
       params[:tools] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        params[:tool_choice] = build_forced_tool_choice(tool_klass.tool_name)
+      end
     end
 
     params[:stream] = true if model_completion.stream_response?
@@ -92,24 +98,6 @@ private
     end
   end
 
-  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|
-      {
-        "name" => tool_use["name"],
-        "arguments" => tool_use["input"]
-      }
-    end
-  end
-
   def extract_citations(resp)
     return [] if resp&.dig("content").nil?
 

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

@@ -3,6 +3,7 @@
 class Raif::Llms::Bedrock < Raif::Llm
   include Raif::Concerns::Llms::Bedrock::MessageFormatting
   include Raif::Concerns::Llms::Bedrock::ToolFormatting
+  include Raif::Concerns::Llms::Bedrock::ResponseToolCalls
 
   def perform_model_completion!(model_completion, &block)
     if Raif.config.aws_bedrock_model_name_prefix.present?
@@ -72,6 +73,11 @@ private
     if supports_native_tool_use?
       tools = build_tools_parameter(model_completion)
       params[:tool_config] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        params[:tool_config][:tool_choice] = build_forced_tool_choice(tool_klass.tool_name)
+      end
     end
 
     params
@@ -121,26 +127,6 @@ private
     end
   end
 
-  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|
-      {
-        "name" => content.tool_use.name,
-        "arguments" => content.tool_use.input
-      }
-    end
-  end
-
   def streaming_chunk_handler(model_completion, &block)
     return unless model_completion.stream_response?
 

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

@@ -0,0 +1,140 @@
+# frozen_string_literal: true
+
+class Raif::Llms::Google < Raif::Llm
+  include Raif::Concerns::Llms::Google::MessageFormatting
+  include Raif::Concerns::Llms::Google::ToolFormatting
+  include Raif::Concerns::Llms::Google::ResponseToolCalls
+
+  def perform_model_completion!(model_completion, &block)
+    params = build_request_parameters(model_completion)
+    endpoint = build_endpoint(model_completion)
+
+    response = connection.post(endpoint) do |req|
+      req.body = params
+      req.options.on_data = streaming_chunk_handler(model_completion, &block) if model_completion.stream_response?
+    end
+
+    unless model_completion.stream_response?
+      update_model_completion(model_completion, response.body)
+    end
+
+    model_completion
+  end
+
+private
+
+  def connection
+    @connection ||= Faraday.new(url: "https://generativelanguage.googleapis.com/v1beta") do |f|
+      f.headers["x-goog-api-key"] = Raif.config.google_api_key
+      f.request :json
+      f.response :json
+      f.response :raise_error
+    end
+  end
+
+  def build_endpoint(model_completion)
+    if model_completion.stream_response?
+      "models/#{model_completion.model_api_name}:streamGenerateContent?alt=sse"
+    else
+      "models/#{model_completion.model_api_name}:generateContent"
+    end
+  end
+
+  def streaming_response_type
+    Raif::StreamingResponses::Google
+  end
+
+  def update_model_completion(model_completion, response_json)
+    model_completion.raw_response = if model_completion.response_format_json?
+      extract_json_response(response_json)
+    else
+      extract_text_response(response_json)
+    end
+
+    model_completion.response_array = response_json&.dig("candidates", 0, "content", "parts")
+    model_completion.response_tool_calls = extract_response_tool_calls(response_json)
+    model_completion.citations = extract_citations(response_json)
+    model_completion.completion_tokens = response_json&.dig("usageMetadata", "candidatesTokenCount")
+    model_completion.prompt_tokens = response_json&.dig("usageMetadata", "promptTokenCount")
+    model_completion.total_tokens = response_json&.dig("usageMetadata", "totalTokenCount") ||
+      (model_completion.completion_tokens.to_i + model_completion.prompt_tokens.to_i)
+    model_completion.save!
+  end
+
+  def build_request_parameters(model_completion)
+    params = {
+      contents: model_completion.messages
+    }
+
+    if model_completion.system_prompt.present?
+      params[:system_instruction] = { parts: [{ text: model_completion.system_prompt }] }
+    end
+
+    params[:generationConfig] = build_generation_config(model_completion)
+
+    if supports_native_tool_use?
+      tools = build_tools_parameter(model_completion)
+      params[:tools] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        params[:toolConfig] = { functionCallingConfig: build_forced_tool_choice(tool_klass.tool_name) }
+      end
+    end
+
+    params
+  end
+
+  def build_generation_config(model_completion)
+    config = {}
+
+    temperature = model_completion.temperature || default_temperature
+    config[:temperature] = temperature.to_f if temperature.present?
+
+    max_tokens = model_completion.max_completion_tokens || default_max_completion_tokens
+    config[:maxOutputTokens] = max_tokens if max_tokens.present?
+
+    # Use native JSON schema support for structured output
+    if model_completion.response_format_json? && model_completion.json_response_schema.present?
+      config[:responseMimeType] = "application/json"
+      config[:responseSchema] = sanitize_schema_for_google(model_completion.json_response_schema)
+    end
+
+    config
+  end
+
+  def extract_text_response(resp)
+    parts = resp&.dig("candidates", 0, "content", "parts")
+    return if parts.blank?
+
+    parts.select { |p| p.key?("text") }.map { |p| p["text"] }.join
+  end
+
+  def extract_json_response(resp)
+    # Google AI supports native JSON schema output, so the response should be in the text field
+    extract_text_response(resp)
+  end
+
+  def extract_citations(resp)
+    # Google AI returns grounding metadata for search results
+    grounding_metadata = resp&.dig("candidates", 0, "groundingMetadata")
+    return [] if grounding_metadata.blank?
+
+    citations = []
+
+    # Extract from grounding chunks
+    grounding_chunks = grounding_metadata["groundingChunks"] || []
+    grounding_chunks.each do |chunk|
+      web = chunk["web"]
+      next unless web.present?
+
+      citations << {
+        "url" => Raif::Utils::HtmlFragmentProcessor.strip_tracking_parameters(web["uri"]),
+        "title" => web["title"]
+      }
+    end
+
+    citations.uniq { |citation| citation["url"] }
+  end
+
+end

data/app/models/raif/llms/open_ai_base.rb

@@ -28,11 +28,25 @@ class Raif::Llms::OpenAiBase < Raif::Llm
 private
 
   def connection
-    @connection ||= Faraday.new(url: "https://api.openai.com/v1") do |f|
-      f.headers["Authorization"] = "Bearer #{Raif.config.open_ai_api_key}"
-      f.request :json
-      f.response :json
-      f.response :raise_error
+    @connection ||= begin
+      conn = Faraday.new(url: Raif.config.open_ai_base_url, request: Raif.default_request_options) do |f|
+        case Raif.config.open_ai_auth_header_style
+        when :bearer
+          f.headers["Authorization"] = "Bearer #{Raif.config.open_ai_api_key}"
+        when :api_key
+          f.headers["api-key"] = Raif.config.open_ai_api_key
+        else
+          raise Raif::Errors::InvalidConfigError,
+            "Raif.config.open_ai_auth_header_style must be either :bearer or :api_key"
+        end
+
+        f.request :json
+        f.response :json
+        f.response :raise_error
+      end
+
+      conn.params["api-version"] = Raif.config.open_ai_api_version if Raif.config.open_ai_api_version.present?
+      conn
     end
   end
 

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

@@ -3,6 +3,7 @@
 class Raif::Llms::OpenAiCompletions < Raif::Llms::OpenAiBase
   include Raif::Concerns::Llms::OpenAiCompletions::MessageFormatting
   include Raif::Concerns::Llms::OpenAiCompletions::ToolFormatting
+  include Raif::Concerns::Llms::OpenAiCompletions::ResponseToolCalls
 
 private
 
@@ -26,17 +27,6 @@ private
     )
   end
 
-  def extract_response_tool_calls(resp)
-    return if resp.dig("choices", 0, "message", "tool_calls").blank?
-
-    resp.dig("choices", 0, "message", "tool_calls").map do |tool_call|
-      {
-        "name" => tool_call["function"]["name"],
-        "arguments" => JSON.parse(tool_call["function"]["arguments"])
-      }
-    end
-  end
-
   def build_request_parameters(model_completion)
     formatted_system_prompt = format_system_prompt(model_completion)
 
@@ -60,6 +50,11 @@ private
     if supports_native_tool_use?
       tools = build_tools_parameter(model_completion)
       parameters[:tools] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        parameters[:tool_choice] = build_forced_tool_choice(tool_klass.tool_name)
+      end
     end
 
     if model_completion.stream_response?

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

@@ -3,6 +3,7 @@
 class Raif::Llms::OpenAiResponses < Raif::Llms::OpenAiBase
   include Raif::Concerns::Llms::OpenAiResponses::MessageFormatting
   include Raif::Concerns::Llms::OpenAiResponses::ToolFormatting
+  include Raif::Concerns::Llms::OpenAiResponses::ResponseToolCalls
 
 private
 
@@ -27,22 +28,6 @@ private
     )
   end
 
-  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 << {
-        "name" => output_item["name"],
-        "arguments" => JSON.parse(output_item["arguments"])
-      }
-    end
-
-    tool_calls.any? ? tool_calls : nil
-  end
-
   def extract_raw_response(resp)
     text_outputs = []
 
@@ -110,6 +95,11 @@ private
     if supports_native_tool_use?
       tools = build_tools_parameter(model_completion)
       parameters[:tools] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        parameters[:tool_choice] = build_forced_tool_choice(tool_klass.tool_name)
+      end
     end
 
     # Add response format if needed. Default will be { "type": "text" }

data/app/models/raif/llms/open_router.rb

@@ -3,6 +3,7 @@
 class Raif::Llms::OpenRouter < Raif::Llm
   include Raif::Concerns::Llms::OpenAiCompletions::MessageFormatting
   include Raif::Concerns::Llms::OpenAiCompletions::ToolFormatting
+  include Raif::Concerns::Llms::OpenAiCompletions::ResponseToolCalls
   include Raif::Concerns::Llms::OpenAi::JsonSchemaValidation
 
   def perform_model_completion!(model_completion, &block)
@@ -23,7 +24,7 @@ class Raif::Llms::OpenRouter < Raif::Llm
 private
 
   def connection
-    @connection ||= Faraday.new(url: "https://openrouter.ai/api/v1") do |f|
+    @connection ||= Faraday.new(url: "https://openrouter.ai/api/v1", request: Raif.default_request_options) do |f|
       f.headers["Authorization"] = "Bearer #{Raif.config.open_router_api_key}"
       f.headers["HTTP-Referer"] = Raif.config.open_router_site_url if Raif.config.open_router_site_url.present?
       f.headers["X-Title"] = Raif.config.open_router_app_name if Raif.config.open_router_app_name.present?
@@ -85,6 +86,11 @@ private
       end
 
       params[:tools] = tools unless tools.blank?
+
+      if model_completion.tool_choice.present?
+        tool_klass = model_completion.tool_choice.constantize
+        params[:tool_choice] = build_forced_tool_choice(tool_klass.tool_name)
+      end
     end
 
     if model_completion.stream_response?
@@ -93,7 +99,9 @@ private
       params[:stream_options] = { include_usage: true }
     end
 
-    if model_completion.response_format_json?
+    # OpenRouter will sometimes complain about combining response_format json and tool calling.
+    # If we're telling it to use the json_response tool, then the json_object response_format should be irrelevant.
+    if model_completion.response_format_json? && params[:tools].blank?
       params[:response_format] = { type: "json_object" }
       model_completion.response_format_parameter = "json_object"
     end
@@ -119,16 +127,4 @@ private
       extract_text_response(resp)
     end
   end
-
-  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|
-      {
-        "name" => tool_call["function"]["name"],
-        "arguments" => JSON.parse(tool_call["function"]["arguments"])
-      }
-    end
-  end
 end

data/app/models/raif/model_completion.rb

@@ -1,14 +1,68 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_model_completions
+#
+#  id                        :bigint           not null, primary key
+#  available_model_tools     :jsonb            not null
+#  citations                 :jsonb
+#  completed_at              :datetime
+#  completion_tokens         :integer
+#  failed_at                 :datetime
+#  failure_error             :string
+#  failure_reason            :text
+#  llm_model_key             :string           not null
+#  max_completion_tokens     :integer
+#  messages                  :jsonb            not null
+#  model_api_name            :string           not null
+#  output_token_cost         :decimal(10, 6)
+#  prompt_token_cost         :decimal(10, 6)
+#  prompt_tokens             :integer
+#  raw_response              :text
+#  response_array            :jsonb
+#  response_format           :integer          default("text"), not null
+#  response_format_parameter :string
+#  response_tool_calls       :jsonb
+#  retry_count               :integer          default(0), not null
+#  source_type               :string
+#  started_at                :datetime
+#  stream_response           :boolean          default(FALSE), not null
+#  system_prompt             :text
+#  temperature               :decimal(5, 3)
+#  tool_choice               :string
+#  total_cost                :decimal(10, 6)
+#  total_tokens              :integer
+#  created_at                :datetime         not null
+#  updated_at                :datetime         not null
+#  response_id               :string
+#  source_id                 :bigint
+#
+# Indexes
+#
+#  index_raif_model_completions_on_completed_at  (completed_at)
+#  index_raif_model_completions_on_created_at    (created_at)
+#  index_raif_model_completions_on_failed_at     (failed_at)
+#  index_raif_model_completions_on_source        (source_type,source_id)
+#  index_raif_model_completions_on_started_at    (started_at)
+#
 class Raif::ModelCompletion < Raif::ApplicationRecord
   include Raif::Concerns::LlmResponseParsing
   include Raif::Concerns::HasAvailableModelTools
+  include Raif::Concerns::BooleanTimestamp
+
+  boolean_timestamp :started_at
+  boolean_timestamp :completed_at
+  boolean_timestamp :failed_at
 
   belongs_to :source, polymorphic: true, optional: true
 
   validates :llm_model_key, presence: true, inclusion: { in: ->{ Raif.available_llm_keys.map(&:to_s) } }
   validates :model_api_name, presence: true
 
+  # Scope to find completions that have response tool calls
+  scope :with_response_tool_calls, -> { where_json_not_blank(:response_tool_calls) }
+
   delegate :json_response_schema, to: :source, allow_nil: true
 
   before_save :set_total_tokens
@@ -41,6 +95,13 @@ class Raif::ModelCompletion < Raif::ApplicationRecord
     end
   end
 
+  def record_failure!(exception)
+    self.failed_at = Time.current
+    self.failure_error = exception.class.name
+    self.failure_reason = exception.message.truncate(255)
+    save!
+  end
+
 private
 
   def llm_config

data/app/models/raif/model_tool.rb

@@ -3,7 +3,7 @@
 class Raif::ModelTool
   include Raif::Concerns::JsonSchemaDefinition
 
-  delegate :tool_name, :tool_description, :tool_arguments_schema, :example_model_invocation, to: :class
+  delegate :tool_name, :tool_description, :example_model_invocation, to: :class
 
   class << self
     # The description of the tool that will be provided to the model
@@ -76,8 +76,9 @@ class Raif::ModelTool
       false
     end
 
-    def invoke_tool(tool_arguments:, source:)
+    def invoke_tool(provider_tool_call_id:, tool_arguments:, source:)
       tool_invocation = Raif::ModelToolInvocation.new(
+        provider_tool_call_id: provider_tool_call_id,
         source: source,
         tool_type: name,
         tool_arguments: tool_arguments
@@ -96,4 +97,11 @@ class Raif::ModelTool
     end
   end
 
+  # Instance method to get the tool arguments schema
+  # For instance-dependent schemas, builds the schema with this instance as context
+  # For class-level schemas, returns the class-level schema
+  def tool_arguments_schema
+    schema_for_instance(:tool_arguments)
+  end
+
 end

data/app/models/raif/model_tool_invocation.rb

@@ -1,5 +1,25 @@
 # frozen_string_literal: true
 
+# == Schema Information
+#
+# Table name: raif_model_tool_invocations
+#
+#  id                    :bigint           not null, primary key
+#  completed_at          :datetime
+#  failed_at             :datetime
+#  result                :jsonb            not null
+#  source_type           :string           not null
+#  tool_arguments        :jsonb            not null
+#  tool_type             :string           not null
+#  created_at            :datetime         not null
+#  updated_at            :datetime         not null
+#  provider_tool_call_id :string
+#  source_id             :bigint           not null
+#
+# Indexes
+#
+#  index_raif_model_tool_invocations_on_source  (source_type,source_id)
+#
 class Raif::ModelToolInvocation < Raif::ApplicationRecord
   belongs_to :source, polymorphic: true
 
@@ -22,14 +42,26 @@ class Raif::ModelToolInvocation < Raif::ApplicationRecord
     @tool ||= tool_type.constantize
   end
 
-  def as_llm_message
-    "Invoking tool: #{tool_name} with arguments: #{tool_arguments.to_json}"
+  # Returns tool call in the format expected by LLM message formatting
+  # @param assistant_message [String, nil] Optional assistant message accompanying the tool call
+  # @return [Hash] Hash representation for JSONB storage and LLM APIs
+  def as_tool_call_message(assistant_message: nil)
+    Raif::Messages::ToolCall.new(
+      provider_tool_call_id: provider_tool_call_id,
+      name: tool_name,
+      arguments: tool_arguments,
+      assistant_message: assistant_message
+    ).to_h
   end
 
-  def result_llm_message
-    return unless tool.respond_to?(:observation_for_invocation)
-
-    tool.observation_for_invocation(self)
+  # Returns tool result in the format expected by LLM message formatting
+  # @return [Hash] Hash representation for JSONB storage and LLM APIs
+  def as_tool_call_result_message
+    Raif::Messages::ToolCallResult.new(
+      provider_tool_call_id: provider_tool_call_id,
+      name: tool_name,
+      result: result
+    ).to_h
   end
 
   def to_partial_path