ruby_llm-agents 3.8.0 → 3.10.0

data/lib/ruby_llm/agents/results/speech_result.rb

@@ -23,6 +23,8 @@ module RubyLLM
     #
     # @api public
     class SpeechResult
+      include Trackable
+
       # @!group Audio Content
 
       # @!attribute [r] audio
@@ -229,6 +231,10 @@ module RubyLLM
 
         # Execution record
         @execution_id = attributes[:execution_id]
+
+        # Tracking
+        @agent_class_name = attributes[:agent_class_name]
+        register_with_tracker
       end
 
       # Loads the associated Execution record from the database

data/lib/ruby_llm/agents/results/trackable.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # Mixin that registers a result object with the active Tracker.
+    #
+    # Included in every result class so that RubyLLM::Agents.track
+    # can collect results automatically.
+    #
+    # @api private
+    module Trackable
+      def self.included(base)
+        base.attr_reader :agent_class_name unless base.method_defined?(:agent_class_name)
+      end
+
+      private
+
+      # Call from the end of initialize to register with the active tracker
+      def register_with_tracker
+        tracker = Thread.current[:ruby_llm_agents_tracker]
+        tracker << self if tracker
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/results/transcription_result.rb

@@ -26,6 +26,8 @@ module RubyLLM
     #
     # @api public
     class TranscriptionResult
+      include Trackable
+
       # @!group Content
 
       # @!attribute [r] text
@@ -250,6 +252,10 @@ module RubyLLM
 
         # Execution record
         @execution_id = attributes[:execution_id]
+
+        # Tracking
+        @agent_class_name = attributes[:agent_class_name]
+        register_with_tracker
       end
 
       # Loads the associated Execution record from the database

data/lib/ruby_llm/agents/text/embedder.rb

@@ -337,11 +337,14 @@ module RubyLLM
         embed_options = {model: context&.model || resolved_model}
         embed_options[:dimensions] = resolved_dimensions if resolved_dimensions
 
-        # Pass scoped RubyLLM context for thread-safe per-tenant API keys
+        # Use scoped RubyLLM::Context for thread-safe per-tenant API keys.
+        # RubyLLM::Context#embed creates an Embedding with the scoped config.
         llm_ctx = context&.llm
-        embed_options[:context] = llm_ctx if llm_ctx.is_a?(RubyLLM::Context)
-
-        response = RubyLLM.embed(preprocessed, **embed_options)
+        response = if llm_ctx.is_a?(RubyLLM::Context)
+          llm_ctx.embed(preprocessed, **embed_options)
+        else
+          RubyLLM.embed(preprocessed, **embed_options)
+        end
 
         # ruby_llm returns vectors as an array (even for single text)
         vectors = response.vectors

data/lib/ruby_llm/agents/tool.rb

@@ -0,0 +1,169 @@
+# frozen_string_literal: true
+
+require "timeout"
+
+module RubyLLM
+  module Agents
+    # Base class for tools that need access to the agent's execution context.
+    #
+    # Inherits from RubyLLM::Tool and adds:
+    # - `context` accessor: read agent params, tenant, execution ID
+    # - `timeout` DSL: per-tool timeout in seconds
+    # - Error handling: exceptions become error strings for the LLM
+    # - Tool execution tracking: records each tool call in the database
+    #
+    # Users implement `execute()` — the standard RubyLLM convention.
+    # This class overrides `call()` to wrap execution with its features.
+    #
+    # @example Defining a tool
+    #   class BashTool < RubyLLM::Agents::Tool
+    #     description "Run a shell command"
+    #     timeout 30
+    #
+    #     param :command, desc: "The command to run", required: true
+    #
+    #     def execute(command:)
+    #       context.container_id  # reads agent param
+    #       # ... run command ...
+    #     end
+    #   end
+    #
+    # @example Using with an agent
+    #   class CodingAgent < ApplicationAgent
+    #     param :container_id, required: true
+    #     tools [BashTool]
+    #   end
+    #
+    #   CodingAgent.call(query: "list files", container_id: "abc123")
+    #
+    class Tool < RubyLLM::Tool
+      # The execution context, set before each call.
+      # Provides access to agent params, tenant, execution ID.
+      #
+      # @return [ToolContext, nil]
+      attr_reader :context
+
+      class << self
+        # Sets or gets the per-tool timeout in seconds.
+        #
+        # @param value [Integer, nil] Timeout in seconds (setter)
+        # @return [Integer, nil] The configured timeout (getter)
+        def timeout(value = nil)
+          if value
+            @timeout = value
+          else
+            @timeout
+          end
+        end
+      end
+
+      # Wraps RubyLLM's call() with context, timeout, tracking, and error handling.
+      #
+      # RubyLLM's Chat calls tool.call(args) during the tool loop.
+      # We set up context, create a tracking record, apply timeout,
+      # then delegate to super (which validates args and calls execute).
+      #
+      # @param args [Hash] Tool arguments from the LLM
+      # @return [String, Tool::Halt] The tool result or a Halt signal
+      def call(args)
+        pipeline_context = Thread.current[:ruby_llm_agents_caller_context]
+        @context = pipeline_context ? ToolContext.new(pipeline_context) : nil
+
+        record = start_tool_tracking(pipeline_context, args)
+
+        check_cancelled!(pipeline_context)
+
+        timeout_seconds = self.class.timeout
+        timeout_seconds ||= RubyLLM::Agents.configuration.default_tool_timeout
+
+        result = if timeout_seconds
+          Timeout.timeout(timeout_seconds) { super }
+        else
+          super
+        end
+
+        complete_tool_tracking(record, result, status: "success")
+        result
+      rescue Timeout::Error
+        complete_tool_tracking(record, nil, status: "timed_out", error: "Timed out after #{timeout_seconds}s")
+        "TIMEOUT: Tool did not complete within #{timeout_seconds}s."
+      rescue RubyLLM::Agents::CancelledError
+        complete_tool_tracking(record, nil, status: "cancelled")
+        raise # Let cancellation propagate to BaseAgent
+      rescue => e
+        complete_tool_tracking(record, nil, status: "error", error: e.message)
+        "ERROR (#{e.class}): #{e.message}"
+      end
+
+      private
+
+      # Creates a "running" ToolExecution record before the tool runs.
+      # Silently skips if no execution_id or ToolExecution is not available.
+      #
+      # @param pipeline_context [Pipeline::Context, nil]
+      # @param args [Hash] The tool arguments
+      # @return [ToolExecution, nil]
+      def start_tool_tracking(pipeline_context, args)
+        return nil unless pipeline_context&.execution_id
+        return nil unless defined?(ToolExecution)
+
+        @tool_iteration = (@tool_iteration || 0) + 1
+
+        ToolExecution.create!(
+          execution_id: pipeline_context.execution_id,
+          tool_name: name,
+          iteration: @tool_iteration,
+          status: "running",
+          input: normalize_input(args),
+          started_at: Time.current
+        )
+      rescue => e
+        # Don't let tracking failures break tool execution
+        Rails.logger.debug("[RubyLLM::Agents::Tool] Tracking failed: #{e.message}") if defined?(Rails) && Rails.logger
+        nil
+      end
+
+      # Updates the ToolExecution record after the tool completes.
+      #
+      # @param record [ToolExecution, nil]
+      # @param result [Object, nil] The tool result
+      # @param status [String] Final status
+      # @param error [String, nil] Error message
+      def complete_tool_tracking(record, result, status:, error: nil)
+        return unless record
+
+        completed_at = Time.current
+        duration_ms = record.started_at ? ((completed_at - record.started_at) * 1000).to_i : nil
+        output_str = result.is_a?(RubyLLM::Tool::Halt) ? result.content.to_s : result.to_s
+
+        record.update!(
+          status: status,
+          output: truncate_output(output_str),
+          output_bytes: output_str.bytesize,
+          error_message: error,
+          completed_at: completed_at,
+          duration_ms: duration_ms
+        )
+      rescue => e
+        Rails.logger.debug("[RubyLLM::Agents::Tool] Tracking update failed: #{e.message}") if defined?(Rails) && Rails.logger
+      end
+
+      def check_cancelled!(pipeline_context)
+        return unless pipeline_context
+        on_cancelled = pipeline_context[:on_cancelled]
+        return unless on_cancelled.respond_to?(:call)
+        raise CancelledError, "Execution cancelled" if on_cancelled.call
+      end
+
+      def normalize_input(args)
+        return {} if args.nil?
+        args.respond_to?(:to_h) ? args.to_h : {}
+      end
+
+      def truncate_output(str)
+        max = RubyLLM::Agents.configuration.try(:tool_result_max_length) || 10_000
+        (str.length > max) ? str[0, max] + "... (truncated)" : str
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/tool_context.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # Read-only wrapper around Pipeline::Context for tool authors.
+    #
+    # Exposes agent params and execution metadata to tools without
+    # leaking pipeline internals. Supports both method-style and
+    # hash-style access to agent params.
+    #
+    # @example Method-style access
+    #   context.container_id   # reads agent param
+    #   context.tenant_id      # fixed attribute
+    #
+    # @example Hash-style access
+    #   context[:container_id]
+    #
+    class ToolContext
+      # Execution record ID — links tool calls to the agent execution
+      #
+      # @return [Integer, nil]
+      def id
+        @ctx.execution_id
+      end
+
+      # Tenant ID from the pipeline context
+      #
+      # @return [String, nil]
+      def tenant_id
+        @ctx.tenant_id
+      end
+
+      # Agent class name
+      #
+      # @return [String, nil]
+      def agent_type
+        @ctx.agent_class&.name
+      end
+
+      # Hash-style access to agent params
+      #
+      # @param key [Symbol, String] The param key
+      # @return [Object, nil]
+      def [](key)
+        @agent_options[key.to_sym] || @agent_options[key.to_s]
+      end
+
+      def initialize(pipeline_context)
+        @ctx = pipeline_context
+        @agent_options = @ctx.agent_instance&.send(:options) || {}
+      end
+
+      private
+
+      # Method-style access to agent params
+      def method_missing(method_name, *args)
+        key = method_name.to_sym
+        if @agent_options.key?(key) || @agent_options.key?(key.to_s)
+          self[key]
+        else
+          super
+        end
+      end
+
+      def respond_to_missing?(method_name, include_private = false)
+        key = method_name.to_sym
+        @agent_options.key?(key) || @agent_options.key?(key.to_s) || super
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/track_report.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # Aggregated read-only report returned by RubyLLM::Agents.track.
+    #
+    # Provides totals and breakdowns across all agent calls made
+    # inside the tracked block.
+    #
+    # @example
+    #   report = RubyLLM::Agents.track do
+    #     TranscribeAgent.call(with: audio_path)
+    #     ChatAgent.call(message: "hello")
+    #   end
+    #   report.total_cost   # => 0.0078
+    #   report.call_count   # => 2
+    #
+    # @api public
+    class TrackReport
+      attr_reader :value, :error, :results, :request_id
+      attr_reader :started_at, :completed_at
+
+      def initialize(value:, error:, results:, request_id:, started_at:, completed_at:)
+        @value = value
+        @error = error
+        @results = results.freeze
+        @request_id = request_id
+        @started_at = started_at
+        @completed_at = completed_at
+      end
+
+      def successful?
+        @error.nil?
+      end
+
+      def failed?
+        !successful?
+      end
+
+      def call_count
+        @results.size
+      end
+
+      def total_cost
+        @results.sum { |r| r.total_cost || 0 }
+      end
+
+      def input_cost
+        @results.sum { |r| r.input_cost || 0 }
+      end
+
+      def output_cost
+        @results.sum { |r| r.output_cost || 0 }
+      end
+
+      def total_tokens
+        @results.sum { |r| r.total_tokens }
+      end
+
+      def input_tokens
+        @results.sum { |r| r.input_tokens || 0 }
+      end
+
+      def output_tokens
+        @results.sum { |r| r.output_tokens || 0 }
+      end
+
+      def duration_ms
+        return nil unless @started_at && @completed_at
+        ((@completed_at - @started_at) * 1000).to_i
+      end
+
+      def all_successful?
+        @results.all?(&:success?)
+      end
+
+      def any_errors?
+        @results.any?(&:error?)
+      end
+
+      def errors
+        @results.select(&:error?)
+      end
+
+      def successful
+        @results.select(&:success?)
+      end
+
+      def models_used
+        @results.filter_map(&:chosen_model_id).uniq
+      end
+
+      def cost_breakdown
+        @results.map do |r|
+          {
+            agent: r.respond_to?(:agent_class_name) ? r.agent_class_name : nil,
+            model: r.chosen_model_id,
+            cost: r.total_cost || 0,
+            tokens: r.total_tokens,
+            duration_ms: r.duration_ms
+          }
+        end
+      end
+
+      def to_h
+        {
+          successful: successful?,
+          value: value,
+          error: error&.message,
+          request_id: request_id,
+          call_count: call_count,
+          total_cost: total_cost,
+          input_cost: input_cost,
+          output_cost: output_cost,
+          total_tokens: total_tokens,
+          input_tokens: input_tokens,
+          output_tokens: output_tokens,
+          duration_ms: duration_ms,
+          started_at: started_at,
+          completed_at: completed_at,
+          models_used: models_used,
+          cost_breakdown: cost_breakdown
+        }
+      end
+    end
+  end
+end

data/lib/ruby_llm/agents/tracker.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Agents
+    # Internal collector used by RubyLLM::Agents.track to accumulate
+    # Result objects produced during a tracked block.
+    #
+    # Not part of the public API — users interact with TrackReport instead.
+    #
+    # @api private
+    class Tracker
+      attr_reader :results, :defaults, :request_id, :tags
+
+      def initialize(defaults: {}, request_id: nil, tags: {})
+        @results = []
+        @defaults = defaults
+        @request_id = request_id || generate_request_id
+        @tags = tags
+      end
+
+      def <<(result)
+        @results << result
+      end
+
+      private
+
+      def generate_request_id
+        "track_#{SecureRandom.hex(8)}"
+      end
+    end
+  end
+end