llm-shell 0.9.2 → 0.10.0

data/lib/llm/shell/internal/llm.rb/lib/llm/buffer.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # {LLM::Buffer LLM::Buffer} provides an Enumerable object that
+  # yields each message in a conversation on-demand, and only sends
+  # a request to the LLM when a response is needed.
+  class Buffer
+    include Enumerable
+
+    ##
+    # @param [LLM::Provider] provider
+    # @return [LLM::Buffer]
+    def initialize(provider)
+      @provider = provider
+      @pending = []
+      @completed = []
+    end
+
+    ##
+    # @yield [LLM::Message]
+    #  Yields each message in the conversation thread
+    # @raise (see LLM::Provider#complete)
+    # @return [void]
+    def each(...)
+      if block_given?
+        empty! unless @pending.empty?
+        @completed.each { yield(_1) }
+      else
+        enum_for(:each, ...)
+      end
+    end
+
+    ##
+    # Returns an array of unread messages
+    # @see LLM::Message#read?
+    # @see LLM::Message#read!
+    # @return [Array<LLM::Message>]
+    def unread
+      reject(&:read?)
+    end
+
+    ##
+    # Find a message (in descending order)
+    # @return [LLM::Message, nil]
+    def find(...)
+      reverse_each.find(...)
+    end
+
+    ##
+    # Returns the last message(s) in the buffer
+    # @param [Integer, nil] n
+    #  The number of messages to return
+    # @return [LLM::Message, Array<LLM::Message>, nil]
+    def last(n = nil)
+      if @pending.empty?
+        n.nil? ? @completed.last : @completed.last(n)
+      else
+        n.nil? ? to_a.last : to_a.last(n)
+      end
+    end
+
+    ##
+    # @param [[LLM::Message, Hash, Symbol]] item
+    #  A message and its parameters
+    # @return [void]
+    def <<(item)
+      @pending << item
+      self
+    end
+    alias_method :push, :<<
+
+    ##
+    # @param [Integer, Range] index
+    #  The message index
+    # @return [LLM::Message, nil]
+    #  Returns a message, or nil
+    def [](index)
+      if @pending.empty?
+        if Range === index
+          slice = @completed[index]
+          (slice.nil? || slice.size < index.size) ? to_a[index] : slice
+        else
+          @completed[index]
+        end
+      else
+        to_a[index]
+      end
+    end
+
+    ##
+    # @return [String]
+    def inspect
+      "#<#{self.class.name}:0x#{object_id.to_s(16)} " \
+      "completed_count=#{@completed.size} pending_count=#{@pending.size}>"
+    end
+
+    ##
+    # Returns true when the buffer is empty
+    # @return [Boolean]
+    def empty?
+      @pending.empty? and @completed.empty?
+    end
+
+    ##
+    # @example
+    #   llm = LLM.openai(key: ENV["KEY"])
+    #   bot = LLM::Bot.new(llm, stream: $stdout)
+    #   bot.chat "Hello", role: :user
+    #   bot.messages.flush
+    # @see LLM::Bot#drain
+    # @note
+    #   This method is especially useful when using the streaming API.
+    # Drains the buffer and returns all messages as an array
+    # @return [Array<LLM::Message>]
+    def drain
+      to_a
+    end
+    alias_method :flush, :drain
+
+    private
+
+    def empty!
+      message, params, method = @pending.pop
+      if method == :complete
+        complete!(message, params)
+      elsif method == :respond
+        respond!(message, params)
+      else
+        raise LLM::Error, "Unknown method: #{method}"
+      end
+    end
+
+    def complete!(message, params)
+      oldparams = @pending.map { _1[1] }
+      pendings  = @pending.map { _1[0] }
+      messages  = [*@completed, *pendings]
+      role = message.role
+      completion = @provider.complete(
+        message.content,
+        [*oldparams, params.merge(role:, messages:)].inject({}, &:merge!)
+      )
+      @completed.concat([*pendings, message, *completion.choices[0]])
+      @pending.clear
+    end
+
+    def respond!(message, params)
+      oldparams = @pending.map { _1[1] }
+      pendings  = @pending.map { _1[0] }
+      messages  = [*pendings]
+      role = message.role
+      params = [
+        *oldparams,
+        params.merge(input: messages),
+        @response ? {previous_response_id: @response.response_id} : {}
+      ].inject({}, &:merge!)
+      @response = @provider.responses.create(message.content, params.merge(role:))
+      @completed.concat([*pendings, message, *@response.choices[0]])
+      @pending.clear
+    end
+  end
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/client.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # @api private
+  module Client
+    private
+
+    ##
+    # @api private
+    def persistent_client
+      LLM.lock(:clients) do
+        if clients[client_id]
+          clients[client_id]
+        else
+          require "net/http/persistent" unless defined?(Net::HTTP::Persistent)
+          client = Net::HTTP::Persistent.new(name: self.class.name)
+          client.read_timeout = timeout
+          clients[client_id] = client
+        end
+      end
+    end
+
+    ##
+    # @api private
+    def transient_client
+      client = Net::HTTP.new(host, port)
+      client.read_timeout = timeout
+      client.use_ssl = ssl
+      client
+    end
+
+    def client_id = "#{host}:#{port}:#{timeout}:#{ssl}"
+    def clients = self.class.clients
+  end
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/error.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # The superclass of all LLM errors
+  class Error < RuntimeError
+    def initialize(...)
+      block_given? ? yield(self) : nil
+      super
+    end
+  end
+
+  ##
+  # The superclass of all HTTP protocol errors
+  class ResponseError < Error
+    ##
+    # @return [Net::HTTPResponse]
+    #  Returns the response associated with an error
+    attr_accessor :response
+
+    def message
+      [super, response.body].join("\n")
+    end
+  end
+
+  ##
+  # HTTPUnauthorized
+  UnauthorizedError = Class.new(ResponseError)
+
+  ##
+  # HTTPTooManyRequests
+  RateLimitError = Class.new(ResponseError)
+
+  ##
+  # HTTPServerError
+  ServerError = Class.new(ResponseError)
+
+  ##
+  # When no images are found in a response
+  NoImageError = Class.new(ResponseError)
+
+  ##
+  # When an given an input object that is not understood
+  FormatError = Class.new(Error)
+
+  ##
+  # When given a prompt object that is not understood
+  PromptError = Class.new(FormatError)
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/eventhandler.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module LLM
+  ##
+  # @private
+  class EventHandler
+    ##
+    # @param [#parse!] parser
+    # @return [LLM::EventHandler]
+    def initialize(parser)
+      @parser = parser
+    end
+
+    ##
+    # "data:" event callback
+    # @param [LLM::EventStream::Event] event
+    # @return [void]
+    def on_data(event)
+      return if event.end?
+      chunk = JSON.parse(event.value)
+      @parser.parse!(chunk)
+    rescue JSON::ParserError
+    end
+
+    ##
+    # Callback for when *any* of chunk of data
+    # is received, regardless of whether it has
+    # a field name or not. Primarily for ollama,
+    # which does emit Server-Sent Events (SSE).
+    # @param [LLM::EventStream::Event] event
+    # @return [void]
+    def on_chunk(event)
+      return if event.end?
+      chunk = JSON.parse(event.chunk)
+      @parser.parse!(chunk)
+    rescue JSON::ParserError
+    end
+
+    ##
+    # Returns a fully constructed response body
+    # @return [LLM::Object]
+    def body = @parser.body
+  end
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/eventstream/event.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+module LLM::EventStream
+  ##
+  # @private
+  class Event
+    FIELD_REGEXP = /[^:]+/
+    VALUE_REGEXP = /(?<=: ).+/
+
+    ##
+    # Returns the field name
+    # @return [Symbol]
+    attr_reader :field
+
+    ##
+    # Returns the field value
+    # @return [String]
+    attr_reader :value
+
+    ##
+    # Returns the full chunk
+    # @return [String]
+    attr_reader :chunk
+
+    ##
+    # @param [String] chunk
+    # @return [LLM::EventStream::Event]
+    def initialize(chunk)
+      @field = chunk[FIELD_REGEXP]
+      @value = chunk[VALUE_REGEXP]
+      @chunk = chunk
+    end
+
+    ##
+    # Returns true when the event represents an "id" chunk
+    # @return [Boolean]
+    def id?
+      @field == "id"
+    end
+
+    ##
+    # Returns true when the event represents a "data" chunk
+    # @return [Boolean]
+    def data?
+      @field == "data"
+    end
+
+    ##
+    # Returns true when the event represents an "event" chunk
+    # @return [Boolean]
+    def event?
+      @field == "event"
+    end
+
+    ##
+    # Returns true when the event represents a "retry" chunk
+    # @return [Boolean]
+    def retry?
+      @field == "retry"
+    end
+
+    ##
+    # Returns true when a chunk represents the end of the stream
+    # @return [Boolean]
+    def end?
+      @value == "[DONE]"
+    end
+  end
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/eventstream/parser.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module LLM::EventStream
+  ##
+  # @private
+  class Parser
+    ##
+    # @return [LLM::EventStream::Parser]
+    def initialize
+      @buffer = StringIO.new
+      @events = Hash.new { |h, k| h[k] = [] }
+      @offset = 0
+      @visitors = []
+    end
+
+    ##
+    # Register a visitor
+    # @param [#on_data] visitor
+    # @return [void]
+    def register(visitor)
+      @visitors << visitor
+    end
+
+    ##
+    # Subscribe to an event
+    # @param [Symbol] evtname
+    # @param [Proc] block
+    # @return [void]
+    def on(evtname, &block)
+      @events[evtname.to_s] << block
+    end
+
+    ##
+    # Append an event to the internal buffer
+    # @return [void]
+    def <<(event)
+      io = StringIO.new(event)
+      IO.copy_stream io, @buffer
+      each_line { parse!(_1) }
+    end
+
+    ##
+    # Returns the internal buffer
+    # @return [String]
+    def body
+      @buffer.string
+    end
+
+    ##
+    # Free the internal buffer
+    # @return [void]
+    def free
+      @buffer.truncate(0)
+      @buffer.rewind
+    end
+
+    private
+
+    def parse!(event)
+      event = Event.new(event)
+      dispatch(event)
+    end
+
+    def dispatch(event)
+      @visitors.each { dispatch_visitor(_1, event) }
+      @events[event.field].each { _1.call(event) }
+    end
+
+    def dispatch_visitor(visitor, event)
+      method = "on_#{event.field}"
+      if visitor.respond_to?(method)
+        visitor.public_send(method, event)
+      elsif visitor.respond_to?("on_chunk")
+        visitor.on_chunk(event)
+      end
+    end
+
+    def each_line
+      string.each_line.with_index do
+        next if _2 < @offset
+        yield(_1)
+        @offset += 1
+      end
+    end
+
+    def string = @buffer.string
+  end
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/eventstream.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+##
+# @private
+module LLM::EventStream
+  require_relative "eventstream/parser"
+  require_relative "eventstream/event"
+end

data/lib/llm/shell/internal/llm.rb/lib/llm/file.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+##
+# {LLM::File LLM::File} represents a local file. It can be used
+# as a prompt with certain providers (eg: Ollama, Gemini),
+# and as an input with certain methods. It is usually not necessary
+# to create an instance of LLM::File directly.
+class LLM::File
+  ##
+  # @return [String]
+  #  Returns the path to the file
+  attr_reader :path
+
+  def initialize(path)
+    @path = path
+  end
+
+  ##
+  # @return [String]
+  #  Returns basename of the file
+  def basename
+    File.basename(path)
+  end
+
+  ##
+  # @return [String]
+  #  Returns the MIME type of the file
+  def mime_type
+    LLM::Mime[path]
+  end
+
+  ##
+  # @return [Boolean]
+  #  Returns true if the file is an image
+  def image?
+    mime_type.start_with?("image/")
+  end
+
+  ##
+  # @return [Boolean]
+  #  Returns true if the file is a PDF document
+  def pdf?
+    mime_type == "application/pdf"
+  end
+
+  ##
+  # @return [Integer]
+  #  Returns the size of the file in bytes
+  def bytesize
+    File.size(path)
+  end
+
+  ##
+  # @return [String]
+  #  Returns the file contents in base64
+  def to_b64
+    [File.binread(path)].pack("m0")
+  end
+
+  ##
+  # @return [String]
+  #  Returns the file contents in base64 URL format
+  def to_data_uri
+    "data:#{mime_type};base64,#{to_b64}"
+  end
+
+  ##
+  # @return [File]
+  #  Yields an IO object suitable to be streamed
+  def with_io
+    io = File.open(path, "rb")
+    yield(io)
+  ensure
+    io.close
+  end
+end
+
+##
+# @param [String, File, LLM::Response] obj
+#  The path to the file, or an existing file reference
+# @return [LLM::File]
+def LLM.File(obj)
+  case obj
+  when File
+    obj.close unless obj.closed?
+    LLM.File(obj.path)
+  when LLM::File, LLM::Response then obj
+  when String then LLM::File.new(obj)
+  else raise TypeError, "don't know how to handle #{obj.class} objects"
+  end
+end