llm.rb 0.4.1 → 0.5.0

data/lib/json/schema.rb

@@ -3,7 +3,25 @@
 module JSON
 end unless defined?(JSON)
 
+##
+# The {JSON::Schema JSON::Schema} class represents a JSON schema,
+# and provides methods that let you describe and produce a schema
+# that can be used in various contexts that include the validation
+# and generation of JSON data.
+#
+# @see https://json-schema.org/ JSON Schema Specification
+# @see https://tour.json-schema.org/ JSON Schema Tour
+#
+# @example
+#  schema = JSON::Schema.new
+#  schema.object({
+#    name: schema.string.enum("John", "Jane").required,
+#    age: schema.integer.required,
+#    hobbies: schema.array(schema.string, schema.null).required,
+#    address: schema.object({street: schema.string}).required,
+#  })
 class JSON::Schema
+  require_relative "schema/version"
   require_relative "schema/leaf"
   require_relative "schema/object"
   require_relative "schema/array"
@@ -15,59 +33,52 @@ class JSON::Schema
 
   ##
   # Returns an object
-  # @param properties [Hash] A hash of properties
-  # @param rest [Hash] Any other options
+  # @param [Hash] properties A hash of properties
   # @return [JSON::Schema::Object]
-  def object(properties, **rest)
-    Object.new(properties, **rest)
+  def object(properties)
+    Object.new(properties)
   end
 
   ##
   # Returns an array
-  # @param items [Array] An array of items
-  # @param rest [Hash] Any other options
+  # @param [Array] items An array of items
   # @return [JSON::Schema::Array]
-  def array(items, **rest)
-    Array.new(items, **rest)
+  def array(*items)
+    Array.new(*items)
   end
 
   ##
   # Returns a string
-  # @param rest [Hash] Any other options
   # @return [JSON::Schema::String]
-  def string(...)
-    String.new(...)
+  def string
+    String.new
   end
 
   ##
   # Returns a number
-  # @param rest [Hash] Any other options
   # @return [JSON::Schema::Number] a number
-  def number(...)
-    Number.new(...)
+  def number
+    Number.new
   end
 
   ##
   # Returns an integer
-  # @param rest [Hash] Any other options
   # @return [JSON::Schema::Integer]
-  def integer(...)
-    Integer.new(...)
+  def integer
+    Integer.new
   end
 
   ##
   # Returns a boolean
-  # @param rest [Hash] Any other options
   # @return [JSON::Schema::Boolean]
-  def boolean(...)
-    Boolean.new(...)
+  def boolean
+    Boolean.new
   end
 
   ##
   # Returns null
-  # @param rest [Hash] Any other options
   # @return [JSON::Schema::Null]
-  def null(...)
-    Null.new(...)
+  def null
+    Null.new
   end
 end

data/lib/llm/buffer.rb

@@ -28,6 +28,22 @@ module LLM
       @completed.each { yield(_1) }
     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
+
     ##
     # @param [[LLM::Message, Hash]] item
     #  A message and its parameters
@@ -48,7 +64,7 @@ module LLM
     private
 
     def empty!
-      message, params, method = @pending[-1]
+      message, params, method = @pending.pop
       if method == :complete
         complete!(message, params)
       elsif method == :respond
@@ -59,24 +75,26 @@ module LLM
     end
 
     def complete!(message, params)
-      messages = @pending[0..-2].map { _1[0] }
+      pendings = @pending.map { _1[0] }
+      messages = [*@completed, *pendings]
       completion = @provider.complete(
         message.content,
         message.role,
         **params.merge(messages:)
       )
-      @completed.concat([*messages, message, completion.choices[0]])
+      @completed.concat([*pendings, message, completion.choices[0]])
       @pending.clear
     end
 
     def respond!(message, params)
-      input = @pending[0..-2].map { _1[0] }
-      @response = @provider.responses.create(
-        message.content,
-        message.role,
-        **params.merge(input:).merge(@response ? {previous_response_id: @response.id} : {})
-      )
-      @completed.concat([*input, message, @response.outputs[0]])
+      pendings = @pending.map { _1[0] }
+      input = [*pendings]
+      params = [
+        params.merge(input:),
+        @response ? {previous_response_id: @response.id} : {}
+      ].inject({}, &:merge!)
+      @response = @provider.responses.create(message.content, message.role, **params)
+      @completed.concat([*pendings, message, @response.outputs[0]])
       @pending.clear
     end
   end

data/lib/llm/chat.rb

@@ -37,7 +37,7 @@ module LLM
       @provider = provider
       @params = params.merge!(model:, schema:)
       @lazy = false
-      @messages = []
+      @messages = [].extend(Array)
     end
 
     ##
@@ -108,14 +108,39 @@ module LLM
       @lazy
     end
 
+    ##
+    # @return [String]
     def inspect
       "#<#{self.class.name}:0x#{object_id.to_s(16)} " \
       "@provider=#{@provider.class}, @params=#{@params.inspect}, " \
       "@messages=#{@messages.inspect}, @lazy=#{@lazy.inspect}>"
     end
 
+    ##
+    # Returns an array of functions that have yet to be called
+    # @return [Array<LLM::Function>]
+    def functions
+      messages
+        .select(&:assistant?)
+        .flat_map(&:functions)
+        .reject(&:called?)
+    end
+
     private
 
+    ##
+    # @private
+    module Array
+      def find(...)
+        reverse_each.find(...)
+      end
+
+      def unread
+        reject(&:read?)
+      end
+    end
+    private_constant :Array
+
     def respond!(prompt, role, params)
       @provider.responses.create(
         prompt,

data/lib/llm/core_ext/ostruct.rb

@@ -8,17 +8,19 @@ class OpenStruct
     #   obj = OpenStruct.from_hash(person: {name: 'John'})
     #   obj.person.name  # => 'John'
     #   obj.person.class # => OpenStruct
-    # @param [Hash] hash_obj
+    # @param [Hash, Array] obj
     #   A Hash object
     # @return [OpenStruct]
-    #   An OpenStruct object initialized by visiting `hash_obj` with
-    #   recursion
-    def from_hash(hash_obj)
-      visited_object = {}
-      hash_obj.each do |key, value|
-        visited_object[key] = walk(value)
+    #   An OpenStruct object initialized by visiting `obj` with recursion
+    def from_hash(obj)
+      case obj
+      when self then from_hash(obj.to_h)
+      when Array then obj.map { |v| from_hash(v) }
+      else
+        visited = {}
+        obj.each { visited[_1] = walk(_2) }
+        new(visited)
       end
-      new(visited_object)
     end
 
     private
@@ -34,4 +36,8 @@ class OpenStruct
     end
   end
   extend FromHash
+
+  def to_json(...)
+    to_h.to_json(...)
+  end
 end

data/lib/llm/file.rb

@@ -72,5 +72,10 @@ end
 #  The path to a file
 # @return [LLM::File]
 def LLM.File(path)
-  LLM::File.new(path)
+  case path
+  when LLM::File, LLM::Response::File
+    path
+  else
+    LLM::File.new(path)
+  end
 end

data/lib/llm/function.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+class LLM::Function
+  class Return < Struct.new(:id, :value)
+  end
+
+  ##
+  # Returns function arguments
+  # @return [Array, nil]
+  attr_accessor :arguments
+
+  ##
+  # Returns the function ID
+  # @return [String, nil]
+  attr_accessor :id
+
+  ##
+  # @param [String] name The function name
+  # @yieldparam [LLM::Function] self The function object
+  def initialize(name, &b)
+    @name = name
+    @schema = JSON::Schema.new
+    yield(self)
+  end
+
+  ##
+  # Set the function description
+  # @param [String] str The function description
+  # @return [void]
+  def description(str)
+    @description = str
+  end
+
+  ##
+  # @yieldparam [JSON::Schema] schema The schema object
+  # @return [void]
+  def params
+    @params = yield(@schema)
+  end
+
+  ##
+  # Set the function implementation
+  # @param [Proc] b The function implementation
+  # @return [void]
+  def define(&b)
+    @runner = b
+  end
+
+  ##
+  # Call the function
+  # @param [Array] args The arguments to pass to the function
+  # @return [Object] The result of the function call
+  def call
+    Return.new id, @runner.call(arguments)
+  ensure
+    @called = true
+  end
+
+  ##
+  # Returns true when a function has been called
+  # @return [Boolean]
+  def called?
+    @called
+  end
+
+  ##
+  # @return [Hash]
+  def format(provider)
+    case provider.class.to_s
+    when "LLM::Gemini"
+      {name: @name, description: @description, parameters: @params}.compact
+    when "LLM::Anthropic"
+      {name: @name, description: @description, input_schema: @params}.compact
+    else
+      {
+        type: "function", name: @name,
+        function: {name: @name, description: @description, parameters: @params}
+      }.compact
+    end
+  end
+end

data/lib/llm/message.rb

@@ -64,12 +64,57 @@ module LLM
       role == "assistant" || role == "model"
     end
 
+    ##
+    # @return [Array<LLM::Function>]
+    def functions
+      @functions ||= tool_calls.map do |fn|
+        function = LLM.functions[fn.name].dup
+        function.tap { _1.id = fn.id }
+        function.tap { _1.arguments = fn.arguments }
+      end
+    end
+
+    ##
+    # @return [Boolean]
+    #  Returns true when the message requests a function call
+    def tool_call?
+      tool_calls.any?
+    end
+
+    ##
+    # Returns true when the message is a system message
+    # @return [Boolean]
+    def system?
+      role == "system"
+    end
+
+    ##
+    # Marks the message as read
+    # @return [void]
+    def read!
+      @read = true
+    end
+
+    ##
+    # Returns true when the message has been read
+    # @return [Boolean]
+    def read?
+      @read
+    end
+
     ##
     # Returns a string representation of the message
     # @return [String]
     def inspect
       "#<#{self.class.name}:0x#{object_id.to_s(16)} " \
-      "role=#{role.inspect} content=#{content.inspect}>"
+      "tool_call=#{tool_calls.any?} role=#{role.inspect} " \
+      "content=#{content.inspect}>"
+    end
+
+    private
+
+    def tool_calls
+      @tool_calls ||= OpenStruct.from_hash(@extra[:tool_calls] || [])
     end
   end
 end

data/lib/llm/providers/anthropic/format/completion_format.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module LLM::Anthropic::Format
+  ##
+  # @private
+  class CompletionFormat
+    ##
+    # @param [LLM::Message, Hash] message
+    #  The message to format
+    def initialize(message)
+      @message = message
+    end
+
+    ##
+    # Formats the message for the Anthropic chat completions API
+    # @return [Hash]
+    def format
+      catch(:abort) do
+        if Hash === message
+          {role: message[:role], content: format_content(message[:content])}
+        else
+          format_message
+        end
+      end
+    end
+
+    private
+
+    def format_message
+      if message.tool_call?
+        {role: message.role, content: message.extra[:original_tool_calls]}
+      else
+        {role: message.role, content: format_content(content)}
+      end
+    end
+
+    ##
+    # @param [String, URI] content
+    #  The content to format
+    # @return [String, Hash]
+    #  The formatted content
+    def format_content(content)
+      case content
+      when Hash
+        content.empty? ? throw(:abort, nil) : [content]
+      when Array
+        content.empty? ? throw(:abort, nil) : content.flat_map { format_content(_1) }
+      when URI
+        [{type: :image, source: {type: "url", url: content.to_s}}]
+      when LLM::File
+        if content.image?
+          [{type: :image, source: {type: "base64", media_type: content.mime_type, data: content.to_b64}}]
+        else
+          raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
+                                          "is not an image, and therefore not supported by the " \
+                                          "Anthropic API"
+        end
+      when String
+        [{type: :text, text: content}]
+      when LLM::Message
+        format_content(content.content)
+      when LLM::Function::Return
+        {type: "tool_result", tool_use_id: content.id, content: content.value}
+      else
+        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
+                                       "is not supported by the Anthropic API"
+      end
+    end
+
+    def message = @message
+    def content = message.content
+  end
+end

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

@@ -4,49 +4,23 @@ class LLM::Anthropic
   ##
   # @private
   module Format
+    require_relative "format/completion_format"
+
     ##
     # @param [Array<LLM::Message>] messages
     #  The messages to format
     # @return [Array<Hash>]
     def format(messages)
-      messages.map do
-        if Hash === _1
-          {role: _1[:role], content: format_content(_1[:content])}
-        else
-          {role: _1.role, content: format_content(_1.content)}
-        end
+      messages.filter_map do
+        CompletionFormat.new(_1).format
       end
     end
 
     private
 
-    ##
-    # @param [String, URI] content
-    #  The content to format
-    # @return [String, Hash]
-    #  The formatted content
-    def format_content(content)
-      case content
-      when Array
-        content.flat_map { format_content(_1) }
-      when URI
-        [{type: :image, source: {type: "url", url: content.to_s}}]
-      when LLM::File
-        if content.image?
-          [{type: :image, source: {type: "base64", media_type: content.mime_type, data: content.to_b64}}]
-        else
-          raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
-                                          "is not an image, and therefore not supported by the " \
-                                          "Anthropic API"
-        end
-      when String
-        [{type: :text, text: content}]
-      when LLM::Message
-        format_content(content.content)
-      else
-        raise LLM::Error::PromptError, "The given object (an instance of #{content.class}) " \
-                                       "is not supported by the Anthropic API"
-      end
+    def format_tools(tools)
+      return {} unless tools
+      {tools: tools.map { _1.format(self) }}
     end
   end
 end

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

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module LLM::Anthropic::ResponseParser
+  ##
+  # @private
+  class CompletionParser
+    def initialize(body)
+      @body = OpenStruct.from_hash(body)
+    end
+
+    def format(response)
+      {
+        model:,
+        prompt_tokens:,
+        completion_tokens:,
+        total_tokens:,
+        choices: format_choices(response)
+      }
+    end
+
+    private
+
+    def format_choices(response)
+      texts.map.with_index do |choice, index|
+        extra = {index:, response:, tool_calls: format_tool_calls(tools), original_tool_calls: tools}
+        LLM::Message.new(role, choice.text, extra)
+      end
+    end
+
+    def format_tool_calls(tools)
+      (tools || []).filter_map do |tool|
+        tool = {
+          id: tool.id,
+          name: tool.name,
+          arguments: tool.input
+        }
+        OpenStruct.new(tool)
+      end
+    end
+
+    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 parts = body.content
+    def texts = parts.select { _1["type"] == "text" }
+    def tools = parts.select { _1["type"] == "tool_use" }
+  end
+end

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

@@ -17,15 +17,7 @@ class LLM::Anthropic
     #  The response body from the LLM provider
     # @return [Hash]
     def parse_completion(body)
-      {
-        model: body["model"],
-        choices: body["content"].map do
-          # TODO: don't hardcode role
-          LLM::Message.new("assistant", _1["text"], {response: self})
-        end,
-        prompt_tokens: body.dig("usage", "input_tokens"),
-        completion_tokens: body.dig("usage", "output_tokens")
-      }
+      CompletionParser.new(body).format(self)
     end
   end
 end

data/lib/llm/providers/anthropic.rb

@@ -7,6 +7,7 @@ module LLM
   class Anthropic < Provider
     require_relative "anthropic/error_handler"
     require_relative "anthropic/response_parser"
+    require_relative "anthropic/response_parser/completion_parser"
     require_relative "anthropic/format"
     require_relative "anthropic/models"
     include Format
@@ -49,11 +50,11 @@ module LLM
     # @raise [LLM::Error::PromptError]
     #  When given an object a provider does not understand
     # @return (see LLM::Provider#complete)
-    def complete(prompt, role = :user, model: default_model, max_tokens: 1024, **params)
-      params = {max_tokens:, model:}.merge!(params)
+    def complete(prompt, role = :user, model: default_model, max_tokens: 1024, tools: nil, **params)
+      params = [{max_tokens:, model:}, format_tools(tools), params].inject({}, &:merge!).compact
       req = Net::HTTP::Post.new("/v1/messages", headers)
       messages = [*(params.delete(:messages) || []), Message.new(role, prompt)]
-      body = JSON.dump({messages: format(messages)}.merge!(params))
+      body = JSON.dump({messages: [format(messages)].flatten}.merge!(params))
       set_body_stream(req, StringIO.new(body))
       res = request(@http, req)
       Response::Completion.new(res).extend(response_parser)

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

@@ -34,7 +34,7 @@ class LLM::Gemini
     #   res = llm.audio.create_transcription(file: LLM::File("/rocket.mp3"))
     #   res.text # => "A dog on a rocket to the moon"
     # @see https://ai.google.dev/gemini-api/docs/audio Gemini docs
-    # @param [LLM::File, LLM::Response::File] file The input audio
+    # @param [String, LLM::File, LLM::Response::File] file The input audio
     # @param [String] model The model to use
     # @param [Hash] params Other parameters (see Gemini docs)
     # @raise (see LLM::Provider#request)
@@ -43,7 +43,7 @@ class LLM::Gemini
       res = @provider.complete [
         "Your task is to transcribe the contents of an audio file",
         "Your response should include the transcription, and nothing else",
-        file
+        LLM.File(file)
       ], :user, model:, **params
       LLM::Response::AudioTranscription
         .new(res)
@@ -58,7 +58,7 @@ class LLM::Gemini
     #   res = llm.audio.create_translation(file: LLM::File("/bismillah.mp3"))
     #   res.text # => "In the name of Allah, the Beneficent, the Merciful."
     # @see https://ai.google.dev/gemini-api/docs/audio Gemini docs
-    # @param [LLM::File, LLM::Response::File] file The input audio
+    # @param [String, LLM::File, LLM::Response::File] file The input audio
     # @param [String] model The model to use
     # @param [Hash] params Other parameters (see Gemini docs)
     # @raise (see LLM::Provider#request)
@@ -67,7 +67,7 @@ class LLM::Gemini
       res = @provider.complete [
         "Your task is to translate the contents of an audio file into English",
         "Your response should include the translation, and nothing else",
-        file
+        LLM.File(file)
       ], :user, model:, **params
       LLM::Response::AudioTranslation
         .new(res)