raif 1.0.0 → 1.2.0

data/app/models/raif/llm.rb

@@ -3,13 +3,17 @@
 module Raif
   class Llm
     include ActiveModel::Model
+    include Raif::Concerns::Llms::MessageFormatting
 
     attr_accessor :key,
       :api_name,
       :default_temperature,
       :default_max_completion_tokens,
       :supports_native_tool_use,
-      :provider_settings
+      :provider_settings,
+      :input_token_cost,
+      :output_token_cost,
+      :supported_provider_managed_tools
 
     validates :key, presence: true
     validates :api_name, presence: true
@@ -18,13 +22,26 @@ module Raif
 
     alias_method :supports_native_tool_use?, :supports_native_tool_use
 
-    def initialize(key:, api_name:, model_provider_settings: {}, supports_native_tool_use: true, temperature: nil, max_completion_tokens: nil)
+    def initialize(
+      key:,
+      api_name:,
+      model_provider_settings: {},
+      supported_provider_managed_tools: [],
+      supports_native_tool_use: true,
+      temperature: nil,
+      max_completion_tokens: nil,
+      input_token_cost: nil,
+      output_token_cost: nil
+    )
       @key = key
       @api_name = api_name
       @provider_settings = model_provider_settings
       @supports_native_tool_use = supports_native_tool_use
       @default_temperature = temperature || 0.7
       @default_max_completion_tokens = max_completion_tokens
+      @input_token_cost = input_token_cost
+      @output_token_cost = output_token_cost
+      @supported_provider_managed_tools = supported_provider_managed_tools.map(&:to_s)
     end
 
     def name
@@ -32,7 +49,7 @@ module Raif
     end
 
     def chat(message: nil, messages: nil, response_format: :text, available_model_tools: [], source: nil, system_prompt: nil, temperature: nil,
-      max_completion_tokens: nil)
+      max_completion_tokens: 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
@@ -55,13 +72,13 @@ module Raif
         return
       end
 
-      messages = [{ role: "user", content: message }] if message.present?
+      messages = [{ "role" => "user", "content" => message }] if message.present?
 
       temperature ||= default_temperature
       max_completion_tokens ||= default_max_completion_tokens
 
       model_completion = Raif::ModelCompletion.new(
-        messages: messages,
+        messages: format_messages(messages),
         system_prompt: system_prompt,
         response_format: response_format,
         source: source,
@@ -69,20 +86,102 @@ module Raif
         model_api_name: api_name,
         temperature: temperature,
         max_completion_tokens: max_completion_tokens,
-        available_model_tools: available_model_tools
+        available_model_tools: available_model_tools,
+        stream_response: block_given?
       )
 
-      perform_model_completion!(model_completion)
+      retry_with_backoff(model_completion) do
+        perform_model_completion!(model_completion, &block)
+      end
+
       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}")
+      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)
+      raise e
     end
 
-    def perform_model_completion!(model_completion)
-      raise NotImplementedError, "Raif::Llm subclasses must implement #perform_model_completion!"
+    def perform_model_completion!(model_completion, &block)
+      raise NotImplementedError, "#{self.class.name} must implement #perform_model_completion!"
     end
 
     def self.valid_response_formats
       VALID_RESPONSE_FORMATS
     end
 
+    def supports_provider_managed_tool?(tool_klass)
+      supported_provider_managed_tools&.include?(tool_klass.to_s)
+    end
+
+    def validate_provider_managed_tool_support!(tool)
+      unless supports_provider_managed_tool?(tool)
+        raise Raif::Errors::UnsupportedFeatureError,
+          "Invalid provider-managed tool: #{tool.name} for #{key}"
+      end
+    end
+
+  private
+
+    def retry_with_backoff(model_completion)
+      retries = 0
+      max_retries = Raif.config.llm_request_max_retries
+      base_delay = 3
+      max_delay = 30
+
+      begin
+        yield
+      rescue *Raif.config.llm_request_retriable_exceptions => e
+        retries += 1
+        if retries <= max_retries
+          delay = [base_delay * (2**(retries - 1)), max_delay].min
+          Raif.logger.warn("Retrying LLM API request after error: #{e.message}. Attempt #{retries}/#{max_retries}. Waiting #{delay} seconds...")
+          model_completion.increment!(:retry_count)
+          sleep delay
+          retry
+        else
+          Raif.logger.error("LLM API request failed after #{max_retries} retries. Last error: #{e.message}")
+          raise
+        end
+      end
+    end
+
+    def streaming_response_type
+      raise NotImplementedError, "#{self.class.name} must implement #streaming_response_type"
+    end
+
+    def streaming_chunk_handler(model_completion, &block)
+      return unless model_completion.stream_response?
+
+      streaming_response = streaming_response_type.new
+      event_parser = EventStreamParser::Parser.new
+      accumulated_delta = ""
+
+      proc do |chunk, _size, _env|
+        event_parser.feed(chunk) do |event_type, data, _id, _reconnect_time|
+          if data.blank? || data == "[DONE]"
+            update_model_completion(model_completion, streaming_response.current_response_json)
+            next
+          end
+
+          event_data = JSON.parse(data)
+          delta, finish_reason = streaming_response.process_streaming_event(event_type, event_data)
+
+          accumulated_delta += delta if delta.present?
+
+          if accumulated_delta.length >= Raif.config.streaming_update_chunk_size_threshold || finish_reason.present?
+            update_model_completion(model_completion, streaming_response.current_response_json)
+
+            if accumulated_delta.present?
+              block.call(model_completion, accumulated_delta, event_data)
+              accumulated_delta = ""
+            end
+          end
+        end
+      end
+    end
+
   end
 end

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

@@ -1,47 +1,57 @@
 # frozen_string_literal: true
 
 class Raif::Llms::Anthropic < Raif::Llm
+  include Raif::Concerns::Llms::Anthropic::MessageFormatting
+  include Raif::Concerns::Llms::Anthropic::ToolFormatting
 
-  def perform_model_completion!(model_completion)
-    params = build_api_parameters(model_completion)
-
+  def perform_model_completion!(model_completion, &block)
+    params = build_request_parameters(model_completion)
     response = connection.post("messages") do |req|
-      req.body = params.to_json
+      req.body = params
+      req.options.on_data = streaming_chunk_handler(model_completion, &block) if model_completion.stream_response?
     end
 
-    resp = JSON.parse(response.body, symbolize_names: true)
-
-    # Handle API errors
-    unless response.success?
-      error_message = resp[:error]&.dig(:message) || "Anthropic API error: #{response.status}"
-      raise Raif::Errors::Anthropic::ApiError, error_message
+    unless model_completion.stream_response?
+      update_model_completion(model_completion, response.body)
     end
 
-    model_completion.raw_response = if model_completion.response_format_json?
-      extract_json_response(resp)
-    else
-      extract_text_response(resp)
-    end
-
-    model_completion.response_tool_calls = extract_response_tool_calls(resp)
-    model_completion.completion_tokens = resp&.dig(:usage, :output_tokens)
-    model_completion.prompt_tokens = resp&.dig(:usage, :input_tokens)
-    model_completion.save!
-
     model_completion
   end
 
+private
+
   def connection
     @connection ||= Faraday.new(url: "https://api.anthropic.com/v1") do |f|
-      f.headers["Content-Type"] = "application/json"
       f.headers["x-api-key"] = Raif.config.anthropic_api_key
       f.headers["anthropic-version"] = "2023-06-01"
+      f.request :json
+      f.response :json
+      f.response :raise_error
     end
   end
 
-protected
+  def streaming_response_type
+    Raif::StreamingResponses::Anthropic
+  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_id = response_json&.dig("id")
+    model_completion.response_array = response_json&.dig("content")
+    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("usage", "output_tokens")
+    model_completion.prompt_tokens = response_json&.dig("usage", "input_tokens")
+    model_completion.total_tokens = model_completion.completion_tokens.to_i + model_completion.prompt_tokens.to_i
+    model_completion.save!
+  end
 
-  def build_api_parameters(model_completion)
+  def build_request_parameters(model_completion)
     params = {
       model: model_completion.model_api_name,
       messages: model_completion.messages,
@@ -51,70 +61,75 @@ protected
 
     params[:system] = model_completion.system_prompt if model_completion.system_prompt.present?
 
-    # Add tools to the request if needed
-    tools = []
-
-    # If we're looking for a JSON response, add a tool to the request that the model can use to provide a JSON response
-    if model_completion.response_format_json? && model_completion.json_response_schema.present?
-      tools << {
-        name: "json_response",
-        description: "Generate a structured JSON response based on the provided schema.",
-        input_schema: model_completion.json_response_schema
-      }
+    if supports_native_tool_use?
+      tools = build_tools_parameter(model_completion)
+      params[:tools] = tools unless tools.blank?
     end
 
-    # 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|
-        tools << {
-          name: tool.tool_name,
-          description: tool.tool_description,
-          input_schema: tool.tool_arguments_schema
-        }
-      end
-    end
-
-    params[:tools] = tools if tools.any?
+    params[:stream] = true if model_completion.stream_response?
 
     params
   end
 
   def extract_text_response(resp)
-    resp&.dig(:content)&.first&.dig(:text)
+    return if resp&.dig("content").blank?
+
+    resp.dig("content").select{|v| v["type"] == "text" }.map{|v| v["text"] }.join("\n")
   end
 
   def extract_json_response(resp)
-    return extract_text_response(resp) if resp&.dig(:content).nil?
+    return extract_text_response(resp) if resp&.dig("content").nil?
 
     # Look for tool_use blocks in the content array
-    tool_name = "json_response"
-    tool_response = resp&.dig(:content)&.find do |content|
-      content[:type] == "tool_use" && content[:name] == tool_name
+    tool_response = resp&.dig("content")&.find do |content|
+      content["type"] == "tool_use" && content["name"] == "json_response"
     end
 
     if tool_response
-      JSON.generate(tool_response[:input])
+      JSON.generate(tool_response["input"])
     else
       extract_text_response(resp)
     end
   end
 
   def extract_response_tool_calls(resp)
-    return if resp&.dig(:content).nil?
+    return if resp&.dig("content").nil?
 
     # Find any tool_use content blocks
-    tool_uses = resp&.dig(:content)&.select do |content|
-      content[:type] == "tool_use"
+    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]
+        "name" => tool_use["name"],
+        "arguments" => tool_use["input"]
       }
     end
   end
 
+  def extract_citations(resp)
+    return [] if resp&.dig("content").nil?
+
+    citations = []
+
+    # Look through content blocks for citations
+    resp.dig("content").each do |content|
+      next unless content["type"] == "text" && content["citations"].present?
+
+      content["citations"].each do |citation|
+        next unless citation["type"] == "web_search_result_location"
+
+        citations << {
+          "url" => Raif::Utils::HtmlFragmentProcessor.strip_tracking_parameters(citation["url"]),
+          "title" => citation["title"]
+        }
+      end
+    end
+
+    citations.uniq{|citation| citation["url"] }
+  end
+
 end

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

@@ -0,0 +1,165 @@
+# frozen_string_literal: true
+
+class Raif::Llms::Bedrock < Raif::Llm
+  include Raif::Concerns::Llms::Bedrock::MessageFormatting
+  include Raif::Concerns::Llms::Bedrock::ToolFormatting
+
+  def perform_model_completion!(model_completion, &block)
+    if Raif.config.aws_bedrock_model_name_prefix.present?
+      model_completion.model_api_name = "#{Raif.config.aws_bedrock_model_name_prefix}.#{model_completion.model_api_name}"
+    end
+
+    params = build_request_parameters(model_completion)
+
+    if model_completion.stream_response?
+      bedrock_client.converse_stream(params) do |stream|
+        stream.on_error_event do |event|
+          raise Raif::Errors::StreamingError.new(
+            message: event.error_message,
+            type: event.event_type,
+            code: event.error_code,
+            event: event
+          )
+        end
+
+        handler = streaming_chunk_handler(model_completion, &block)
+        stream.on_event do |event|
+          handler.call(event)
+        end
+      end
+    else
+      response = bedrock_client.converse(params)
+      update_model_completion(model_completion, response)
+    end
+
+    model_completion
+  end
+
+private
+
+  def bedrock_client
+    @bedrock_client ||= Aws::BedrockRuntime::Client.new(region: Raif.config.aws_bedrock_region)
+  end
+
+  def update_model_completion(model_completion, resp)
+    model_completion.raw_response = if model_completion.response_format_json?
+      extract_json_response(resp)
+    else
+      extract_text_response(resp)
+    end
+
+    model_completion.response_array = resp.output.message.content
+    model_completion.response_tool_calls = extract_response_tool_calls(resp)
+    model_completion.completion_tokens = resp.usage.output_tokens
+    model_completion.prompt_tokens = resp.usage.input_tokens
+    model_completion.total_tokens = resp.usage.total_tokens
+    model_completion.save!
+  end
+
+  def build_request_parameters(model_completion)
+    # The AWS Bedrock SDK requires symbols for keys
+    messages_param = model_completion.messages.map(&:deep_symbolize_keys)
+    replace_tmp_base64_data_with_bytes(messages_param)
+
+    params = {
+      model_id: model_completion.model_api_name,
+      inference_config: { max_tokens: model_completion.max_completion_tokens || 8192 },
+      messages: messages_param
+    }
+
+    params[:system] = [{ text: model_completion.system_prompt }] if model_completion.system_prompt.present?
+
+    if supports_native_tool_use?
+      tools = build_tools_parameter(model_completion)
+      params[:tool_config] = tools unless tools.blank?
+    end
+
+    params
+  end
+
+  def replace_tmp_base64_data_with_bytes(messages)
+    # The AWS Bedrock SDK requires data sent as bytes (and doesn't support base64 like everyone else)
+    # The ModelCompletion stores the messages as JSON though, so it can't be raw bytes.
+    # We store the image data as base64, so we need to convert that to bytes before sending to AWS.
+    messages.each do |message|
+      message[:content].each do |content|
+        next unless content[:image] || content[:document]
+
+        type_key = content[:image] ? :image : :document
+        base64_data = content[type_key][:source].delete(:tmp_base64_data)
+        content[type_key][:source][:bytes] = Base64.strict_decode64(base64_data)
+      end
+    end
+  end
+
+  def extract_text_response(resp)
+    message = resp.output.message
+
+    # Find the first text content block
+    text_block = message.content&.find do |content|
+      content.respond_to?(:text) && content.text.present?
+    end
+
+    text_block&.text
+  end
+
+  def extract_json_response(resp)
+    # Get the message from the response object
+    message = resp.output.message
+
+    return extract_text_response(resp) if message.content.nil?
+
+    # Look for tool_use blocks in the content array
+    tool_response = message.content.find do |content|
+      content.respond_to?(:tool_use) && content.tool_use.present? && content.tool_use.name == "json_response"
+    end
+
+    if tool_response&.tool_use
+      JSON.generate(tool_response.tool_use.input)
+    else
+      extract_text_response(resp)
+    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?
+
+    streaming_response = Raif::StreamingResponses::Bedrock.new
+    accumulated_delta = ""
+
+    proc do |event|
+      delta, finish_reason = streaming_response.process_streaming_event(event.class, event)
+      accumulated_delta += delta if delta.present?
+
+      if accumulated_delta.length >= Raif.config.streaming_update_chunk_size_threshold || finish_reason.present?
+        update_model_completion(model_completion, streaming_response.current_response)
+
+        if accumulated_delta.present?
+          block.call(model_completion, accumulated_delta, event)
+          accumulated_delta = ""
+        end
+      end
+    end
+  end
+
+end

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

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+class Raif::Llms::OpenAiBase < Raif::Llm
+  include Raif::Concerns::Llms::OpenAi::JsonSchemaValidation
+
+  def perform_model_completion!(model_completion, &block)
+    if supports_temperature?
+      model_completion.temperature ||= default_temperature
+    else
+      Raif.logger.warn "Temperature is not supported for #{api_name}. Ignoring temperature parameter."
+      model_completion.temperature = nil
+    end
+
+    parameters = build_request_parameters(model_completion)
+
+    response = connection.post(api_path) do |req|
+      req.body = parameters
+      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://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
+    end
+  end
+
+  def format_system_prompt(model_completion)
+    formatted_system_prompt = model_completion.system_prompt.to_s.strip
+
+    # If the response format is JSON, we need to include "as json" in the system prompt.
+    # OpenAI requires this and will throw an error if it's not included.
+    if model_completion.response_format_json?
+      # Ensure system prompt ends with a period if not empty
+      if formatted_system_prompt.present? && !formatted_system_prompt.end_with?(".", "?", "!")
+        formatted_system_prompt += "."
+      end
+      formatted_system_prompt += " Return your response as JSON."
+      formatted_system_prompt.strip!
+    end
+
+    formatted_system_prompt
+  end
+
+  def supports_structured_outputs?
+    # Not all OpenAI models support structured outputs:
+    # https://platform.openai.com/docs/guides/structured-outputs?api-mode=chat#supported-models
+    provider_settings.key?(:supports_structured_outputs) ? provider_settings[:supports_structured_outputs] : true
+  end
+
+  def supports_temperature?
+    provider_settings.key?(:supports_temperature) ? provider_settings[:supports_temperature] : true
+  end
+
+end

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

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+class Raif::Llms::OpenAiCompletions < Raif::Llms::OpenAiBase
+  include Raif::Concerns::Llms::OpenAiCompletions::MessageFormatting
+  include Raif::Concerns::Llms::OpenAiCompletions::ToolFormatting
+
+private
+
+  def api_path
+    "chat/completions"
+  end
+
+  def streaming_response_type
+    Raif::StreamingResponses::OpenAiCompletions
+  end
+
+  def update_model_completion(model_completion, response_json)
+    model_completion.update!(
+      response_id: response_json["id"],
+      response_tool_calls: extract_response_tool_calls(response_json),
+      raw_response: response_json.dig("choices", 0, "message", "content"),
+      response_array: response_json["choices"],
+      completion_tokens: response_json.dig("usage", "completion_tokens"),
+      prompt_tokens: response_json.dig("usage", "prompt_tokens"),
+      total_tokens: response_json.dig("usage", "total_tokens")
+    )
+  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)
+
+    messages = model_completion.messages
+    messages_with_system = if formatted_system_prompt.blank?
+      messages
+    else
+      [{ "role" => "system", "content" => formatted_system_prompt }] + messages
+    end
+
+    parameters = {
+      model: api_name,
+      messages: messages_with_system
+    }
+
+    if supports_temperature?
+      parameters[:temperature] = model_completion.temperature.to_f
+    end
+
+    # If the LLM supports native tool use and there are available tools, add them to the parameters
+    if supports_native_tool_use?
+      tools = build_tools_parameter(model_completion)
+      parameters[:tools] = tools unless tools.blank?
+    end
+
+    if model_completion.stream_response?
+      parameters[:stream] = true
+      # Ask for usage stats in the last chunk
+      parameters[:stream_options] = { include_usage: true }
+    end
+
+    # Add response format if needed
+    response_format = determine_response_format(model_completion)
+    parameters[:response_format] = response_format if response_format
+    model_completion.response_format_parameter = response_format[:type] if response_format
+
+    parameters
+  end
+
+  def determine_response_format(model_completion)
+    # Only configure response format for JSON outputs
+    return unless model_completion.response_format_json?
+
+    if model_completion.json_response_schema.present? && supports_structured_outputs?
+      validate_json_schema!(model_completion.json_response_schema)
+
+      {
+        type: "json_schema",
+        json_schema: {
+          name: "json_response_schema",
+          strict: true,
+          schema: model_completion.json_response_schema
+        }
+      }
+    else
+      # Default JSON mode for OpenAI models that don't support structured outputs or no schema is provided
+      { type: "json_object" }
+    end
+  end
+
+end