ruby_llm 2.1.0 → 2.2.0

data/lib/ruby_llm/chat/file_references.rb

@@ -0,0 +1,259 @@
+# frozen_string_literal: true
+
+require 'json'
+require 'stringio'
+
+module RubyLLM
+  class Chat
+    # What Chat#with_file_store changes in a conversation. Every attachment is
+    # kept in the chat's FileStore, and each request shows it in full only
+    # while it is new: on the latest question, and on read_file results
+    # since then. Before that, and on other tools' results, the model sees
+    # one line naming the file, and reads it again with the file tools. A
+    # tool result longer than the inline limit is saved as a file, and
+    # the model gets its start and the path. An attachment the provider
+    # cannot take, such as a Word file, goes as the text the store's
+    # converters read from it, or as that line when none can. An attachment
+    # a converter on this machine reads for free has its text extracted
+    # when it is kept, so search_files finds it at once.
+    class FileReferences # :nodoc:
+      INLINE_LIMIT = 10_000
+      PREVIEW = 2_000
+      SPILLED = Regexp.new('\n\n\[This result has \d+ characters, so only the start is shown\. All of it is in ' \
+                           '(\S+): read more with read_file, or search it with search_files\.\]\z')
+      READ_TOOL = 'read_file'
+
+      # Stands for the store a chat picks when given none: a new
+      # FileStore::Memory, or a persisted chat's files in Active Storage.
+      DEFAULT_STORE = Object.new.tap { |default| def default.inspect = 'the default store' }.freeze
+
+      attr_reader :inline_limit, :tools, :store
+
+      class << self
+        def validate(store, inline_limit:, tools: FileStore::Tools::DEFAULT, store_lambda: false)
+          unless store?(store) || (store_lambda && store.is_a?(Proc))
+            raise ArgumentError, 'with_file_store takes a folder path, :memory, a RubyLLM::FileStore, or nil, ' \
+                                 "got #{store.inspect}"
+          end
+          FileStore::Tools.names(tools)
+          return if inline_limit == false || (inline_limit.is_a?(Integer) && inline_limit.positive?)
+
+          raise ArgumentError, 'inline_limit: takes a positive number of characters or false'
+        end
+
+        # Returns the path of the file holding the whole result whose start
+        # is +content+, or +nil+ when the result was sent whole.
+        def spilled_path(content)
+          content.to_s[SPILLED, 1]
+        end
+
+        def resolve(store)
+          return FileStore::Memory.new if store == :memory
+          return FileStore::Directory.new(store) if path?(store)
+
+          store
+        end
+
+        private
+
+        def store?(store)
+          store.is_a?(FileStore) || store.equal?(DEFAULT_STORE) || store == :memory || path?(store)
+        end
+
+        def path?(store)
+          store.is_a?(String) || store.respond_to?(:to_path)
+        end
+      end
+
+      def initialize(store:, inline_limit:, written:, tools: FileStore::Tools::DEFAULT, config: -> { RubyLLM.config })
+        @store = store
+        @config = config
+        @inline_limit = inline_limit
+        @tools = FileStore::Tools.names(tools)
+        @written = written
+        @stored = {}.compare_by_identity
+        @described = {}.compare_by_identity
+      end
+
+      def view(list)
+        current = list.rindex { |message| message.role == :user } || 0
+        reads = read_paths(list)
+        list.each_with_index.map do |message, index|
+          message.attachments.empty? ? message : view_message(message, current: index >= current, reads:)
+        end
+      end
+
+      def spill(tool, tool_call, content)
+        limit = tool.respond_to?(:inline_limit) ? tool.inline_limit : nil
+        limit = @inline_limit if limit.nil?
+        return content unless limit && content.is_a?(String) && content.length > limit
+
+        text, extension = readable(content)
+        file = write("results/#{safe_name(tool_call.name)}-#{safe_id(tool_call.id)}.#{extension}", text)
+        "#{text[0, PREVIEW]}\n\n[This result has #{text.length} characters, so only the start is shown. " \
+          "All of it is in #{file.path}: read more with read_file, or search it with search_files.]"
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not save a long tool result in the chat's files: #{e.class}" }
+        content
+      end
+
+      def store_attachment(attachment)
+        @stored[attachment] ||= store.store_attachment(attachment) do |file|
+          @written.call(file)
+          extract(file)
+        end
+      end
+
+      def describe(attachment, config: RubyLLM.config)
+        file = store_attachment(attachment)
+        (as_text(file, config) || reference(file.path, file)).tap { |replacement| @described[replacement] = true }
+      end
+
+      # Records a read for each text attachment in +list+, the messages a
+      # request sent, that went as itself and matches the file in the store.
+      # Anything else the provider got for it, a reference, converted text,
+      # or an upload, is a new object that was never stored here.
+      def record_shown(list)
+        list.flat_map(&:attachments).each do |attachment|
+          file = @stored[attachment]
+          record_read(attachment, file) if file && attachment.text?
+        end
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not record the files a request showed: #{e.class}" }
+      end
+
+      # Returns the StoredFiles of the attachments on +list+, the chat's
+      # messages, as they are now, keeping any not kept yet. Leaves out
+      # read_file results: they show a file to whoever read it, and a store
+      # with an owner may not show that file to this one.
+      def attached_files(list)
+        reads = read_paths(list)
+        list.reject { |message| message.role == :tool && reads.key?(message.tool_call_id) }.flat_map(&:attachments)
+            .filter_map { |attachment| attached_file(attachment) }.uniq(&:path)
+      end
+
+      # Removes the replacements #describe made from +replacements+, a
+      # chat's memo, so another store or limit describes them again.
+      def forget(replacements)
+        replacements.delete_if { |_attachment, replacement| @described.key?(replacement) }
+      end
+
+      private
+
+      def attached_file(attachment)
+        store[store_attachment(attachment).path] unless attachment.provider_file?
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not keep an attachment in the chat's files: #{e.class}" }
+        nil
+      end
+
+      def record_read(attachment, file)
+        seen = FileStore.digest(attachment.content.to_s.b)
+        current = store[file.path]
+        return unless current && seen == (current.digest || FileStore.digest(store.read(current.path)))
+
+        store.record_read(current.path, seen) unless store.read_digest(current.path) == seen
+      end
+
+      def extract(file)
+        return if file.text? || store.cached_text(file)
+
+        text = FileStore::Converters.convert(store.attachment(file.path, file:), config: @config.call, local: true)
+        store.cache_text(file, text) if text
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not extract the text of #{file.path}: #{e.class}" }
+      end
+
+      def view_message(message, current:, reads:)
+        read = (reads[message.tool_call_id] if message.tool_result?)
+        shown = current && (!message.tool_result? || read)
+        attachments = message.attachments.map { |attachment| view_of(attachment, shown:, read:) }
+        attachments == message.attachments ? message : message.with_attachments(attachments)
+      end
+
+      def view_of(attachment, shown:, read:)
+        return attachment if attachment.provider_file?
+        return (shown ? attachment : reference(read, nil)) if read
+
+        file = store_attachment(attachment)
+        return reference(file.path, file) unless shown
+
+        too_long?(attachment) ? start_of(attachment, file) : attachment
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not keep an attachment in the chat's files: #{e.class}" }
+        attachment
+      end
+
+      def too_long?(attachment)
+        @inline_limit && attachment.text? && attachment.content.to_s.length > @inline_limit
+      end
+
+      def start_of(attachment, file)
+        text_attachment(file, attachment.content.to_s)
+      end
+
+      def as_text(file, config)
+        text = store.text(file.path, config:)
+        return if text.nil? || text.strip.empty?
+
+        text_attachment(file, text)
+      rescue StandardError => e
+        RubyLLM.logger.warn { "Could not read #{file.path} as text: #{e.class}" }
+        nil
+      end
+
+      def text_attachment(file, text)
+        partial = @inline_limit && text.length > @inline_limit
+        if partial
+          text = "#{text[0, @inline_limit]}\n\n[#{file.path} has #{text.length} characters, so only the start " \
+                 'is shown, and you have not seen the rest. Read more with read_file, or search it with search_files.]'
+        end
+        attachment = Attachment.new(StringIO.new("#{file.path}, read as text:\n\n#{text}"),
+                                    filename: "#{File.basename(file.path)}.txt")
+        partial ? attachment.mark_partial : attachment
+      end
+
+      def reference(path, file)
+        text = if file
+                 "#{path} (#{file.content_type}, #{file.human_size}) is in this chat's files and not shown " \
+                   'here. Open it with read_file, or search it with search_files.'
+               else
+                 "#{path} was shown here. Open it again with read_file to see it."
+               end
+        Attachment.new(StringIO.new(text), filename: "#{File.basename(path)}.txt")
+      end
+
+      # Returns the path each read_file call asked for, by call id.
+      def read_paths(list)
+        list.flat_map { |message| message.tool_calls.to_a }
+            .select { |call| call.name == READ_TOOL }
+            .to_h { |call| [call.id, read_path(call)] }
+      end
+
+      def read_path(call)
+        path = call.arguments.to_h.transform_keys(&:to_s)['path'].to_s
+        path.empty? ? 'the file' : path
+      end
+
+      def write(path, content)
+        store.write(path, content).tap { |file| @written.call(file) }
+      end
+
+      def readable(content)
+        return [content, 'txt'] unless content.lstrip.start_with?('{', '[')
+
+        [JSON.pretty_generate(JSON.parse(content)), 'json']
+      rescue JSON::ParserError
+        [content, 'txt']
+      end
+
+      def safe_name(name)
+        name.to_s.gsub(/[^A-Za-z0-9_-]/, '_')[0, 64]
+      end
+
+      def safe_id(id)
+        id.to_s.gsub(/[^A-Za-z0-9]/, '').chars.last(8).join
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/message_queue.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Holds a chat's queued messages in memory. Any thread may queue or
+    # withdraw while the conversation loop takes them, and each message is
+    # taken or withdrawn once. A message that fails to join stays queued,
+    # with the ones taken after it, and those that joined are already
+    # yielded.
+    class MessageQueue # :nodoc:
+      def initialize(messages = [], &add)
+        @add = add
+        @messages = messages.map { |message| message.requeue(self) }
+        @lock = Mutex.new
+      end
+
+      def push(content:, attachments:, steer:, sender:)
+        queued = QueuedMessage.new(content:, attachments:, steer:, sender:, queue: self)
+        @lock.synchronize { @messages << queued }
+        queued
+      end
+
+      def delete(queued)
+        id = queued.respond_to?(:id) ? queued.id : queued
+        @lock.synchronize do
+          index = @messages.index { |message| message.id == id }
+          @messages.delete_at(index) if index
+        end
+      end
+
+      def to_a
+        @lock.synchronize { @messages.dup }
+      end
+
+      def none?(steering_only: false)
+        @lock.synchronize { @messages.none? { |message| message.steer? || !steering_only } }
+      end
+
+      def take(steering_only: false)
+        taken = @lock.synchronize do
+          ready, @messages = @messages.partition { |message| message.steer? || !steering_only }
+          ready
+        end
+        taken.each_with_index do |queued, index|
+          joined = index
+          message = @add.call(queued)
+          joined = index + 1
+          yield message, queued
+        rescue StandardError
+          @lock.synchronize { @messages.unshift(*taken.drop(joined)) }
+          raise
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/ownership.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Attributes every provider attempt a chat's operations make, including
+    # the tools they run, the converters a render runs, and the callbacks
+    # of a completion a batch delivers, to the owner set with
+    # Chat#with_owner, through the same fiber storage RubyLLM.with_owner
+    # uses. A chat without an owner leaves the surrounding one in place.
+    module Ownership # :nodoc:
+      [*Rescuable::OPERATIONS, :render, :add_completion].each do |operation|
+        define_method(operation) do |*args, **kwargs, &block|
+          Accounting::Usage.owned_by(owner) { super(*args, **kwargs, &block) }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/rescuable.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Runs a chat's operations through the Agent::rescue_from handlers of the
+    # agent that configured it, however the chat was reached. Only the
+    # outermost operation rescues: #ask runs #complete, and a handler's value
+    # must not stand in for a step of a loop that is still running.
+    module Rescuable # :nodoc:
+      OPERATIONS = %i[ask say ask_later complete generate run_tools step count_tokens compact].freeze
+
+      attr_writer :rescuer
+
+      # The agents whose callback macros run on this chat, by agent class.
+      def agent_callback_receivers
+        @agent_callback_receivers ||= {}
+      end
+
+      OPERATIONS.each do |operation|
+        define_method(operation) do |*args, **kwargs, &block|
+          return super(*args, **kwargs, &block) unless @rescuer
+
+          rescuing_with(@rescuer) { super(*args, **kwargs, &block) }
+        end
+      end
+
+      def rescuing_with(agent)
+        return yield if @rescuing
+
+        begin
+          @rescuing = true
+          yield
+        rescue StandardError => e
+          @rescuing = false
+          agent.rescue_with_handler(e)
+        ensure
+          @rescuing = false
+        end
+      end
+
+      def initialize_copy(source)
+        super
+        @rescuing = false
+        @agent_callback_receivers = @agent_callback_receivers&.dup
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/run.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # What the conversation loop tells the world around it, and asks of it.
+    # A plain chat reports usage through +usage.ruby_llm+ and has nothing
+    # else to consult. The Rails integration swaps in ActiveRecord::ChatRun,
+    # which reaches the record that owns the run. A copy of the chat shares
+    # it.
+    class Run # :nodoc:
+      def record_usage(entry)
+        Accounting::Usage.record(entry)
+      end
+
+      def record_progress(_partial); end
+
+      def chat_id; end
+
+      def cache_key; end
+
+      def owner; end
+
+      def cancellation_requested?
+        false
+      end
+
+      def interrupted_by?(_error)
+        false
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/tool_call_ids.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'digest'
+require 'json'
+
+module RubyLLM
+  class Chat
+    class ToolCallIds # :nodoc:
+      def self.answer(result, messages)
+        calls = unanswered(messages)
+        call = calls.find { |candidate| candidate.id == result.tool_call_id } ||
+               calls.reverse.find { |candidate| candidate.provider_call_id == result.provider_call_id }
+        call ? result.with_tool_call_id(call.id) : result
+      end
+
+      def self.unanswered(messages)
+        answered = messages.select(&:tool_result?).to_set(&:tool_call_id)
+        messages.select(&:tool_call?).flat_map { |message| message.tool_calls.to_a }
+                .reject { |call| answered.include?(call.id) }
+      end
+      private_class_method :unanswered
+
+      def initialize(namespace: nil)
+        @ids = {}
+        @namespace = namespace
+      end
+
+      def assign(message)
+        return message unless message.tool_call?
+
+        calls = message.tool_calls.transform_values do |call|
+          call.provider_call_id ? call.with_id(id_for(call.provider_call_id)) : call
+        end
+        calls = ToolCall.index(calls.values) unless message.is_a?(Chunk)
+        message.with_tool_calls(calls)
+      end
+
+      private
+
+      def id_for(id)
+        @ids[id] ||= @namespace ? Digest::SHA256.hexdigest(JSON.generate([*@namespace, id])) : SecureRandom.uuid
+      end
+    end
+  end
+end

data/lib/ruby_llm/chat/tool_call_state.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  class Chat
+    # Holds what a chat knows about its tool calls beyond the transcript, in
+    # memory. A copy of the chat gets a copy of the state. The Rails
+    # integration swaps in ActiveRecord::ChatToolCallState, which keeps it
+    # on the tool call records so any process can read it.
+    class ToolCallState # :nodoc:
+      def initialize
+        @inputs = {}
+        @decisions = {}
+        @notes = {}
+        @remembered = []
+        @spent_escalations = []
+        @lock = Mutex.new
+      end
+
+      def initialize_copy(source)
+        super
+        @inputs = @inputs.dup
+        @decisions = @decisions.dup
+        @notes = @notes.dup
+        @remembered = @remembered.dup
+        @spent_escalations = @spent_escalations.dup
+      end
+
+      def input(tool_call)
+        @inputs[tool_call.id]
+      end
+
+      def record_input(tool_call, input)
+        @inputs[tool_call.id] = input
+      end
+
+      # Replaces the input +tool_call+ is paused on with what the block
+      # returns for it, while the call is still paused. Answers recorded
+      # from several threads at once take turns, each on the input the one
+      # before left.
+      def settle(tool_call)
+        exclusively do
+          input = @inputs[tool_call.id]
+          @inputs[tool_call.id] = yield(input) if input
+        end
+      end
+
+      # Runs the block while no other thread settles or decides through
+      # this state, so a check and the write it leads to stay together.
+      def exclusively(&)
+        @lock.synchronize(&)
+      end
+
+      def record_decision(id, decision)
+        @decisions[id] = decision
+      end
+
+      def decision(tool_call)
+        @decisions[tool_call.id]
+      end
+
+      # The note a person gave when declining +tool_call+, kept apart from
+      # its input so a declined call never looks paused on a question.
+      def decline_note(tool_call)
+        @notes[tool_call.id]
+      end
+
+      def record_decline_note(tool_call, note)
+        @notes[tool_call.id] = note
+      end
+
+      def denied?(tool_call)
+        decision(tool_call) == false
+      end
+
+      def remember(identity)
+        @remembered |= [identity]
+      end
+
+      def remembered?(call)
+        @remembered.include?([call.namespace, call.name.to_s])
+      end
+
+      def escalation_spent?(tool_call)
+        @spent_escalations.include?(tool_call.id)
+      end
+
+      def spend_escalation(tool_call)
+        @spent_escalations << tool_call.id
+      end
+
+      def claim(requested, **)
+        requested
+      end
+
+      # Frees the claims on +ids+ once their round ends, whether it ended
+      # normally or by +stopped_by+, the signal or exit that stopped it.
+      # A round that +finished+ left only calls that paused unanswered.
+      def release(_ids, _stopped_by = nil, finished: false); end
+
+      # Notes that +tool_call+, claimed with its round, starts running now.
+      def started(_tool_call); end
+
+      # Whether +tool_call+, claimed with its round, had started in a
+      # process that stopped before saving its result, so it is answered
+      # as interrupted instead of running again.
+      def interrupted?(_tool_call)
+        false
+      end
+
+      def running_elsewhere?(_tool_call)
+        false
+      end
+    end
+  end
+end