omniai 1.6.6 → 1.8.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0bc92ec17e5ca87590eca2181453b6a3e54c2d15bb6dda0c80c604a443e59660
-  data.tar.gz: e4eb4ad22437711fbac2c538edc644861107985662f24a236b317d229802f7b7
+  metadata.gz: 7325d58430db1dd52ecbbf7964ad9a02a3fe5ae89be51779582a672bbc8c8222
+  data.tar.gz: 2464d2ae08fc94fad891acae6eb1ec7fbe163414a3dfb6583fd89bca52dc92d1
 SHA512:
-  metadata.gz: 75b7b695353ec0899b786ee057beaedcd8ab21ba15efca6f3a21c4e22730dc73b7676fbdfece11f624a1dce77638ad2d7cc4a6ff054c2180edb8f423b20ca6ee
-  data.tar.gz: 9c7ab97db7e4d8085774c9fd76803b8abfe53905b6d9b6d141caed311d019347d3f0e880e957597b529de8e479d1b69415213e82ca1daba34bf00d7049eddbae
+  metadata.gz: 2d966778690e811bab403217378a748c778e043c48cdb946d9f2b3503ae3b48a165c63ab57bbba65ffa864cd9395388181b57b3c689efb548b864d58f931a757
+  data.tar.gz: a82e9b1b092c4df179a4d2b931f6d2c217ba1976504fd40179bc9c0549ab9c9fd349e3eedb79df2c3e30cde1a1ea3b236b99449074c10da381171e1d560314e8

data/Gemfile

@@ -4,11 +4,13 @@ source 'https://rubygems.org'
 
 gemspec
 
+gem 'factory_bot'
 gem 'logger'
 gem 'rake'
 gem 'rspec'
 gem 'rspec_junit_formatter'
 gem 'rubocop'
+gem 'rubocop-factory_bot'
 gem 'rubocop-rake'
 gem 'rubocop-rspec'
 gem 'simplecov'

data/README.md

@@ -128,7 +128,7 @@ Generating a completion is as simple as sending in the text:
 
 ```ruby
 completion = client.chat('Tell me a joke.')
-completion.choice.message.content # 'Why don't scientists trust atoms? They make up everything!'
+completion.text # 'Why don't scientists trust atoms? They make up everything!'
 ```
 
 #### Completions using a Complex Prompt
@@ -145,7 +145,7 @@ completion = client.chat do |prompt|
     message.file('./hamster.jpeg', "image/jpeg")
   end
 end
-completion.choice.message.content  # 'They are photos of a cat, a cat, and a hamster.'
+completion.text  # 'They are photos of a cat, a cat, and a hamster.'
 ```
 
 #### Completions using Streaming via Proc
@@ -154,7 +154,7 @@ A real-time stream of messages can be generated by passing in a proc:
 
 ```ruby
 stream = proc do |chunk|
-  print(chunk.choice.delta.content) # '...'
+  print(chunk.text) # '...'
 end
 client.chat('Tell me a joke.', stream:)
 ```
@@ -229,6 +229,29 @@ tempfile.close
 tempfile.unlink
 ```
 
+### Embeddings
+
+Clients that support generating embeddings (e.g. OpenAI, Mistral, etc.) convert text to embeddings via the following:
+
+```ruby
+response = client.embed('The quick brown fox jumps over a lazy dog')
+response.usage # <OmniAI::Embed::Usage prompt_tokens=5 total_tokens=5>
+response.embedding # [0.1, 0.2, ...] >
+```
+
+Batches of text can also be converted to embeddings via the following:
+
+```ruby
+response = client.embed([
+  '',
+  '',
+])
+response.usage # <OmniAI::Embed::Usage prompt_tokens=5 total_tokens=5>
+response.embeddings.each do |embedding|
+  embedding # [0.1, 0.2, ...]
+end
+```
+
 ## CLI
 
 OmniAI packages a basic command line interface (CLI) to allow for exploration of various APIs. A detailed CLI documentation can be found via help:
@@ -263,3 +286,32 @@ Type 'exit' or 'quit' to abort.
 ```
 The warmest place on earth is Africa.
 ```
+
+### Embed
+
+#### w/ input
+
+```bash
+omniai embed "The quick brown fox jumps over a lazy dog."
+```
+
+```
+0.0
+...
+```
+
+#### w/o input
+
+```bash
+omniai embed --provider="openai" --model="text-embedding-ada-002"
+```
+
+```
+Type 'exit' or 'quit' to abort.
+# Whe quick brown fox jumps over a lazy dog.
+```
+
+```
+0.0
+...
+```

data/lib/omniai/chat/choice.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # For for.
+    class Choice
+      # @return [Integer]
+      attr_accessor :index
+
+      # @return [Message]
+      attr_accessor :message
+
+      # @param message [Message]
+      # @param index [Integer]
+      # @param tool_call_list [Array<ToolCall>]
+      def initialize(message:, index: 0)
+        @message = message
+        @index = index
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} index=#{@index} message=#{@message.inspect}>"
+      end
+
+      # @param data [Hash]
+      # @param context [OmniAI::Context] optional
+      #
+      # @return [Choice]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:choice)
+        return deserialize.call(data, context:) if deserialize
+
+        index = data['index']
+        message = Message.deserialize(data['message'] || data['delta'], context:)
+
+        new(message:, index:)
+      end
+
+      # @param context [OmniAI::Context] optional
+      # @return [Hash]
+      def serialize(context: nil)
+        serialize = context&.serializer(:choice)
+        return serialize.call(self, context:) if serialize
+
+        {
+          index:,
+          message: message.serialize(context:),
+        }
+      end
+
+      # @return [Message]
+      def delta
+        message
+      end
+
+      # @return [Array<Content>, String]
+      def content
+        message.content
+      end
+
+      # @return [Array<ToolCall, nil]
+      def tool_call_list
+        message.tool_call_list
+      end
+    end
+  end
+end

data/lib/omniai/chat/content.rb

@@ -16,14 +16,22 @@ module OmniAI
       #
       # @return [String]
       def serialize(context: nil)
-        raise NotImplementedError, ' # {self.class}#serialize undefined'
+        raise NotImplementedError, "#{self.class}#serialize undefined"
       end
 
-      # @param data [hash]
+      # @param data [Hash, Array, String]
       # @param context [Context] optional
       #
       # @return [Content]
       def self.deserialize(data, context: nil)
+        return data if data.nil?
+        return data.map { |data| deserialize(data, context:) } if data.is_a?(Array)
+
+        deserialize = context&.deserializer(:content)
+        return deserialize.call(data, context:) if deserialize
+
+        return data if data.is_a?(String)
+
         raise ArgumentError, "untyped data=#{data.inspect}" unless data.key?('type')
 
         case data['type']

data/lib/omniai/chat/file.rb

@@ -21,8 +21,8 @@ module OmniAI
       # @return [String]
       def fetch!
         case @io
-        when IO then @io.read
-        else ::File.binread(@io)
+        when String then ::File.binread(@io)
+        else @io.read
         end
       end
 
@@ -33,7 +33,7 @@ module OmniAI
           content = fetch!
           Text.new("<file>#{filename}: #{content}</file>").serialize(context:)
         else
-          serializer = context&.serializers&.[](:file)
+          serializer = context&.serializer(:file)
           return serializer.call(self, context:) if serializer
 
           { type: "#{kind}_url", "#{kind}_url": { url: data_uri } }

data/lib/omniai/chat/function.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # A function that includes a name / arguments.
+    class Function
+      # @return [String]
+      attr_accessor :name
+
+      # @return [Hash]
+      attr_accessor :arguments
+
+      # @param name [String]
+      # @param arguments [Hash]
+      def initialize(name:, arguments:)
+        @name = name
+        @arguments = arguments
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} name=#{name.inspect} arguments=#{arguments.inspect}>"
+      end
+
+      # @param data [Hash]
+      # @param context [Context] optional
+      #
+      # @return [Function]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:function)
+        return deserialize.call(data, context:) if deserialize
+
+        name = data['name']
+        arguments = begin
+          JSON.parse(data['arguments']) if data['arguments']
+        rescue JSON::ParserError
+          data['arguments']
+        end
+
+        new(name:, arguments:)
+      end
+
+      # @param context [Context] optional
+      #
+      # @return [Hash]
+      def serialize(context: nil)
+        serializer = context&.serializer(:function)
+        return serializer.call(self, context:) if serializer
+
+        {
+          name: @name,
+          arguments: (JSON.generate(@arguments) if @arguments),
+        }
+      end
+    end
+  end
+end

data/lib/omniai/chat/message/builder.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    class Message
+      # Used to build a message.
+      class Builder
+        # @param role [String]
+        # @return [Message]
+        def self.build(role:)
+          builder = new(role:)
+          yield(builder)
+          builder.build
+        end
+
+        # @param role [String]
+        def initialize(role: Role::USER)
+          @role = role
+          @content = []
+        end
+
+        # @return [Message]
+        def build
+          Message.new(content: @content, role: @role)
+        end
+
+        # @example
+        #   message.text('What are these photos of?')
+        #
+        # @param value [String]
+        #
+        # @return [Text]
+        def text(value)
+          Text.new(value).tap do |text|
+            @content << text
+          end
+        end
+
+        # @example
+        #   message.url('https://example.com/hamster.jpg', type: "image/jpeg")
+        #
+        # @param uri [String]
+        # @param type [String]
+        #
+        # @return [URL]
+        def url(uri, type)
+          URL.new(uri, type).tap do |url|
+            @content << url
+          end
+        end
+
+        # @example
+        #   message.file(File.open('hamster.jpg'), type: "image/jpeg")
+        #
+        # @param io [IO]
+        # @param type [String]
+        #
+        # @return [File]
+        def file(io, type)
+          File.new(io, type).tap do |file|
+            @content << file
+          end
+        end
+      end
+    end
+  end
+end

data/lib/omniai/chat/message.rb

@@ -10,20 +10,50 @@ module OmniAI
     #       message.url 'https://example.com/cat.jpg', type: "image/jpeg"
     #       message.url 'https://example.com/dog.jpg', type: "image/jpeg"
     #       message.file File.open('hamster.jpg'), type: "image/jpeg"
+    #       message.
     #     end
     #   end
     class Message
+      # @example
+      #   prompt.build('What is the capital of Canada?')
+      #
+      # @example
+      #   prompt.build(role: 'user') do |message|
+      #     message.text 'What is the capital of Canada?'
+      #   end
+      #
+      # @param content [String, nil]
+      # @param role [Symbol]
+      #
+      # @yield [builder]
+      # @yieldparam builder [Message::Builder]
+      #
+      # @return [Message]
+      def self.build(content = nil, role: Role::USER, &block)
+        raise ArgumentError, 'content or block is required' if content.nil? && block.nil?
+
+        Builder.build(role:) do |builder|
+          builder.text(content) if content
+          block&.call(builder)
+        end
+      end
+
       # @return [Array<Content>, String]
       attr_accessor :content
 
       # @return [String]
       attr_accessor :role
 
+      # @return [Array<ToolCall>, nil]
+      attr_accessor :tool_call_list
+
       # @param content [String, nil]
       # @param role [String]
-      def initialize(content: nil, role: Role::USER)
-        @content = content || []
+      # @param tool_call_list [Array<ToolCall>, nil]
+      def initialize(content:, role: Role::USER, tool_call_list: nil)
+        @content = content
         @role = role
+        @tool_call_list = tool_call_list
       end
 
       # @return [String]
@@ -48,13 +78,14 @@ module OmniAI
       #
       # @return [Message]
       def self.deserialize(data, context: nil)
-        deserialize = context&.deserializers&.[](:message)
+        deserialize = context&.deserializer(:message)
         return deserialize.call(data, context:) if deserialize
 
-        new(
-          content: data['content'].map { |content| Content.deserialize(content, context:) },
-          role: data['role']
-        )
+        role = data['role']
+        content = Content.deserialize(data['content'], context:)
+        tool_call_list = data['tool_calls']&.map { |subdata| ToolCall.deserialize(subdata, context:) }
+
+        new(content:, role:, tool_call_list:)
       end
 
       # Usage:
@@ -66,12 +97,13 @@ module OmniAI
       #
       # @return [Hash]
       def serialize(context: nil)
-        serializer = context&.serializers&.[](:message)
+        serializer = context&.serializer(:message)
         return serializer.call(self, context:) if serializer
 
-        content = @content.is_a?(String) ? @content : @content.map { |content| content.serialize(context:) }
+        content = @content.is_a?(Array) ? @content.map { |content| content.serialize(context:) } : @content
+        tool_calls = @tool_call_list&.map { |tool_call| tool_call.serialize(context:) }
 
-        { role: @role, content: }
+        { role: @role, content:, tool_calls: }.compact
       end
 
       # @return [Boolean]
@@ -89,45 +121,32 @@ module OmniAI
         role?(Role::USER)
       end
 
-      # Usage:
-      #
-      #   message.text('What are these photos of?')
-      #
-      # @param value [String]
-      #
-      # @return [Text]
-      def text(value)
-        Text.new(value).tap do |text|
-          @content << text
-        end
+      # @return [Boolean]
+      def tool?
+        role?(Role::TOOL)
       end
 
-      # Usage:
-      #
-      #   message.url('https://example.com/hamster.jpg', type: "image/jpeg")
-      #
-      # @param uri [String]
-      # @param type [String]
-      #
-      # @return [URL]
-      def url(uri, type)
-        URL.new(uri, type).tap do |url|
-          @content << url
-        end
+      # @return [Boolean]
+      def text?
+        !text.nil?
       end
 
-      # Usage:
-      #
-      #   message.file(File.open('hamster.jpg'), type: "image/jpeg")
-      #
-      # @param io [IO]
-      # @param type [String]
-      #
-      # @return [File]
-      def file(io, type)
-        File.new(io, type).tap do |file|
-          @content << file
-        end
+      # @return [String, nil]
+      def text
+        return if @content.nil?
+        return @content if @content.is_a?(String)
+
+        parts = arrayify(@content).filter { |content| content.is_a?(Text) }
+        parts.map(&:text).join("\n") unless parts.empty?
+      end
+
+      # @param object [Object]
+      # @return [Array]
+      def arrayify(object)
+        return if object.nil?
+        return object if object.is_a?(Array)
+
+        [object]
       end
     end
   end

data/lib/omniai/chat/payload.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # A chunk or completion.
+    class Payload
+      # @return [Array<Choice>]
+      attr_accessor :choices
+
+      # @return [Usage, nil]
+      attr_accessor :usage
+
+      # @param choices [Array<Choice>]
+      # @param usage [Usage, nil]
+      def initialize(choices:, usage: nil)
+        @choices = choices
+        @usage = usage
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} choices=#{choices.inspect} usage=#{usage.inspect}>"
+      end
+
+      # @param data [Hash]
+      # @param context [OmniAI::Context] optional
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:payload)
+        return deserialize.call(data, context:) if deserialize
+
+        choices = data['choices'].map { |choice_data| Choice.deserialize(choice_data, context:) }
+        usage = Usage.deserialize(data['usage'], context:) if data['usage']
+
+        new(choices:, usage:)
+      end
+
+      # @param context [OmniAI::Context] optional
+      # @return [Hash]
+      def serialize(context:)
+        serialize = context&.serializer(:payload)
+        return serialize.call(self, context:) if serialize
+
+        {
+          choices: choices.map { |choice| choice.serialize(context:) },
+          usage: usage&.serialize(context:),
+        }
+      end
+
+      # @param index [Integer]
+      # @return [Choice]
+      def choice(index: 0)
+        @choices[index]
+      end
+
+      # @param index [Integer]
+      # @return [Message]
+      def message(index: 0)
+        choice(index:).message
+      end
+
+      # @return [Array<Message>]
+      def messages
+        @choices.map(&:message)
+      end
+
+      # @param index [Integer]
+      # @return [String, nil]
+      def text(index: 0)
+        message(index:).text
+      end
+
+      # @param index [Integer]
+      # @return [Boolean]
+      def text?(index: 0)
+        message(index:).text?
+      end
+
+      # @param index [Integer]
+      # @return [Array<ToolCall>]
+      def tool_call_list(index:)
+        message(index:).tool_call_list
+      end
+    end
+  end
+end

data/lib/omniai/chat/prompt.rb

@@ -57,6 +57,11 @@ module OmniAI
         @messages = messages
       end
 
+      # @return [Prompt]
+      def dup
+        self.class.new(messages: @messages.dup)
+      end
+
       # @return [String]
       def inspect
         "#<#{self.class.name} messages=#{@messages.inspect}>"
@@ -75,57 +80,66 @@ module OmniAI
       #
       # @return [Array<Hash>]
       def serialize(context: nil)
-        serializer = context&.serializers&.[](:prompt)
+        serializer = context&.serializer(:prompt)
         return serializer.call(self, context:) if serializer
 
         @messages.map { |message| message.serialize(context:) }
       end
 
-      # Usage:
+      # @example
+      #   prompt.message 'What is the capital of Canada?', role: '...'
       #
-      #   prompt.message('What is the capital of Canada?')
+      # @example
+      #   prompt.message role: '...' do |message|
+      #     message.text 'What is the capital of Canada?'
+      #   end
       #
       # @param content [String, nil]
       # @param role [Symbol]
       #
-      # @yield [Message]
+      # @yield [builder]
+      # @yieldparam builder [Message::Builder]
+      #
       # @return [Message]
-      def message(content = nil, role: :user, &block)
-        raise ArgumentError, 'content or block is required' if content.nil? && block.nil?
-
-        Message.new(content:, role:).tap do |message|
-          block&.call(message)
+      def message(content = nil, role: Role::USER, &)
+        Message.build(content, role:, &).tap do |message|
           @messages << message
         end
       end
 
-      # Usage:
-      #
-      #   prompt.system('You are a helpful assistant.')
+      # @example
+      #   prompt.system 'You are a helpful assistant.'
       #
+      # @example
       #   prompt.system do |message|
       #     message.text 'You are a helpful assistant.'
       #   end
       #
       # @param content [String, nil]
       #
-      # @yield [Message]
+      # @yield [builder]
+      # @yieldparam builder [Message::Builder]
+      #
       # @return [Message]
       def system(content = nil, &)
         message(content, role: Role::SYSTEM, &)
       end
 
-      # Usage:
-      #
+      # @example
       #   prompt.user('What is the capital of Canada?')
       #
+      # @example
       #   prompt.user do |message|
       #     message.text 'What is the capital of Canada?'
+      #     message.url 'https://...', type: "image/gif"
+      #     message.file File.open('...'), type: "image/gif"
       #   end
       #
       # @param content [String, nil]
       #
-      # @yield [Message]
+      # @yield [builder]
+      # @yieldparam builder [Message::Builder]
+      #
       # @return [Message]
       def user(content = nil, &)
         message(content, role: Role::USER, &)