ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/batch.rb

@@ -95,6 +95,11 @@ module RubyLLM
         new(provider:, store: config.batch_store, **provider.find_batch(id))
       end
 
+      def awaiting_response?(chat) # :nodoc:
+        last = chat.messages.reverse.find { |message| message.role != :system }
+        AWAITING_ROLES.include?(last&.role)
+      end
+
       private
 
       def submit_chats(records)
@@ -104,7 +109,7 @@ module RubyLLM
         payload = { provider: provider.slug, provider_class: provider.name, requests: chats.size }
         RubyLLM.instrument('batch.ruby_llm', payload, config: provider.config) do |event|
           requests = chats.each_with_index.map do |chat, index|
-            { custom_id: index.to_s, model: chat.model.id, payload: chat.render }
+            { custom_id: index.to_s, model: chat.model.id, payload: chat.render, headers: chat.headers }
           end
           store = provider.config.batch_store
           batch = new(provider:, chats:, store:, **provider.create_batch(requests))
@@ -153,7 +158,7 @@ module RubyLLM
       end
 
       def awaiting_model?(chat)
-        !chat.complete? && AWAITING_ROLES.include?(chat.messages.last&.role)
+        !chat.complete? && awaiting_response?(chat)
       end
 
       def shared_provider(chats)
@@ -161,19 +166,25 @@ module RubyLLM
         raise ArgumentError, "A batch takes one provider per submission, got: #{slugs.join(', ')}" if slugs.size > 1
 
         provider = chats.first.provider
+        unless chats.all? { |chat| chat.provider.same_credentials?(provider) }
+          raise ArgumentError, 'A batch takes one set of credentials per submission'
+        end
+
         raise Error, "#{provider.slug} doesn't support batch requests" unless provider.batches?
 
         provider
       end
     end
 
-    def initialize(provider:, chats: nil, requests: nil, batch_protocol: nil, store: nil, **attributes) # :nodoc:
+    def initialize(provider:, chats: nil, requests: nil, batch_protocol: nil, store: nil, delivered: [], # :nodoc:
+                   delivery: nil, **attributes)
       @provider = provider
       @chats = chats
       @requests = requests
       @batch_protocol = protocol_name(batch_protocol)
       @store = store
-      @delivered = {}
+      @delivered = delivered.to_h { |index| [index, true] }
+      @delivery = delivery
       @statuses = []
       apply(attributes)
     end
@@ -241,8 +252,9 @@ module RubyLLM
     def messages
       return @messages if @messages
 
+      @deferred = false
       collected = collect_results
-      @messages = collected if @completed
+      @messages = collected if @completed && !@deferred
       collected
     end
 
@@ -291,9 +303,11 @@ module RubyLLM
       validate_result_indices(results)
       slots = Array.new(result_slot_count(results))
 
-      results.each do |index, result, failure_status|
+      results.each do |index, result, failure_status, billed|
+        result = Chat::ToolCallIds.new(namespace: [@provider.slug, id, index]).assign(result) if result.is_a?(Message)
         slots[index] = result
         deliver(index, result, failure_status)
+        deliver_failed_usage(index, **billed) if billed
       end
 
       fill_missing_statuses(slots.size) if complete?
@@ -325,7 +339,9 @@ module RubyLLM
     end
 
     # Collecting early keeps reading fresh, so a result already delivered
-    # comes back on every later poll: hand each one over once.
+    # comes back on every later poll: hand each one over once. A result with
+    # no chat or request waits for completion, because a poller that finds
+    # the batch by id gets a new object each time.
     def deliver(index, result, failure_status)
       statuses[index] = result ? :succeeded : failure_status
       return unless result
@@ -339,69 +355,119 @@ module RubyLLM
 
     def deliver_embedding(index, embedding)
       request = requests&.[](index)
-      delivered = @delivered[index]
-      attach_batch_usage(
-        embedding,
-        operation: :embedding,
-        model: request&.model,
-        category: :embeddings,
-        instrument: !request.nil? && !delivered
-      )
-      return if delivered
+      entry = attach_batch_usage(embedding, operation: :embedding, model: request&.model, category: :embeddings)
+      return if settled?(index, request)
 
       @delivered[index] = true
+      report_batch_usage(entry) if entry
       request&.result = embedding
     end
 
     def deliver_message(index, message)
       chat = chats&.[](index)
-      delivered = @delivered[index] || (chat && already_in_chat?(chat, message))
-      attach_batch_usage(
-        message,
-        operation: :chat,
-        model: chat&.model,
-        category: :text_tokens,
-        instrument: !chat.nil? && !delivered
-      )
-      return if delivered
+      entry = attach_batch_usage(message, operation: :chat, model: chat&.model, category: :text_tokens,
+                                          owner: chat&.owner)
+      return if settled?(index, chat) || (chat && already_in_chat?(chat, message))
+
+      delivered = exclusively(index) do
+        @delivered[index] = true
+        report_batch_usage(entry, chat:) if entry
+        chat&.add_completion(message, record_usage: true)
+      end
+      delivered ? @delivered[index] = true : @deferred = true
+    end
 
-      @delivered[index] = true
-      chat&.add_completion(message, record_usage: true)
+    # A slot RubyLLM could not parse still cost what the provider reported.
+    def deliver_failed_usage(index, tokens:, model:)
+      chat = chats&.[](index)
+      entry = failed_usage_entry(tokens, model, chat)
+      return if entry.nil? || settled?(index, chat)
+
+      delivered = exclusively(index, failed: true) do
+        report_batch_usage(entry, chat:)
+        chat&.add_usage_entry(entry)
+      end
+      delivered ? @delivered[index] = true : @deferred = true
+    end
+
+    def settled?(index, owner)
+      @delivered[index] || !(owner || complete?)
+    end
+
+    # A store may run this in its own turn, returning false when it could
+    # not, so the next poll delivers the result.
+    def exclusively(index, **, &deliver)
+      return @delivery.call(index, **, &deliver) if @delivery
+
+      deliver.call
+      true
     end
 
-    def attach_batch_usage(result, operation:, model:, category:, instrument:)
+    def attach_batch_usage(result, operation:, model:, category:, owner: nil)
       return unless result.ruby_llm_usage_entries.empty?
 
-      model ||= RubyLLM.models.find(result.model, provider: @provider.slug, config: @provider.config)
-      entry = Accounting::Usage::Entry.new(
+      model ||= result_model(result.model)
+      entry = batch_usage_entry(
+        operation:, model_id: result.model || model.id, model:, status: :succeeded, tokens: result.tokens,
+        category:, message: result.is_a?(Message) ? result : nil, owner:
+      )
+      result.ruby_llm_usage_entries = [entry]
+      entry
+    end
+
+    def failed_usage_entry(tokens, model_id, chat)
+      model = chat&.model || (result_model(model_id) if model_id)
+      model_id ||= model&.id
+      return unless model_id
+
+      batch_usage_entry(operation: :chat, model_id:, model:, status: :failed, tokens:, category: :text_tokens,
+                        owner: chat&.owner)
+    end
+
+    def batch_usage_entry(operation:, model_id:, model:, status:, tokens:, category:, message: nil, owner: nil)
+      Accounting::Usage::Entry.new(
         operation:,
         provider: @provider.slug,
-        model: result.model || model.id,
-        status: :succeeded,
-        tokens: result.tokens,
-        cost: @provider.batch_cost(result.tokens, model:, category:),
-        message: result.is_a?(Message) ? result : nil,
-        owner: Accounting::Usage.owner
+        model: model_id,
+        status:,
+        tokens:,
+        cost: @provider.batch_cost(tokens, model:, category:),
+        message:,
+        owner: owner || Accounting::Usage.owner,
+        workflow: Accounting::Usage.workflow
       )
-      result.ruby_llm_usage_entries = [entry]
-      return unless instrument
+    end
 
-      if operation == :chat
+    # A chat writes its own ledger row when it takes the answer. Anything
+    # else, such as a batch picked up by id, gets a row with no chat.
+    def report_batch_usage(entry, chat: nil)
+      if chat
         Accounting::Usage.instrument(entry, config: @provider.config)
       else
         Accounting::Usage.report(entry, config: @provider.config)
       end
     end
 
-    # A plain answer is the chat's last message once it arrives. A tool-call
+    def result_model(model_id)
+      RubyLLM.models.find(model_id, provider: @provider.slug, config: @provider.config)
+    rescue ModelNotFoundError
+      nil
+    end
+
+    # A plain answer is the chat's last non-system message once it arrives. A tool-call
     # answer is not: running its tools adds messages after it, so we match on its
     # tool-call ids instead.
     def already_in_chat?(chat, message)
-      if message.tool_call?
-        chat.messages.any? { |m| m.tool_call? && m.tool_calls.keys.intersect?(message.tool_calls.keys) }
-      else
-        !AWAITING_ROLES.include?(chat.messages.last&.role)
-      end
+      return !self.class.awaiting_response?(chat) unless message.tool_call?
+
+      tool_calls_delivered?(chat, message)
+    end
+
+    def tool_calls_delivered?(chat, message)
+      calls = chat.messages.flat_map { |existing| existing.tool_calls.to_a }
+      legacy_ids = calls.select { |call| call.provider_call_id == call.id }.map(&:id)
+      calls.map(&:id).intersect?(message.tool_calls.keys) ||
+        legacy_ids.intersect?(message.tool_calls.map(&:provider_call_id))
     end
 
     module Helpers # :nodoc:
@@ -416,6 +482,25 @@ module RubyLLM
         status.to_s.match?(/cancel/i) ? :cancelled : :failed
       end
 
+      # One unreadable answer fails its own slot, so the others still arrive.
+      # Pass a chat answer's +body+ to keep the usage the provider billed.
+      def parsed_batch_result(index, custom_id, body = nil, parser: self)
+        [index, yield]
+      rescue Error => e
+        [index, nil, batch_failure(custom_id, e.message), *billed_batch_usage(body, parser)]
+      end
+
+      def billed_batch_usage(body, parser)
+        return [] unless body
+
+        tokens = Message.new(role: :assistant, content: nil, **parser.send(:parse_completion_usage, body)).tokens
+        tokens.to_h.any? ? [{ tokens:, model: body['model'] }] : []
+      end
+
+      def batch_response_succeeded?(response)
+        response.is_a?(Hash) && response['status_code'].to_i.between?(200, 299)
+      end
+
       def batch_error_message(line)
         response = line['response']
         body = response['body'] if response.is_a?(Hash)
@@ -425,7 +510,8 @@ module RubyLLM
         batch_error_value(line['error']) ||
           line['error_message'] ||
           batch_error_value(body_error) ||
-          batch_error_value(response_error)
+          batch_error_value(response_error) ||
+          batch_error_value(body)
       end
 
       def batch_error_value(error)

data/lib/ruby_llm/cached_content.rb

@@ -60,11 +60,11 @@ module RubyLLM
       self
     end
 
-    # Extends the cache's lifetime to +ttl:+ seconds from now, given as an
-    # Integer or a provider duration string such as <tt>"600s"</tt>.
+    # Extends the cache's lifetime to +ttl:+ from now, given as seconds, a
+    # Duration, or a provider duration string such as <tt>"600s"</tt>.
     # Updates #expires_at and returns +self+.
     #
-    #   cache.renew(ttl: 3600)
+    #   cache.renew(ttl: 1.hour)
     #
     def renew(ttl:)
       refreshed = @provider_instance.extend_cache(name, ttl: ttl)
@@ -81,7 +81,7 @@ module RubyLLM
     # +content+ is the text to cache; pass file attachments with +with:+
     # the way Chat#ask accepts them. +instructions:+ caches a system
     # prompt alongside the content, and +ttl:+ sets the cache lifetime in
-    # seconds (Integer) or as a provider duration string such as
+    # seconds, as a Duration, or as a provider duration string such as
     # <tt>"300s"</tt>. When +provider:+ is omitted, the model's default
     # provider is used. The content must exceed the model's minimum
     # cacheable token count.

data/lib/ruby_llm/chat/agent_run.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Wraps a run of a chat an Agent configured in one +agent.ruby_llm+
+    # event, whether it is reached through the chat, the agent, or a
+    # persisted record. Only the outermost operation of a run opens one, so
+    # #ask, #complete, and the moves they make report a single event.
+    module AgentRun # :nodoc:
+      OPERATIONS = %i[complete generate run_tools step].freeze
+
+      OPERATIONS.each do |operation|
+        define_method(operation) do |*args, **kwargs, &block|
+          chat = agent_run_chat
+          return super(*args, **kwargs, &block) unless chat
+
+          chat.running_agent { super(*args, **kwargs, &block) }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/context_budget.rb

@@ -0,0 +1,119 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Fits tool results to shares of the model's context window, as set with
+    # Chat#with_context_budget. Counting tokens takes a request, so sizes go
+    # by characters, at two a token: JSON and HTML, which most large tool
+    # results are, run that dense.
+    class ContextBudget # :nodoc:
+      SHARES = %i[per_result per_request].freeze
+      CHARACTERS_PER_TOKEN = 2
+      # The cut moves in steps of this share of the request budget, so
+      # the request prefix, and the prompt cache, hold for many rounds.
+      STEP = 0.15
+      LEFT_OUT = '[This earlier tool result was left out to fit the context window. ' \
+                 'Call the tool again if you need it.]'
+      TOO_LONG = '[This tool result was left out to fit the context window. ' \
+                 'Ask for less, such as a narrower query or the next page.]'
+      NOTE_ROOM = 200
+
+      def self.normalize(budget, **shares)
+        return if budget.nil?
+
+        budget = budget.to_h.transform_keys(&:to_sym).merge(shares)
+        unknown = budget.keys - SHARES - [:window]
+        raise ArgumentError, "Unknown context budget option: #{unknown.join(', ')}" if unknown.any?
+        unless budget.keys.intersect?(SHARES)
+          raise ArgumentError, 'with_context_budget takes per_result:, per_request:, or both'
+        end
+
+        budget.slice(*SHARES).each { |name, share| validate_share(name, share) }
+        validate_window(budget[:window]) if budget.key?(:window)
+        budget.freeze
+      end
+
+      def self.validate_share(name, share)
+        return if share.is_a?(Numeric) && share.positive? && share <= 1
+
+        raise ArgumentError, "#{name} must be a share of the context window above 0 and at most 1"
+      end
+
+      def self.validate_window(window)
+        return if window.is_a?(Integer) && window.positive?
+
+        raise ArgumentError, 'window must be a positive whole number of tokens'
+      end
+      private_class_method :validate_share, :validate_window
+
+      def initialize(budget, context_window)
+        @budget = budget
+        window = budget[:window] || context_window
+        @window = window.to_i if window.to_i.positive?
+      end
+
+      def window?
+        !@window.nil?
+      end
+
+      def cap(content)
+        limit = characters(:per_result)
+        return content unless limit && content.is_a?(String) && content.length > limit
+
+        shown = [limit - NOTE_ROOM, 0].max
+        "#{content[0, shown]}\n\n[Result truncated: #{shown} of #{content.length} characters shown. " \
+          'Ask for a narrower query or the next page.]'
+      end
+
+      def fit(list)
+        limit = characters(:per_request)
+        return list unless limit
+
+        excess = size_of_all(list) - limit
+        return list unless excess.positive?
+
+        step = limit * STEP
+        latest = list.rindex { |message| message.role == :assistant } || -1
+        fitted = leave_out(list, (excess / step).ceil * step, LEFT_OUT) { |index| index < latest }
+        over = size_of_all(fitted) - limit
+        return fitted unless over.positive?
+
+        leave_out(fitted, over, TOO_LONG) { |index| index > latest }
+      end
+
+      private
+
+      # Replaces the oldest results the block picks by index with +note+,
+      # until +characters+ are freed. A request stops here for older rounds,
+      # in steps. The latest round's results, which the model has not seen
+      # yet, are cut only as far as the request needs: left whole, a round
+      # over the budget would fail this request and every one after it.
+      def leave_out(list, characters, note)
+        list.each_with_index.map do |message, index|
+          next message unless characters.positive? && yield(index) && leavable?(message, note)
+
+          characters -= message.content.length - note.length
+          message.with(content: note)
+        end
+      end
+
+      def characters(share)
+        return unless @window && @budget[share]
+
+        (@window * @budget[share] * CHARACTERS_PER_TOKEN).to_i
+      end
+
+      def leavable?(message, note)
+        message.tool_result? && message.content.is_a?(String) && message.content.length > note.length
+      end
+
+      def size_of_all(list)
+        list.sum { |message| size_of(message) }
+      end
+
+      def size_of(message)
+        message.content.to_s.length + message.tool_calls.to_h.sum { |_id, call| call.arguments.to_json.length }
+      end
+    end
+  end
+end