raif 1.3.0 → 1.5.0

data/lib/raif/messages.rb

@@ -0,0 +1,180 @@
+# frozen_string_literal: true
+
+module Raif
+  # Message types for agent conversation_history and conversation llm_messages.
+  #
+  # These classes provide a structured API for creating messages that get stored
+  # as JSONB and passed to LLM providers. Each class has:
+  # - Named parameters for initialization
+  # - `to_h` for converting to hash format (for storage/API calls)
+  # - `from_h` class method for deserializing from stored hashes
+  #
+  # @example Creating messages
+  #   message = Raif::Messages::ToolCall.new(
+  #     name: "wikipedia_search",
+  #     arguments: { query: "Ruby" },
+  #     provider_tool_call_id: "call_123"
+  #   )
+  #   conversation_history << message.to_h
+  #
+  # @example Deserializing stored messages
+  #   messages = Raif::Messages.from_array(agent.conversation_history)
+  #   messages.each do |msg|
+  #     case msg
+  #     when Raif::Messages::ToolCall
+  #       puts "Tool: #{msg.name}"
+  #     when Raif::Messages::UserMessage
+  #       puts "User: #{msg.content}"
+  #     end
+  #   end
+  module Messages
+    # User role message
+    class UserMessage
+      attr_reader :content
+
+      # @param content [String] The user's message content
+      def initialize(content:)
+        @content = content
+      end
+
+      # @return [Hash] Hash representation for JSONB storage and LLM APIs
+      def to_h
+        { "role" => "user", "content" => content }
+      end
+
+      # Deserialize from a hash
+      # @param hash [Hash] A hash with "content" key
+      # @return [UserMessage]
+      def self.from_h(hash)
+        new(content: hash["content"])
+      end
+    end
+
+    # Assistant role message
+    class AssistantMessage
+      attr_reader :content
+
+      # @param content [String] The assistant's message content
+      def initialize(content:)
+        @content = content
+      end
+
+      # @return [Hash] Hash representation for JSONB storage and LLM APIs
+      def to_h
+        { "role" => "assistant", "content" => content }
+      end
+
+      # Deserialize from a hash
+      # @param hash [Hash] A hash with "content" key
+      # @return [AssistantMessage]
+      def self.from_h(hash)
+        new(content: hash["content"])
+      end
+    end
+
+    # Tool invocation request from the assistant
+    class ToolCall
+      attr_reader :provider_tool_call_id, :name, :arguments, :assistant_message, :provider_metadata
+
+      # @param name [String] The tool name (snake_case)
+      # @param arguments [Hash] The arguments passed to the tool
+      # @param provider_tool_call_id [String, nil] Provider-assigned ID for the tool call
+      # @param assistant_message [String, nil] Optional assistant message accompanying the tool call
+      # @param provider_metadata [Hash, nil] Provider-specific metadata (e.g., Google's thoughtSignature)
+      def initialize(name:, arguments:, provider_tool_call_id: nil, assistant_message: nil, provider_metadata: nil)
+        @provider_tool_call_id = provider_tool_call_id
+        @name = name
+        @arguments = arguments
+        @assistant_message = assistant_message
+        @provider_metadata = provider_metadata
+      end
+
+      # @return [Hash] Hash representation for JSONB storage and LLM APIs
+      def to_h
+        {
+          "type" => "tool_call",
+          "provider_tool_call_id" => provider_tool_call_id,
+          "name" => name,
+          "arguments" => arguments,
+          "assistant_message" => assistant_message,
+          "provider_metadata" => provider_metadata
+        }.compact
+      end
+
+      # Deserialize from a hash
+      # @param hash [Hash] A hash with tool call fields
+      # @return [ToolCall]
+      def self.from_h(hash)
+        new(
+          name: hash["name"],
+          arguments: hash["arguments"],
+          provider_tool_call_id: hash["provider_tool_call_id"],
+          assistant_message: hash["assistant_message"],
+          provider_metadata: hash["provider_metadata"]
+        )
+      end
+    end
+
+    # Result of a tool invocation
+    class ToolCallResult
+      attr_reader :provider_tool_call_id, :name, :result
+
+      # @param result [Hash, String] The result returned by the tool
+      # @param provider_tool_call_id [String, nil] Provider-assigned ID matching the tool call
+      # @param name [String, nil] The tool name (required by some providers like Google)
+      def initialize(result:, provider_tool_call_id: nil, name: nil)
+        @provider_tool_call_id = provider_tool_call_id
+        @name = name
+        @result = result
+      end
+
+      # @return [Hash] Hash representation for JSONB storage and LLM APIs
+      def to_h
+        {
+          "type" => "tool_call_result",
+          "provider_tool_call_id" => provider_tool_call_id,
+          "name" => name,
+          "result" => result
+        }.compact
+      end
+
+      # Deserialize from a hash
+      # @param hash [Hash] A hash with tool call result fields
+      # @return [ToolCallResult]
+      def self.from_h(hash)
+        new(
+          provider_tool_call_id: hash["provider_tool_call_id"],
+          name: hash["name"],
+          result: hash["result"]
+        )
+      end
+    end
+
+    class << self
+      # Deserialize a single message hash into the appropriate message class
+      # @param hash [Hash] A message hash with either "role" or "type" key
+      # @return [UserMessage, AssistantMessage, ToolCall, ToolCallResult]
+      # @raise [ArgumentError] if the hash doesn't match a known message type
+      def from_h(hash)
+        if hash["type"] == "tool_call"
+          ToolCall.from_h(hash)
+        elsif hash["type"] == "tool_call_result"
+          ToolCallResult.from_h(hash)
+        elsif hash["role"] == "user"
+          UserMessage.from_h(hash)
+        elsif hash["role"] == "assistant"
+          AssistantMessage.from_h(hash)
+        else
+          raise ArgumentError, "Unknown message type: #{hash.inspect}"
+        end
+      end
+
+      # Deserialize an array of message hashes
+      # @param messages [Array<Hash>] Array of message hashes
+      # @return [Array<UserMessage, AssistantMessage, ToolCall, ToolCallResult>]
+      def from_array(messages)
+        messages.map { |msg| from_h(msg) }
+      end
+    end
+  end
+end

data/lib/raif/prompt_studio_comparison_builder.rb

@@ -0,0 +1,138 @@
+# frozen_string_literal: true
+
+module Raif
+  class PromptStudioComparisonBuilder
+    # Attempts to rebuild the prompt from current code for a given record.
+    # Returns a hash with the rendered prompts and any warnings.
+    def self.build(record)
+      new(record).build
+    end
+
+    def initialize(record)
+      @record = record
+    end
+
+    def build
+      {
+        original_prompt: original_prompt,
+        original_system_prompt: original_system_prompt,
+        current_prompt: current_prompt,
+        current_system_prompt: current_system_prompt,
+        prompt_changed: changed?(original_prompt, current_prompt),
+        system_prompt_changed: changed?(original_system_prompt, current_system_prompt),
+        has_stale_references: has_stale_references?,
+        warnings: warnings,
+        original_prompt_tokens: original_prompt_tokens,
+        original_prompt_token_cost: original_prompt_token_cost,
+        current_prompt_token_estimate: current_prompt_token_estimate,
+        current_prompt_cost_estimate: current_prompt_cost_estimate
+      }
+    end
+
+  private
+
+    def original_prompt
+      @original_prompt ||= @record.respond_to?(:prompt) ? @record.prompt : nil
+    end
+
+    def original_system_prompt
+      @original_system_prompt ||= @record.system_prompt
+    end
+
+    def current_prompt
+      return @current_prompt if defined?(@current_prompt)
+
+      @current_prompt = begin
+        @record.build_prompt
+      rescue NotImplementedError
+        nil
+      rescue => e
+        warnings << "Error rendering current prompt: #{e.message}"
+        nil
+      end
+    end
+
+    def current_system_prompt
+      return @current_system_prompt if defined?(@current_system_prompt)
+
+      @current_system_prompt = begin
+        @record.build_system_prompt
+      rescue NotImplementedError
+        nil
+      rescue => e
+        warnings << "Error rendering current system prompt: #{e.message}"
+        nil
+      end
+    end
+
+    def warnings
+      @warnings ||= [].tap do |w|
+        w << I18n.t("raif.admin.prompt_studio.common.warning_stale_reference") if has_stale_references?
+      end
+    end
+
+    def has_stale_references?
+      return @has_stale_references if defined?(@has_stale_references)
+
+      @has_stale_references = detect_stale_references
+    end
+
+    def detect_stale_references
+      return false unless @record.respond_to?(:run_with) && @record.run_with.present?
+
+      @record.run_with.each do |_key, value|
+        if value.is_a?(String) && value.start_with?("gid://")
+          begin
+            return true if GlobalID::Locator.locate(value).nil?
+          rescue StandardError
+            return true
+          end
+        end
+      end
+
+      false
+    end
+
+    def changed?(original, current)
+      original.present? && current.present? && original.strip != current.strip
+    end
+
+    def original_prompt_tokens
+      return unless @record.respond_to?(:raif_model_completion)
+
+      @record.raif_model_completion&.prompt_tokens
+    end
+
+    def original_prompt_token_cost
+      return unless @record.respond_to?(:raif_model_completion)
+
+      @record.raif_model_completion&.prompt_token_cost
+    end
+
+    def current_prompt_token_estimate
+      return unless prompt_changed? || system_prompt_changed?
+
+      Raif::TokenEstimator.estimate_tokens(current_system_prompt, current_prompt)
+    end
+
+    def current_prompt_cost_estimate
+      token_estimate = current_prompt_token_estimate
+      return unless token_estimate
+
+      return unless @record.llm_model_key.present?
+
+      llm_config = Raif.llm_config(@record.llm_model_key.to_sym)
+      return unless llm_config&.dig(:input_token_cost)
+
+      llm_config[:input_token_cost] * token_estimate
+    end
+
+    def prompt_changed?
+      changed?(original_prompt, current_prompt)
+    end
+
+    def system_prompt_changed?
+      changed?(original_system_prompt, current_system_prompt)
+    end
+  end
+end

data/lib/raif/token_estimator.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Raif
+  class TokenEstimator
+    def self.available?
+      return true if defined?(::Tiktoken)
+
+      require "tiktoken_ruby"
+      !!defined?(::Tiktoken)
+    rescue LoadError
+      false
+    end
+
+    # Estimates the total token count for a prompt + system prompt combination.
+    # Returns nil if tiktoken_ruby is not installed.
+    def self.estimate_tokens(*texts)
+      return unless available?
+
+      encoder = encoder_for_model("gpt-4")
+      texts.compact.sum { |text| encoder.encode(text).length }
+    end
+
+    def self.encoder_for_model(model)
+      @encoders ||= {}
+      @encoders[model] ||= ::Tiktoken.encoding_for_model(model)
+    end
+  end
+end

data/lib/raif/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Raif
-  VERSION = "1.3.0"
+  VERSION = "1.5.0"
 end

data/lib/raif.rb

@@ -10,6 +10,9 @@ require "raif/llm_registry"
 require "raif/embedding_model_registry"
 require "raif/json_schema_builder"
 require "raif/migration_checker"
+require "raif/messages"
+require "raif/prompt_studio_comparison_builder"
+require "raif/token_estimator"
 
 require "faraday"
 require "event_stream_parser"
@@ -41,4 +44,12 @@ module Raif
   def self.running_evals?
     ENV["RAIF_RUNNING_EVALS"] == "true"
   end
+
+  def self.default_request_options
+    {
+      open_timeout: config.request_open_timeout,
+      read_timeout: config.request_read_timeout,
+      write_timeout: config.request_write_timeout
+    }.compact
+  end
 end

data/lib/tasks/annotate_rb.rake

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+# This rake task was added by annotate_rb gem.
+
+# Can set `ANNOTATERB_SKIP_ON_DB_TASKS` to be anything to skip this
+if Rails.env.development? && ENV["ANNOTATERB_SKIP_ON_DB_TASKS"].nil?
+  require "annotate_rb"
+
+  AnnotateRb::Core.load_rake_tasks
+end

data/spec/support/rspec_helpers.rb

@@ -3,17 +3,23 @@
 module Raif
   module RspecHelpers
 
-    def stubbed_llm(llm_model_key, &block)
+    def stubbed_llm(llm_model_key, source_instance, &block)
       test_llm = Raif.llm(llm_model_key.to_sym)
 
-      allow(test_llm).to receive(:perform_model_completion!) do |model_completion|
-        result = block.call(model_completion.messages, model_completion)
+      allow(test_llm).to receive(:perform_model_completion!) do |model_completion, &streaming_block|
+        result = block.call(model_completion.messages, model_completion, source_instance)
         model_completion.raw_response = result if result.is_a?(String)
         model_completion.completion_tokens = rand(100..2000)
         model_completion.prompt_tokens = rand(100..2000)
         model_completion.total_tokens = model_completion.completion_tokens + model_completion.prompt_tokens
         model_completion.save!
 
+        if streaming_block && result.is_a?(String)
+          result.chars.each_slice(25) do |chunk|
+            streaming_block.call(model_completion, chunk.join, nil)
+          end
+        end
+
         model_completion
       end
 
@@ -24,10 +30,10 @@ module Raif
       allow(Raif.config).to receive(:llm_api_requests_enabled){ true }
 
       if task.is_a?(Raif::Task)
-        allow(task).to receive(:llm){ stubbed_llm(task.llm_model_key, &block) }
+        allow(task).to receive(:llm){ stubbed_llm(task.llm_model_key, task, &block) }
       else
         allow_any_instance_of(task).to receive(:llm) do |task_instance|
-          stubbed_llm(task_instance.llm_model_key, &block)
+          stubbed_llm(task_instance.llm_model_key, task_instance, &block)
         end
       end
     end
@@ -36,10 +42,10 @@ module Raif
       allow(Raif.config).to receive(:llm_api_requests_enabled){ true }
 
       if conversation.is_a?(Raif::Conversation)
-        allow(conversation).to receive(:llm){ stubbed_llm(conversation.llm_model_key, &block) }
+        allow(conversation).to receive(:llm){ stubbed_llm(conversation.llm_model_key, conversation, &block) }
       else
         allow_any_instance_of(conversation).to receive(:llm) do |conversation_instance|
-          stubbed_llm(conversation_instance.llm_model_key, &block)
+          stubbed_llm(conversation_instance.llm_model_key, conversation_instance, &block)
         end
       end
     end
@@ -48,10 +54,10 @@ module Raif
       allow(Raif.config).to receive(:llm_api_requests_enabled){ true }
 
       if agent.is_a?(Raif::Agent)
-        allow(agent).to receive(:llm){ stubbed_llm(agent.llm_model_key, &block) }
+        allow(agent).to receive(:llm){ stubbed_llm(agent.llm_model_key, agent, &block) }
       else
         allow_any_instance_of(agent).to receive(:llm) do |agent_instance|
-          stubbed_llm(agent_instance.llm_model_key, &block)
+          stubbed_llm(agent_instance.llm_model_key, agent_instance, &block)
         end
       end
     end

data/spec/support/test_task.rb

@@ -30,6 +30,15 @@ class Raif::TestJsonTask < Raif::Task
   end
 end
 
+class Raif::TestCachedTask < Raif::Task
+  enable_anthropic_prompt_caching
+  enable_bedrock_prompt_caching
+
+  def build_prompt
+    "Tell me a joke"
+  end
+end
+
 class Raif::TestHtmlTask < Raif::Task
   llm_response_format :html
   llm_response_allowed_tags %w[p b i u s a]

data/spec/support/test_template_task.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+class Raif::TestTemplateTask < Raif::Task
+  run_with :topic
+
+  after_initialize -> { self.topic ||= "pirates" }
+
+  def topic_description
+    "the topic of #{topic}"
+  end
+end
+
+class Raif::TestTemplateSystemPromptTask < Raif::Task
+  run_with :persona
+
+  after_initialize -> { self.persona ||= "comedian" }
+
+  def build_prompt
+    "Tell me a joke"
+  end
+end
+
+class Raif::TestTemplateConversation < Raif::Conversation
+  attr_writer :persona
+
+  def persona
+    @persona || "helpful assistant"
+  end
+end
+
+class Raif::TestTemplateWithPartialTask < Raif::Task
+  run_with :topic
+
+  after_initialize -> { self.topic ||= "dogs" }
+end
+
+class Raif::TestTemplateAgent < Raif::Agent
+  run_with :agent_role
+
+  after_initialize -> { self.agent_role ||= "researcher" }
+end