ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/active_record/tool_call.rb

@@ -4,8 +4,30 @@ require 'ruby_llm/active_record/payload_helpers'
 
 module RubyLLM
   module ActiveRecord
-    # RubyLLM's private persistence for provider tool calls.
-    class ToolCall < Record # :nodoc:
+    # A tool call a message requested, in RubyLLM's +ruby_llm_tool_calls+
+    # table. An assistant message reaches its calls through
+    # +ruby_llm_tool_calls+, oldest first, and a tool message reaches the
+    # call it answers through +ruby_llm_parent_tool_call+.
+    #
+    #   message.ruby_llm_tool_calls.pluck(:name) # => ["weather", "search"]
+    #   RubyLLM::ActiveRecord::ToolCall.where(name: "weather").count
+    #
+    # Each row belongs to +message+, the assistant message that requested
+    # it, and to +result+, the tool message that answers it once there is
+    # one. Both associations are polymorphic. Its columns:
+    #
+    # [+tool_call_id+] The call's id, unique across the table.
+    # [+name+] The tool's name.
+    # [+namespace+] The namespace of a namespaced tool, or +nil+.
+    # [+arguments+] The arguments the model passed, as a Hash. Empty when
+    #               the model's text was not a JSON object.
+    # [+approval+] <tt>"pending"</tt>, <tt>"approved"</tt>, <tt>"remembered"</tt>,
+    #              or <tt>"denied"</tt>, or +nil+ when no approval was asked.
+    # [+remote+] Whether the provider ran the tool instead of RubyLLM.
+    #
+    # RubyLLM writes these rows. Read and query them, and add scopes and
+    # columns of your own, but do not create or update them.
+    class ToolCall < Record
       self.table_name = 'ruby_llm_tool_calls'
 
       include PayloadHelpers
@@ -15,17 +37,46 @@ module RubyLLM
 
       validates :tool_call_id, :name, presence: true
 
+      # Returns this row as a RubyLLM::ToolCall.
       def to_llm
         RubyLLM::ToolCall.new(
           id: tool_call_id,
+          provider_call_id: self[:provider_call_id],
           name: name,
-          arguments: arguments || {},
+          namespace: self[:namespace],
+          arguments: arguments,
+          malformed_arguments: malformed_arguments,
           thought_signature: thought_signature,
           remote: self[:remote]
         )
       end
 
-      def tool_error_message
+      # Returns the arguments the model passed, as a Hash, or an empty Hash
+      # when the model's text was not a JSON object.
+      def arguments
+        value = self[:arguments]
+        value.is_a?(Hash) ? value : {}
+      end
+
+      def malformed_arguments # :nodoc:
+        value = self[:arguments]
+        value if value.is_a?(String)
+      end
+
+      # Returns whether a process is running this tool call: it has no
+      # result yet and its claim has not expired. The process running it
+      # keeps renewing the claim, which expires within a minute of that
+      # process dying.
+      def running?
+        (result_id.nil? && has_attribute?(:claimed_until) && claimed_until&.future?) || false
+      end
+
+      # Returns whether the call waits on a task or on questions it asked.
+      def paused?
+        has_attribute?(:mcp_state) && mcp_state.present?
+      end
+
+      def tool_error_message # :nodoc:
         payload_error_message(arguments)
       end
     end

data/lib/ruby_llm/active_record/usage.rb

@@ -2,10 +2,43 @@
 
 module RubyLLM
   module ActiveRecord
-    # RubyLLM's private normalized persistence for one provider attempt.
-    class Usage < Record # :nodoc:
+    # One provider attempt in RubyLLM's usage ledger, the +ruby_llm_usages+
+    # table. RubyLLM writes a row for every billable attempt, including
+    # retries, cancelled streams, and one-shot operations without a chat.
+    # A write that fails is logged and skipped, so it never breaks the
+    # operation.
+    #
+    #   chat.ruby_llm_usages.sum(:total_cost)
+    #   RubyLLM::ActiveRecord::Usage.where(owner: current_user).group(:operation).sum(:total_cost)
+    #
+    # Each row has optional polymorphic +chat+, +message+, and +owner+
+    # associations. The +chronological+ scope orders rows oldest first.
+    # Its columns:
+    #
+    # [+operation+] What ran: <tt>"chat"</tt>, <tt>"embedding"</tt>, <tt>"image"</tt>,
+    #               <tt>"speech"</tt>, <tt>"transcription"</tt>, and so on.
+    # [+provider+] The provider's slug.
+    # [+model+] The model the provider ran, or its own resource identifier.
+    # [+status+] <tt>"pending"</tt>, <tt>"succeeded"</tt>, <tt>"failed"</tt>, or <tt>"cancelled"</tt>.
+    # [+input_tokens+, +output_tokens+, +cache_read_tokens+, +cache_write_tokens+, +thinking_tokens+]
+    #   Token counts, in the bucket the provider bills each at.
+    # [+characters+, +audio_seconds+, +search_units+] Units billed instead of tokens.
+    # [+server_tool_use+] Counts of provider-run tool calls, as a Hash.
+    # [+input_cost+, +output_cost+, +cache_read_cost+, +cache_write_cost+, +thinking_cost+]
+    #   The cost of each token bucket, frozen when the attempt finished.
+    # [+characters_cost+, +audio_seconds_cost+, +search_units_cost+] The cost of each unit.
+    # [+total_cost+] The attempt's total cost, or +nil+ when it is unknown.
+    # [+workflow_id+, +workflow_name+, +workflow_step_id+, +workflow_step_name+]
+    #   The workflow and step that made the attempt, if any.
+    #
+    # RubyLLM writes these rows. Read and query them, and add scopes and
+    # columns of your own, but do not create or update them.
+    class Usage < Record
       self.table_name = 'ruby_llm_usages'
 
+      WORKFLOW_COLUMNS = %i[workflow_id workflow_name workflow_step_id workflow_step_name].freeze # :nodoc:
+      UNIT_COLUMNS = ::RubyLLM::Cost::UNITS.flat_map { |unit| [unit, :"#{unit}_cost"] }.freeze # :nodoc:
+
       belongs_to :chat, polymorphic: true, optional: true
       belongs_to :message, polymorphic: true, optional: true
       belongs_to :owner, polymorphic: true, optional: true
@@ -22,7 +55,7 @@ module RubyLLM
       # it. A savepoint keeps a failed insert from aborting the caller's
       # transaction, and a thread that held no connection gives back the
       # one it borrowed.
-      def self.record(entry)
+      def self.record(entry) # :nodoc:
         return unless entry.model
 
         connection_pool.with_connection do
@@ -35,11 +68,11 @@ module RubyLLM
         nil
       end
 
-      def self.ledger_available?
+      def self.ledger_available? # :nodoc:
         table_exists? && columns_hash['chat_id']&.null
       end
 
-      def self.attributes_for(entry)
+      def self.attributes_for(entry) # :nodoc:
         tokens = entry.tokens
         cost = entry.cost
         attributes = {
@@ -60,10 +93,28 @@ module RubyLLM
           total_cost: cost.total
         }
         attributes[:server_tool_use] = tokens.server_tool_use if column_names.include?('server_tool_use')
+        attributes.merge!(unit_attributes(tokens, cost)) if column_names.include?('search_units')
         attributes[:owner] = entry.owner if entry.owner.is_a?(::ActiveRecord::Base) && column_names.include?('owner_id')
+        attributes.merge!(workflow_attributes(entry.workflow)) if column_names.include?('workflow_id')
         attributes
       end
 
+      def self.unit_attributes(tokens, cost) # :nodoc:
+        ::RubyLLM::Cost::UNITS.each_with_object({}) do |unit, attributes|
+          attributes[unit] = tokens.public_send(unit)
+          attributes[:"#{unit}_cost"] = cost.public_send(unit)
+        end
+      end
+
+      def self.workflow_attributes(workflow) # :nodoc:
+        workflow.to_h.slice(*WORKFLOW_COLUMNS)
+      end
+
+      def self.owner_preload # :nodoc:
+        column_names.include?('owner_id') ? :owner : []
+      end
+
+      # Returns the attempt's counts as RubyLLM::Tokens.
       def tokens
         RubyLLM::Tokens.new(
           input: input_tokens,
@@ -71,10 +122,12 @@ module RubyLLM
           cache_read: cache_read_tokens,
           cache_write: cache_write_tokens,
           thinking: thinking_tokens,
+          **recorded_units,
           server_tool_use: (self[:server_tool_use] if has_attribute?(:server_tool_use))
         )
       end
 
+      # Returns the attempt's frozen cost as RubyLLM::Cost.
       def cost
         recorded = {
           input: input_cost,
@@ -82,20 +135,21 @@ module RubyLLM
           cache_read: cache_read_cost,
           cache_write: cache_write_cost,
           thinking: thinking_cost,
+          **recorded_units(suffix: '_cost'),
           total: total_cost
         }.compact
         RubyLLM::Cost.from_h(recorded, tokens: tokens)
       end
 
-      def usage_available?
+      def usage_available? # :nodoc:
         tokens.to_h.any?
       end
 
-      def cost_available?
+      def cost_available? # :nodoc:
         !total_cost.nil?
       end
 
-      def to_entry
+      def to_entry # :nodoc:
         ::RubyLLM::Accounting::Usage::Entry.new(
           operation: operation,
           provider: provider,
@@ -103,9 +157,25 @@ module RubyLLM
           status: status,
           tokens: tokens,
           cost: cost,
-          message: message
+          message: message,
+          owner: (owner if has_attribute?(:owner_id)),
+          workflow: workflow_context
         )
       end
+
+      private
+
+      def recorded_units(suffix: nil)
+        ::RubyLLM::Cost::UNITS.each_with_object({}) do |unit, units|
+          column = :"#{unit}#{suffix}"
+          units[unit] = self[column] if has_attribute?(column)
+        end
+      end
+
+      def workflow_context
+        context = WORKFLOW_COLUMNS.select { |name| has_attribute?(name) }.index_with { |name| self[name] }.compact
+        context unless context.empty?
+      end
     end
   end
 end

data/lib/ruby_llm/agent/callbacks.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # The Agent callback macros, such as Agent::after_message, which
+    # register the Chat callback of the same name on every chat an agent
+    # builds or finds. Each runs on the agent instance, like a Rails
+    # callback on its record. An agent that configures a chat again, as
+    # Agent.new(chat:) does for a chat from Agent::chat, takes over the
+    # callbacks already registered instead of adding them twice. As for
+    # Agent::rescue_from handlers, the agent behind a chat from Agent::chat
+    # or Agent::find is bound to it without running +initialize+.
+    module Callbacks # :nodoc:
+      HOOKS = %i[
+        before_request before_message after_message before_tool_call around_tool_call after_tool_result
+        after_tool_progress before_fallback after_fallback after_file_write
+      ].freeze
+
+      # A declared callback: a method name or a block, with the +if:+ and
+      # +unless:+ conditions that guard it.
+      class Callback # :nodoc:
+        def initialize(hook, method_name, block, conditions)
+          unless method_name.nil? ^ block.nil?
+            raise ArgumentError, "#{hook} takes a method name or a block, not both or neither"
+          end
+
+          unknown = conditions.keys - %i[if unless]
+          raise ArgumentError, "#{hook} takes if: and unless:, got #{unknown.join(', ')}" if unknown.any?
+
+          @handler = method_name || block
+          @if = Array(conditions[:if])
+          @unless = Array(conditions[:unless])
+        end
+
+        # Runs the handler with +args+, passing +block+ on to it, as an
+        # around callback's +run+. A skipped around callback runs +block+.
+        def run(agent, args, &block)
+          return block&.call unless @if.all? { |condition| call(agent, condition, args) }
+          return block&.call if @unless.any? { |condition| call(agent, condition, args) }
+
+          call(agent, @handler, args, &block)
+        end
+
+        private
+
+        def call(agent, handler, args, &block)
+          if handler.is_a?(Symbol) || handler.is_a?(String)
+            method = agent.method(handler)
+            return method.arity.zero? ? method.call(&block) : method.call(*args, &block)
+          end
+          if block
+            method = bound(handler)
+            return method.arity.zero? ? method.bind_call(agent, &block) : method.bind_call(agent, *args, &block)
+          end
+
+          handler.lambda? && handler.arity.zero? ? agent.instance_exec(&handler) : agent.instance_exec(*args, &handler)
+        end
+
+        # instance_exec cannot pass a block on, so a handler that takes one
+        # runs as a method of the agent instead.
+        def bound(handler)
+          @bound ||= Module.new { define_method(:callback, &handler) }.instance_method(:callback)
+        end
+      end
+
+      HOOKS.each do |hook|
+        define_method(hook) do |method_name = nil, **conditions, &block|
+          chat_callbacks << [hook, Callback.new(hook, method_name, block, conditions)]
+          self
+        end
+      end
+
+      def chat_callbacks
+        @chat_callbacks ||= []
+      end
+
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@chat_callbacks, chat_callbacks.dup)
+      end
+
+      def apply_configuration(chat, input_values:, agent: nil, **)
+        super.tap do
+          next if chat_callbacks.empty?
+
+          receivers = (chat.respond_to?(:to_llm) ? chat.to_llm : chat).agent_callback_receivers
+          registered = receivers.key?(self)
+          receivers[self] = agent || allocate.bind(chat, input_values)
+          next if registered
+
+          chat_callbacks.each do |hook, callback|
+            chat.public_send(hook) { |*args, &run| callback.run(receivers[self], args, &run) }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/agent/context_budget.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # Applies the Agent::context_budget macro to every chat an agent builds
+    # or finds.
+    module ContextBudget # :nodoc:
+      def apply_configuration(chat, input_values:, **)
+        super.tap do
+          next if @context_budget.nil?
+
+          chat.with_context_budget(evaluate(@context_budget, runtime_context(chat:, inputs: input_values)) || nil)
+        end
+      end
+
+      private
+
+      def context_budget_option(budget, **shares, &block)
+        return Chat::ContextBudget.normalize(budget, **shares) || false unless block
+        raise ArgumentError, 'context_budget takes options or a block, not both' if shares.any? || budget != {}
+
+        block
+      end
+    end
+  end
+end

data/lib/ruby_llm/agent/join_callbacks.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # The Agent::around_join macro, which registers Chat#around_join
+    # callbacks on every chat an agent builds or finds.
+    module JoinCallbacks # :nodoc:
+      def around_join(&block)
+        raise ArgumentError, 'around_join needs a block' unless block
+
+        around_join_callbacks << block
+      end
+
+      def around_join_callbacks
+        @around_join_callbacks ||= []
+      end
+
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@around_join_callbacks, around_join_callbacks.dup)
+      end
+
+      def apply_configuration(chat, **)
+        super.tap { around_join_callbacks.each { |callback| chat.around_join(&callback) } }
+      end
+    end
+  end
+end

data/lib/ruby_llm/agent/model_options.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # Decides what the Agent::model macro's options apply to. They belong to
+    # the declared model, so a caller's +model:+ replaces them along with it,
+    # and a record on another model does not take them. A record being created
+    # knows what its caller passed; a later one is judged by its saved model.
+    module ModelOptions # :nodoc:
+      DECLARED_WITH_MODEL = %i[provider protocol assume_model_exists].freeze
+
+      private
+
+      def resolved_chat_kwargs(inputs:, options:)
+        replaced = options.key?(:model) && chat_kwargs.key?(:model)
+        return chat_kwargs.except(*DECLARED_WITH_MODEL).merge(options) if replaced
+        return chat_kwargs.merge(options) unless chat_kwargs[:model].is_a?(Proc)
+
+        chat_kwargs.merge(model: evaluate(chat_kwargs[:model], runtime_context(chat: nil, inputs:))).merge(options)
+      end
+
+      def apply_chat_options(chat, options = {})
+        return if options.key?(:model) || !rails_chat_record?(chat) || !declared_model_on?(chat)
+
+        declared = chat_kwargs.except(*options.keys)
+        chat.assume_model_exists = declared[:assume_model_exists] if declared.key?(:assume_model_exists)
+        chat.protocol ||= declared[:protocol] if declared.key?(:protocol)
+      end
+
+      def declared_model_on?(record)
+        return true unless record.model_id
+
+        declared_provider_on?(record) && declared_id_on?(record)
+      end
+
+      def declared_provider_on?(record)
+        declared = chat_kwargs[:provider]
+
+        declared.nil? || record.provider.nil? || declared.to_s == record.provider.to_s
+      end
+
+      def declared_id_on?(record)
+        declared = chat_kwargs[:model]
+        declared = declared.id if declared.is_a?(Model)
+        return true if declared.nil? || declared.is_a?(Proc) || declared == record.model_id
+
+        provider = record.provider || chat_kwargs[:provider]
+        return true if Models::Aliases.resolve(declared, provider) == record.model_id
+
+        Models.find(declared, provider:).id == record.model_id
+      rescue ModelNotFoundError, ConfigurationError
+        false
+      end
+    end
+  end
+end

data/lib/ruby_llm/agent/ownership.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # Applies the Agent::owner macro's owner to every chat an agent builds
+    # or finds, with Chat#with_owner, and hands it to subclasses. The owner
+    # is set once the chat has its context, and before the instructions
+    # and tools that may read it.
+    module Ownership # :nodoc:
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@owner, @owner)
+      end
+
+      private
+
+      def apply_context(chat, runtime, **)
+        super.tap { chat.with_owner(evaluate(owner, runtime)) unless owner.nil? }
+      end
+    end
+  end
+end

data/lib/ruby_llm/agent/tool_conversion.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Agent
+    # The Agent::description and Agent::to_tool macros, which let an agent
+    # class go wherever tools go, as an AgentTool.
+    module ToolConversion # :nodoc:
+      def description(text = nil)
+        return @description unless text
+
+        @description = text
+      end
+
+      def to_tool(**inputs)
+        AgentTool.for(self, inputs).new(inputs)
+      end
+    end
+  end
+end