raif 1.0.0 → 1.2.0

data/app/models/raif/model_tools/fetch_url.rb

@@ -1,57 +1,48 @@
 # frozen_string_literal: true
 
 class Raif::ModelTools::FetchUrl < Raif::ModelTool
+  tool_arguments_schema do
+    string "url", description: "The URL to fetch content from"
+  end
 
-  def self.example_model_invocation
+  example_model_invocation do
     {
       "name": tool_name,
       "arguments": { "url": "https://en.wikipedia.org/wiki/NASA" }
     }
   end
 
-  def self.tool_arguments_schema
-    {
-      type: "object",
-      additionalProperties: false,
-      required: ["url"],
-      properties: {
-        url: {
-          type: "string",
-          description: "The URL to fetch content from"
-        }
-      }
-    }
-  end
-
-  def self.tool_description
+  tool_description do
     "Fetch a URL and return the page content as markdown"
   end
 
-  def self.observation_for_invocation(tool_invocation)
-    return "No results found" unless tool_invocation.result.present?
+  class << self
+    def observation_for_invocation(tool_invocation)
+      return "No results found" unless tool_invocation.result.present?
 
-    <<~OBSERVATION
-      Result Status: #{tool_invocation.result["status"]}
-      Result Content:
-      #{tool_invocation.result["content"]}
-    OBSERVATION
-  end
+      <<~OBSERVATION
+        Result Status: #{tool_invocation.result["status"]}
+        Result Content:
+        #{tool_invocation.result["content"]}
+      OBSERVATION
+    end
 
-  def self.process_invocation(tool_invocation)
-    url = tool_invocation.tool_arguments["url"]
-    response = Faraday.get(url)
+    def process_invocation(tool_invocation)
+      url = tool_invocation.tool_arguments["url"]
+      response = Faraday.get(url)
 
-    readable_content = Raif::Utils::ReadableContentExtractor.new(response.body).extract_readable_content
-    markdown_content = Raif::Utils::HtmlToMarkdownConverter.convert(readable_content)
+      readable_content = Raif::Utils::ReadableContentExtractor.new(response.body).extract_readable_content
+      markdown_content = Raif::Utils::HtmlToMarkdownConverter.convert(readable_content)
 
-    tool_invocation.update!(
-      result: {
-        status: response.status,
-        content: markdown_content
-      }
-    )
+      tool_invocation.update!(
+        result: {
+          status: response.status,
+          content: markdown_content
+        }
+      )
 
-    tool_invocation.result
+      tool_invocation.result
+    end
   end
 
 end

data/app/models/raif/model_tools/provider_managed/base.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class Raif::ModelTools::ProviderManaged::Base < Raif::ModelTool
+  class << self
+    def provider_managed?
+      true
+    end
+  end
+end

data/app/models/raif/model_tools/provider_managed/code_execution.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Raif::ModelTools::ProviderManaged::CodeExecution < Raif::ModelTools::ProviderManaged::Base
+
+end

data/app/models/raif/model_tools/provider_managed/image_generation.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Raif::ModelTools::ProviderManaged::ImageGeneration < Raif::ModelTools::ProviderManaged::Base
+
+end

data/app/models/raif/model_tools/provider_managed/web_search.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+class Raif::ModelTools::ProviderManaged::WebSearch < Raif::ModelTools::ProviderManaged::Base
+
+end

data/app/models/raif/model_tools/wikipedia_search.rb

@@ -1,78 +1,69 @@
 # frozen_string_literal: true
 
 class Raif::ModelTools::WikipediaSearch < Raif::ModelTool
+  tool_arguments_schema do
+    string "query", description: "The query to search Wikipedia for"
+  end
 
-  def self.example_model_invocation
+  example_model_invocation do
     {
       "name" => tool_name,
       "arguments" => { "query": "Jimmy Buffett" }
     }
   end
 
-  def self.tool_arguments_schema
-    {
-      type: "object",
-      additionalProperties: false,
-      required: ["query"],
-      properties: {
-        query: {
-          type: "string",
-          description: "The query to search Wikipedia for"
-        }
-      }
-    }
-  end
-
-  def self.tool_description
+  tool_description do
     "Search Wikipedia for information"
   end
 
-  def self.observation_for_invocation(tool_invocation)
-    return "No results found" unless tool_invocation.result.present?
+  class << self
+    def observation_for_invocation(tool_invocation)
+      return "No results found" unless tool_invocation.result.present?
 
-    JSON.pretty_generate(tool_invocation.result)
-  end
+      JSON.pretty_generate(tool_invocation.result)
+    end
 
-  def self.process_invocation(tool_invocation)
-    query = tool_invocation.tool_arguments["query"]
+    def process_invocation(tool_invocation)
+      query = tool_invocation.tool_arguments["query"]
 
-    conn = Faraday.new(url: "https://en.wikipedia.org/w/api.php")
+      conn = Faraday.new(url: "https://en.wikipedia.org/w/api.php")
 
-    response = conn.get do |req|
-      req.params["action"] = "query"
-      req.params["format"] = "json"
-      req.params["list"] = "search"
-      req.params["srsearch"] = query
-      req.params["srlimit"] = 5 # Limit to 5 results
-      req.params["srprop"] = "snippet"
-    end
+      response = conn.get do |req|
+        req.params["action"] = "query"
+        req.params["format"] = "json"
+        req.params["list"] = "search"
+        req.params["srsearch"] = query
+        req.params["srlimit"] = 5 # Limit to 5 results
+        req.params["srprop"] = "snippet"
+      end
 
-    if response.success?
-      results = JSON.parse(response.body)
-      search_results = results.dig("query", "search") || []
+      if response.success?
+        results = JSON.parse(response.body)
+        search_results = results.dig("query", "search") || []
 
-      # Store the results in the tool_invocation
-      tool_invocation.update!(
-        result: {
-          results: search_results.map do |result|
-            {
-              title: result["title"],
-              snippet: result["snippet"],
-              page_id: result["pageid"],
-              url: "https://en.wikipedia.org/wiki/#{result["title"].gsub(" ", "_")}"
-            }
-          end
-        }
-      )
-    else
-      tool_invocation.update!(
-        result: {
-          error: "Failed to fetch results from Wikipedia API: #{response.status} #{response.reason_phrase}"
-        }
-      )
-    end
+        # Store the results in the tool_invocation
+        tool_invocation.update!(
+          result: {
+            results: search_results.map do |result|
+              {
+                title: result["title"],
+                snippet: result["snippet"],
+                page_id: result["pageid"],
+                url: "https://en.wikipedia.org/wiki/#{result["title"].gsub(" ", "_")}"
+              }
+            end
+          }
+        )
+      else
+        tool_invocation.update!(
+          result: {
+            error: "Failed to fetch results from Wikipedia API: #{response.status} #{response.reason_phrase}"
+          }
+        )
+      end
 
-    tool_invocation.result
+      tool_invocation.result
+    end
   end
 
 end

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

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+class Raif::StreamingResponses::Anthropic
+
+  def initialize
+    @response_json = { "content" => [], "usage" => {} }
+    @finish_reason = nil
+  end
+
+  def process_streaming_event(event_type, event)
+    delta = nil
+    index = event["index"]
+
+    case event_type
+    when "message_start"
+      @response_json = event["message"]
+      @response_json["content"] = []
+      @response_json["usage"] ||= {}
+    when "content_block_start"
+      @response_json["content"][index] = event["content_block"]
+      if event.dig("content_block", "type") == "tool_use"
+        @response_json["content"][index]["input"] = ""
+      end
+    when "content_block_delta"
+      delta_chunk = event["delta"]
+      if delta_chunk["type"] == "text_delta"
+        delta = delta_chunk["text"]
+        @response_json["content"][index]["text"] += delta if delta
+      elsif delta_chunk["type"] == "input_json_delta"
+        @response_json["content"][index]["input"] += delta_chunk["partial_json"]
+      end
+    when "content_block_stop"
+      content_block = @response_json["content"][index]
+      if content_block&.dig("type") == "tool_use"
+        begin
+          content_block["input"] = JSON.parse(content_block["input"])
+        rescue JSON::ParserError
+          # If parsing fails, leave as a string
+        end
+      end
+    when "message_delta"
+      @finish_reason = event.dig("delta", "stop_reason")
+      @response_json["usage"]["output_tokens"] = event.dig("usage", "output_tokens")
+    when "message_stop"
+      @finish_reason = "stop"
+    when "error"
+      error_details = event["error"]
+      raise Raif::Errors::StreamingError.new(
+        message: error_details["message"],
+        type: error_details["type"],
+        event: event
+      )
+    end
+
+    [delta, @finish_reason]
+  end
+
+  def current_response_json
+    @response_json["stop_reason"] = @finish_reason
+    @response_json
+  end
+
+end

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

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+class Raif::StreamingResponses::Bedrock
+
+  def initialize_new_message
+    # Initialize empty AWS response object
+    @message = Aws::BedrockRuntime::Types::Message.new(
+      role: "assistant",
+      content: []
+    )
+
+    @output = Aws::BedrockRuntime::Types::ConverseOutput::Message.new(message: @message)
+
+    @usage = Aws::BedrockRuntime::Types::TokenUsage.new(
+      input_tokens: 0,
+      output_tokens: 0,
+      total_tokens: 0
+    )
+
+    @response = Aws::BedrockRuntime::Types::ConverseResponse.new(
+      output: @output,
+      usage: @usage,
+      stop_reason: nil
+    )
+  end
+
+  def process_streaming_event(event_type, event)
+    delta = nil
+
+    case event.event_type
+    when :message_start
+      initialize_new_message
+    when :content_block_start
+      index = event.content_block_index
+
+      if event.start.is_a?(Aws::BedrockRuntime::Types::ContentBlockStart::ToolUse)
+        tool_use = event.start.tool_use
+        @message.content[index] = Aws::BedrockRuntime::Types::ContentBlock.new(
+          tool_use: Aws::BedrockRuntime::Types::ToolUseBlock.new(
+            tool_use_id: tool_use.tool_use_id,
+            name: tool_use.name,
+            input: ""
+          )
+        )
+      else
+        @message.content[index] = Aws::BedrockRuntime::Types::ContentBlock::Text.new(text: "")
+      end
+    when :content_block_delta
+      index = event.content_block_index
+
+      if event.delta.is_a?(Aws::BedrockRuntime::Types::ContentBlockDelta::Text)
+        @message.content[index] ||= Aws::BedrockRuntime::Types::ContentBlock::Text.new(text: "")
+        delta = event.delta.text
+        @message.content[index].text += delta
+      elsif event.delta.is_a?(Aws::BedrockRuntime::Types::ContentBlockDelta::ToolUse)
+        tool_use = event.delta.tool_use
+        @message.content[index] ||= Aws::BedrockRuntime::Types::ContentBlock.new
+        @message.content[index].tool_use ||= Aws::BedrockRuntime::Types::ToolUseBlock.new(
+          tool_use_id: tool_use.tool_use_id,
+          name: tool_use.name,
+          input: ""
+        )
+
+        @message.content[index].tool_use.input += event.delta.tool_use.input
+      end
+    when :content_block_stop
+      content_block = @message.content[event.content_block_index]
+
+      if content_block&.tool_use&.input.is_a?(String)
+        begin
+          content_block.tool_use.input = JSON.parse(content_block.tool_use.input)
+        rescue JSON::ParserError
+          # If parsing fails, leave as a string
+        end
+      end
+    when :message_stop
+      @response.stop_reason = event.stop_reason
+    when :metadata
+      @response.usage = event.usage if event.respond_to?(:usage)
+    end
+
+    [delta, @response.stop_reason]
+  end
+
+  def current_response
+    @response
+  end
+
+end

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

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+class Raif::StreamingResponses::OpenAiCompletions
+  attr_reader :raw_response, :tool_calls
+
+  def initialize
+    @id = nil
+    @raw_response = ""
+    @tool_calls = []
+    @usage = {}
+    @finish_reason = nil
+    @response_json = {}
+  end
+
+  def process_streaming_event(event_type, event)
+    @id ||= event["id"]
+    delta_chunk = event.dig("choices", 0, "delta")
+    finish_reason = event.dig("choices", 0, "finish_reason")
+    @finish_reason ||= finish_reason
+
+    delta_content = delta_chunk&.dig("content")
+    @raw_response += delta_content if delta_content
+
+    if delta_chunk&.key?("tool_calls")
+      delta_chunk["tool_calls"].each do |tool_call_chunk|
+        index = tool_call_chunk["index"]
+        @tool_calls[index] ||= { "function" => {} }
+        @tool_calls[index]["id"] ||= tool_call_chunk["id"]
+        if tool_call_chunk.dig("function", "name")
+          @tool_calls[index]["function"]["name"] = tool_call_chunk.dig("function", "name")
+        end
+        if tool_call_chunk.dig("function", "arguments")
+          @tool_calls[index]["function"]["arguments"] ||= ""
+          @tool_calls[index]["function"]["arguments"] += tool_call_chunk.dig("function", "arguments")
+        end
+      end
+    end
+
+    @usage = event["usage"] if event["usage"]
+
+    [delta_content, finish_reason]
+  end
+
+  def current_response_json
+    message = {
+      "role" => "assistant",
+      "content" => @raw_response
+    }
+
+    if @tool_calls.any?
+      message["content"] = nil # Per OpenAI spec, content is null if tool_calls are present
+      message["tool_calls"] = @tool_calls.map do |tc|
+        # The streaming format for tool calls is slightly different from the final format.
+        # We need to adjust it here.
+        {
+          "id" => tc["id"],
+          "type" => "function",
+          "function" => {
+            "name" => tc.dig("function", "name"),
+            "arguments" => tc.dig("function", "arguments")
+          }
+        }
+      end
+    end
+
+    {
+      "id" => @id,
+      "choices" => [{
+        "index" => 0,
+        "message" => message,
+        "finish_reason" => @finish_reason
+      }],
+      "usage" => @usage
+    }
+  end
+end

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

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+class Raif::StreamingResponses::OpenAiResponses
+
+  def initialize
+    @output_items = []
+    @finish_reason = nil
+  end
+
+  def process_streaming_event(event_type, event)
+    output_index = event["output_index"]
+    content_index = event["content_index"]
+
+    delta = nil
+
+    case event["type"]
+    when "response.created"
+      @id = event["response"]["id"]
+    when "response.output_item.added", "response.output_item.done"
+      @output_items[output_index] = event["item"]
+    when "response.content_part.added", "response.content_part.done"
+      @output_items[output_index]["content"] ||= []
+      @output_items[output_index]["content"][content_index] = event["part"]
+    when "response.output_text.delta"
+      delta = event["delta"]
+      @output_items[output_index]["content"][content_index]["text"] ||= ""
+      @output_items[output_index]["content"][content_index]["text"] += event["delta"]
+    when "response.output_text.done"
+      @output_items[output_index]["content"][content_index]["text"] = event["text"]
+    when "response.completed"
+      @usage = event["response"]["usage"]
+      @finish_reason = "stop"
+    when "error"
+      raise Raif::Errors::StreamingError.new(
+        message: event["message"],
+        type: event["type"],
+        code: event["code"],
+        event: event
+      )
+    end
+
+    [delta, @finish_reason]
+  end
+
+  # The response we've built up so far.
+  def current_response_json
+    {
+      "id" => @id,
+      "output" => @output_items,
+      "usage" => @usage
+    }
+  end
+
+end

data/app/models/raif/task.rb

@@ -7,6 +7,10 @@ module Raif
     include Raif::Concerns::HasAvailableModelTools
     include Raif::Concerns::InvokesModelTools
     include Raif::Concerns::LlmResponseParsing
+    include Raif::Concerns::LlmTemperature
+    include Raif::Concerns::JsonSchemaDefinition
+
+    llm_temperature 0.7
 
     belongs_to :creator, polymorphic: true
 
@@ -20,12 +24,25 @@ module Raif
 
     delegate :json_response_schema, to: :class
 
-    after_initialize -> { self.available_model_tools ||= [] }
+    scope :completed, -> { where.not(completed_at: nil) }
+    scope :failed, -> { where.not(failed_at: nil) }
+    scope :in_progress, -> { where.not(started_at: nil).where(completed_at: nil, failed_at: nil) }
+    scope :pending, -> { where(started_at: nil, completed_at: nil, failed_at: nil) }
+
+    attr_accessor :files, :images
 
-    def self.llm_response_format(format)
-      raise ArgumentError, "response_format must be one of: #{response_formats.keys.join(", ")}" unless response_formats.keys.include?(format.to_s)
+    after_initialize -> { self.available_model_tools ||= [] }
 
-      after_initialize -> { self.response_format = format }, if: :new_record?
+    def status
+      if completed_at?
+        :completed
+      elsif failed_at?
+        :failed
+      elsif started_at?
+        :in_progress
+      else
+        :pending
+      end
     end
 
     # The primary interface for running a task. It will hit the LLM with the task's prompt and system prompt and return a Raif::Task object.
@@ -34,10 +51,13 @@ module Raif
     # @param creator [Object] The creator of the task (polymorphic association)
     # @param available_model_tools [Array<Class>] Optional array of model tool classes that will be provided to the LLM for it to invoke.
     # @param llm_model_key [Symbol, String] Optional key for the LLM model to use. If blank, Raif.config.default_llm_model_key will be used.
+    # @param images [Array] Optional array of Raif::ModelImageInput objects to include with the prompt.
+    # @param files [Array] Optional array of Raif::ModelFileInput objects to include with the prompt.
     # @param args [Hash] Additional arguments to pass to the instance of the task that is created.
     # @return [Raif::Task, nil] The task instance that was created and run.
-    def self.run(creator:, available_model_tools: [], llm_model_key: nil, **args)
-      task = new(creator:, llm_model_key:, available_model_tools:, started_at: Time.current, **args)
+    def self.run(creator:, available_model_tools: [], llm_model_key: nil, images: [], files: [], **args)
+      task = new(creator:, llm_model_key:, available_model_tools:, started_at: Time.current, images: images, files: files, **args)
+
       task.save!
       task.run
       task
@@ -58,17 +78,19 @@ module Raif
       task
     end
 
-    def run
+    def run(skip_prompt_population: false)
       update_columns(started_at: Time.current) if started_at.nil?
 
-      populate_prompts
-      messages = [{ "role" => "user", "content" => prompt }]
+      populate_prompts unless skip_prompt_population
+      messages = [{ "role" => "user", "content" => message_content }]
+
       mc = llm.chat(
         messages: messages,
         source: self,
         system_prompt: system_prompt,
         response_format: response_format.to_sym,
-        available_model_tools: available_model_tools
+        available_model_tools: available_model_tools,
+        temperature: self.class.temperature
       )
 
       self.raif_model_completion = mc.becomes(Raif::ModelCompletion)
@@ -80,13 +102,18 @@ module Raif
       self
     end
 
+    def re_run
+      update_columns(started_at: Time.current)
+      run(skip_prompt_population: true)
+    end
+
     # Returns the LLM prompt for the task.
     #
     # @param creator [Object] The creator of the task (polymorphic association)
     # @param args [Hash] Additional arguments to pass to the instance of the task that is created.
     # @return [String] The LLM prompt for the task.
     def self.prompt(creator:, **args)
-      new(creator:, **args).prompt
+      new(creator:, **args).build_prompt
     end
 
     # Returns the LLM system prompt for the task.
@@ -95,25 +122,53 @@ module Raif
     # @param args [Hash] Additional arguments to pass to the instance of the task that is created.
     # @return [String] The LLM system prompt for the task.
     def self.system_prompt(creator:, **args)
-      new(creator:, **args).system_prompt
+      new(creator:, **args).build_system_prompt
     end
 
-    def self.json_response_schema
-      nil
+    def self.json_response_schema(&block)
+      if block_given?
+        json_schema_definition(:json_response, &block)
+      elsif schema_defined?(:json_response)
+        schema_for(:json_response)
+      end
     end
 
-  private
-
     def build_prompt
       raise NotImplementedError, "Raif::Task subclasses must implement #build_prompt"
     end
 
     def build_system_prompt
       sp = Raif.config.task_system_prompt_intro
+      sp = sp.call(self) if sp.respond_to?(:call)
       sp += system_prompt_language_preference if requested_language_key.present?
       sp
     end
 
+  private
+
+    def message_content
+      # If there are no images or files, just return the message content can just be a string with the prompt
+      return prompt if images.blank? && files.blank?
+
+      content = [{ "type" => "text", "text" => prompt }]
+
+      images.each do |image|
+        raise Raif::Errors::InvalidModelImageInputError,
+          "Images must be a Raif::ModelImageInput: #{image.inspect}" unless image.is_a?(Raif::ModelImageInput)
+
+        content << image
+      end
+
+      files.each do |file|
+        raise Raif::Errors::InvalidFileInputError,
+          "Files must be a Raif::ModelFileInput: #{file.inspect}" unless file.is_a?(Raif::ModelFileInput)
+
+        content << file
+      end
+
+      content
+    end
+
     def populate_prompts
       self.requested_language_key ||= creator.preferred_language_key if creator.respond_to?(:preferred_language_key)
       self.prompt = build_prompt