llm.rb 0.8.0 → 0.9.0

data/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/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/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/function.rb

@@ -1,32 +1,29 @@
 # frozen_string_literal: true
 
 ##
-# The {LLM::Function LLM::Function} class represents a function that can
-# be called by an LLM. It comes in two forms: a Proc-based function,
-# or a Class-based function.
+# The {LLM::Function LLM::Function} class represents a
+# local function that can be called by an LLM.
 #
-# @example
-#   # Proc-based
+# @example example #1
 #   LLM.function(:system) do |fn|
-#     fn.description "Runs system commands, emits their output"
+#     fn.description "Runs system commands"
 #     fn.params do |schema|
 #       schema.object(command: schema.string.required)
 #     end
 #     fn.define do |params|
-#       Kernel.system(params.command)
+#       {success: Kernel.system(params.command)}
 #     end
 #   end
 #
-# @example
-#   # Class-based
+# @example example #2
 #   class System
 #     def call(params)
-#       Kernel.system(params.command)
+#       {success: Kernel.system(params.command)}
 #     end
 #   end
 #
 #   LLM.function(:system) do |fn|
-#     fn.description "Runs system commands, emits their output"
+#     fn.description "Runs system commands"
 #     fn.params do |schema|
 #       schema.object(command: schema.string.required)
 #     end
@@ -99,7 +96,7 @@ class LLM::Function
   # Returns a value that communicates that the function call was cancelled
   # @example
   #   llm = LLM.openai(key: ENV["KEY"])
-  #   bot = LLM::Chat.new(llm, tools: [fn1, fn2])
+  #   bot = LLM::Bot.new(llm, tools: [fn1, fn2])
   #   bot.chat "I want to run the functions"
   #   bot.chat bot.functions.map(&:cancel)
   # @return [LLM::Function::Return]

data/lib/llm/object/builder.rb

@@ -9,7 +9,7 @@ class LLM::Object
     #   obj = LLM::Object.from_hash(person: {name: 'John'})
     #   obj.person.name  # => 'John'
     #   obj.person.class # => LLM::Object
-    # @param [Hash, Array] obj
+    # @param [Hash, LLM::Object, Array] obj
     #   A Hash object
     # @return [LLM::Object]
     #   An LLM::Object object initialized by visiting `obj` with recursion
@@ -19,20 +19,19 @@ class LLM::Object
       when Array then obj.map { |v| from_hash(v) }
       else
         visited = {}
-        obj.each { visited[_1] = walk(_2) }
+        obj.each { visited[_1] = visit(_2) }
         new(visited)
       end
     end
 
     private
 
-    def walk(value)
-      if Hash === value
-        from_hash(value)
-      elsif Array === value
-        value.map { |v| (Hash === v) ? from_hash(v) : v }
-      else
-        value
+    def visit(value)
+      case value
+      when self then from_hash(value.to_h)
+      when Hash then from_hash(value)
+      when Array then value.map { |v| visit(v) }
+      else value
       end
     end
   end

data/lib/llm/object/kernel.rb

@@ -38,7 +38,7 @@ class LLM::Object
     end
 
     def inspect
-      "#<#{self.class}:0x#{object_id.to_s(16)} @h=#{to_h.inspect}>"
+      "#<#{self.class}:0x#{object_id.to_s(16)} properties=#{to_h.inspect}>"
     end
     alias_method :to_s, :inspect
   end

data/lib/llm/object.rb

@@ -17,7 +17,7 @@ class LLM::Object < BasicObject
   ##
   # @param [Hash] h
   # @return [LLM::Object]
-  def initialize(h)
+  def initialize(h = {})
     @h = h.transform_keys(&:to_sym) || h
   end
 
@@ -51,6 +51,12 @@ class LLM::Object < BasicObject
     to_h.to_json(...)
   end
 
+  ##
+  # @return [Boolean]
+  def empty?
+    @h.empty?
+  end
+
   ##
   # @return [Hash]
   def to_h

data/lib/llm/provider.rb

@@ -21,10 +21,9 @@ class LLM::Provider
   #  Whether to use SSL for the connection
   def initialize(key:, host:, port: 443, timeout: 60, ssl: true)
     @key = key
-    @http = Net::HTTP.new(host, port).tap do |http|
-      http.use_ssl = ssl
-      http.read_timeout = timeout
-    end
+    @client = Net::HTTP.new(host, port)
+    @client.use_ssl = ssl
+    @client.read_timeout = timeout
   end
 
   ##
@@ -78,55 +77,55 @@ class LLM::Provider
   # Starts a new lazy chat powered by the chat completions API
   # @note
   #  This method creates a lazy version of a
-  #  {LLM::Chat LLM::Chat} object.
+  #  {LLM::Bot LLM::Bot} object.
   # @param prompt (see LLM::Provider#complete)
   # @param params (see LLM::Provider#complete)
-  # @return [LLM::Chat]
+  # @return [LLM::Bot]
   def chat(prompt, params = {})
     role = params.delete(:role)
-    LLM::Chat.new(self, params).lazy.chat(prompt, role:)
+    LLM::Bot.new(self, params).chat(prompt, role:)
   end
 
   ##
   # Starts a new chat powered by the chat completions API
   # @note
   #  This method creates a non-lazy version of a
-  #  {LLM::Chat LLM::Chat} object.
+  #  {LLM::Bot LLM::Bot} object.
   # @param prompt (see LLM::Provider#complete)
   # @param params (see LLM::Provider#complete)
   # @raise (see LLM::Provider#complete)
-  # @return [LLM::Chat]
+  # @return [LLM::Bot]
   def chat!(prompt, params = {})
     role = params.delete(:role)
-    LLM::Chat.new(self, params).chat(prompt, role:)
+    LLM::Bot.new(self, params).chat(prompt, role:)
   end
 
   ##
   # Starts a new lazy chat powered by the responses API
   # @note
   #  This method creates a lazy variant of a
-  #  {LLM::Chat LLM::Chat} object.
+  #  {LLM::Bot LLM::Bot} object.
   # @param prompt (see LLM::Provider#complete)
   # @param params (see LLM::Provider#complete)
   # @raise (see LLM::Provider#complete)
-  # @return [LLM::Chat]
+  # @return [LLM::Bot]
   def respond(prompt, params = {})
     role = params.delete(:role)
-    LLM::Chat.new(self, params).lazy.respond(prompt, role:)
+    LLM::Bot.new(self, params).respond(prompt, role:)
   end
 
   ##
   # Starts a new chat powered by the responses API
   # @note
   #  This method creates a non-lazy variant of a
-  #  {LLM::Chat LLM::Chat} object.
+  #  {LLM::Bot LLM::Bot} object.
   # @param prompt (see LLM::Provider#complete)
   # @param params (see LLM::Provider#complete)
   # @raise (see LLM::Provider#complete)
-  # @return [LLM::Chat]
+  # @return [LLM::Bot]
   def respond!(prompt, params = {})
     role = params.delete(:role)
-    LLM::Chat.new(self, params).respond(prompt, role:)
+    LLM::Bot.new(self, params).respond(prompt, role:)
   end
 
   ##
@@ -194,10 +193,7 @@ class LLM::Provider
   # Returns an object that can generate a JSON schema
   # @return [JSON::Schema]
   def schema
-    @schema ||= begin
-      require_relative "../json/schema"
-      JSON::Schema.new
-    end
+    @schema ||= JSON::Schema.new
   end
 
   ##
@@ -216,6 +212,8 @@ class LLM::Provider
 
   private
 
+  attr_reader :client
+
   ##
   # The headers to include with a request
   # @raise [NotImplementedError]
@@ -243,10 +241,21 @@ class LLM::Provider
   end
 
   ##
-  # Initiates a HTTP request
-  # @param [Net::HTTP] http
-  #  The HTTP object to use for the request
-  # @param [Net::HTTPRequest] req
+  # @return [Class]
+  def event_handler
+    LLM::EventHandler
+  end
+
+  ##
+  # @return [Class]
+  #  Returns the provider-specific Server-Side Events (SSE) parser
+  def stream_parser
+    raise NotImplementedError
+  end
+
+  ##
+  # Executes a HTTP request
+  # @param [Net::HTTPRequest] request
   #  The request to send
   # @param [Proc] b
   #  A block to yield the response to (optional)
@@ -260,8 +269,34 @@ class LLM::Provider
   #  When any other unsuccessful status code is returned
   # @raise [SystemCallError]
   #  When there is a network error at the operating system level
-  def request(http, req, &b)
-    res = http.request(req, &b)
+  # @return [Net::HTTPResponse]
+  def execute(request:, stream: nil, &b)
+    res = if stream
+      client.request(request) do |res|
+        handler = event_handler.new stream_parser.new(stream)
+        parser = LLM::EventStream::Parser.new
+        parser.register(handler)
+        res.read_body(parser)
+        # If the handler body is empty, it means the
+        # response was most likely not streamed or
+        # parsing has failed. In that case, we fallback
+        # on the original response body.
+        res.body = handler.body.empty? ? parser.body.dup : handler.body
+      ensure
+        parser&.free
+      end
+    else
+      client.request(request, &b)
+    end
+    handle_response(res)
+  end
+
+  ##
+  # Handles the response from a request
+  # @param [Net::HTTPResponse] res
+  #  The response to handle
+  # @return [Net::HTTPResponse]
+  def handle_response(res)
     case res
     when Net::HTTPOK then res
     else error_handler.new(res).raise_error!

data/lib/llm/providers/anthropic/error_handler.rb

@@ -23,11 +23,11 @@ class LLM::Anthropic
     def raise_error!
       case res
       when Net::HTTPUnauthorized
-        raise LLM::Error::Unauthorized.new { _1.response = res }, "Authentication error"
+        raise LLM::UnauthorizedError.new { _1.response = res }, "Authentication error"
       when Net::HTTPTooManyRequests
-        raise LLM::Error::RateLimit.new { _1.response = res }, "Too many requests"
+        raise LLM::RateLimitError.new { _1.response = res }, "Too many requests"
       else
-        raise LLM::Error::ResponseError.new { _1.response = res }, "Unexpected response"
+        raise LLM::ResponseError.new { _1.response = res }, "Unexpected response"
       end
     end
   end

data/lib/llm/providers/anthropic/models.rb

@@ -40,7 +40,7 @@ class LLM::Anthropic
     def all(**params)
       query = URI.encode_www_form(params)
       req = Net::HTTP::Get.new("/v1/models?#{query}", headers)
-      res = request(http, req)
+      res = execute(request: req)
       LLM::Response::ModelList.new(res).tap { |modellist|
         models = modellist.body["data"].map do |model|
           LLM::Model.from_hash(model).tap { _1.provider = @provider }
@@ -51,12 +51,8 @@ class LLM::Anthropic
 
     private
 
-    def http
-      @provider.instance_variable_get(:@http)
-    end
-
-    [:headers, :request].each do |m|
-      define_method(m) { |*args, &b| @provider.send(m, *args, &b) }
+    [:headers, :execute].each do |m|
+      define_method(m) { |*args, **kwargs, &b| @provider.send(m, *args, **kwargs, &b) }
     end
   end
 end

data/lib/llm/providers/anthropic/response_parser/completion_parser.rb

@@ -41,9 +41,9 @@ module LLM::Anthropic::ResponseParser
     def body = @body
     def role = body.role
     def model = body.model
-    def prompt_tokens = body.usage.input_tokens
-    def completion_tokens = body.usage.output_tokens
-    def total_tokens = body.usage.total_tokens
+    def prompt_tokens = body.usage&.input_tokens
+    def completion_tokens = body.usage&.output_tokens
+    def total_tokens = body.usage&.total_tokens
     def parts = body.content
     def texts = parts.select { _1["type"] == "text" }
     def tools = parts.select { _1["type"] == "tool_use" }

data/lib/llm/providers/anthropic/response_parser.rb

@@ -4,6 +4,7 @@ class LLM::Anthropic
   ##
   # @private
   module ResponseParser
+    require_relative "response_parser/completion_parser"
     def parse_embedding(body)
       {
         model: body["model"],

data/lib/llm/providers/anthropic/stream_parser.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+class LLM::Anthropic
+  ##
+  # @private
+  class StreamParser
+    ##
+    # Returns the fully constructed response body
+    # @return [LLM::Object]
+    attr_reader :body
+
+    ##
+    # @param [#<<] io An IO-like object
+    # @return [LLM::Anthropic::StreamParser]
+    def initialize(io)
+      @body = LLM::Object.new(role: "assistant", content: [])
+      @io = io
+    end
+
+    ##
+    # @param [Hash] chunk
+    # @return [LLM::Anthropic::StreamParser]
+    def parse!(chunk)
+      tap { merge!(chunk) }
+    end
+
+    private
+
+    def merge!(chunk)
+      if chunk["type"] == "message_start"
+        merge_message!(chunk["message"])
+      elsif chunk["type"] == "content_block_start"
+        @body["content"][chunk["index"]] = chunk["content_block"]
+      elsif chunk["type"] == "content_block_delta"
+        if chunk["delta"]["type"] == "text_delta"
+          @body.content[chunk["index"]]["text"] << chunk["delta"]["text"]
+          @io << chunk["delta"]["text"] if @io.respond_to?(:<<)
+        elsif chunk["delta"]["type"] == "input_json_delta"
+          content = @body.content[chunk["index"]]
+          if Hash === content["input"]
+            content["input"] = chunk["delta"]["partial_json"]
+          else
+            content["input"] << chunk["delta"]["partial_json"]
+          end
+        end
+      elsif chunk["type"] == "message_delta"
+        merge_message!(chunk["delta"])
+      elsif chunk["type"] == "content_block_stop"
+        content = @body.content[chunk["index"]]
+        if content["input"]
+          content["input"] = JSON.parse(content["input"])
+        end
+      end
+    end
+
+    def merge_message!(message)
+      message.each do |key, value|
+        @body[key] = if value.respond_to?(:each_pair)
+          merge_message!(value)
+        else
+          value
+        end
+      end
+    end
+  end
+end

data/lib/llm/providers/anthropic.rb

@@ -5,10 +5,10 @@ module LLM
   # The Anthropic class implements a provider for
   # [Anthropic](https://www.anthropic.com)
   class Anthropic < Provider
+    require_relative "anthropic/format"
     require_relative "anthropic/error_handler"
+    require_relative "anthropic/stream_parser"
     require_relative "anthropic/response_parser"
-    require_relative "anthropic/response_parser/completion_parser"
-    require_relative "anthropic/format"
     require_relative "anthropic/models"
     include Format
 
@@ -50,12 +50,13 @@ module LLM
     def complete(prompt, params = {})
       params = {role: :user, model: default_model, max_tokens: 1024}.merge!(params)
       params = [params, format_tools(params)].inject({}, &:merge!).compact
-      role = params.delete(:role)
+      role, stream = params.delete(:role), params.delete(:stream)
+      params[:stream] = true if stream.respond_to?(:<<) || stream == true
       req = Net::HTTP::Post.new("/v1/messages", headers)
       messages = [*(params.delete(:messages) || []), Message.new(role, prompt)]
       body = JSON.dump({messages: [format(messages)].flatten}.merge!(params))
       set_body_stream(req, StringIO.new(body))
-      res = request(@http, req)
+      res = execute(request: req, stream:)
       Response::Completion.new(res).extend(response_parser)
     end
 
@@ -95,6 +96,10 @@ module LLM
       LLM::Anthropic::ResponseParser
     end
 
+    def stream_parser
+      LLM::Anthropic::StreamParser
+    end
+
     def error_handler
       LLM::Anthropic::ErrorHandler
     end

data/lib/llm/providers/gemini/error_handler.rb

@@ -25,14 +25,14 @@ class LLM::Gemini
       when Net::HTTPBadRequest
         reason = body.dig("error", "details", 0, "reason")
         if reason == "API_KEY_INVALID"
-          raise LLM::Error::Unauthorized.new { _1.response = res }, "Authentication error"
+          raise LLM::UnauthorizedError.new { _1.response = res }, "Authentication error"
         else
-          raise LLM::Error::ResponseError.new { _1.response = res }, "Unexpected response"
+          raise LLM::ResponseError.new { _1.response = res }, "Unexpected response"
         end
       when Net::HTTPTooManyRequests
-        raise LLM::Error::RateLimit.new { _1.response = res }, "Too many requests"
+        raise LLM::RateLimitError.new { _1.response = res }, "Too many requests"
       else
-        raise LLM::Error::ResponseError.new { _1.response = res }, "Unexpected response"
+        raise LLM::ResponseError.new { _1.response = res }, "Unexpected response"
       end
     end