omniai 0.1.4 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5d90d000047c670a17478146fa35fca9c993127067b01b11113cc46e739d15e2
-  data.tar.gz: d025855eb6ea4ae79bc7c126a2853fff7a6a0e21d597fa86b2efb6b8d1c3ddf1
+  metadata.gz: 1694fba08fa9ef9bd8dc3496ce1acbc3c321357d082303869df74ee667a83a65
+  data.tar.gz: c0d98caa29e034c610d20654cc175e42eeb69cef414ba79a5fda61d883d2d9f3
 SHA512:
-  metadata.gz: 77247cd74f88fdaa527b377940f32cfaad187c26f8ae4c6f373c708c866ce5c956c0436194786acb474ae07f423ea83511df6bbd5029c91247aff4b5d8ee2891
-  data.tar.gz: c7d3a40a6c21fe1827b162e94eae0dd0d2f66014797e4d878fa6c7aacf17c443412fcefa2f4bc0ad008f7b64b41ab4e2d583109990d6418d77451f1ba67b4fbd
+  metadata.gz: ddfcd3c15811a63b401290577b36170573e4e170064a8cf64686ceacfc1ddb2c3e5af10ab2a7bb1057db9fe691b8359a94da24c0ce774f71ed5df50063b3925b
+  data.tar.gz: 66dcd111279e7de048473789cfae9020a0c6bb9b812bdd10a1bb2c2c12ee246bd9cebf6b7890322fdcc4a72743033746bd9e5a3f8d283b43ae0bbea5315de05e

data/lib/omniai/chat/request.rb

@@ -6,7 +6,7 @@ module OmniAI
     #
     # Usage:
     #
-    #    class OmniAI::OpenAI::ChatGPT::Completion < OmniAI::Chat::Completion
+    #    class OmniAI::OpenAI::ChatGPT::Request < OmniAI::Chat::Request
     #      module Model
     #        CHAT = "davinci"
     #      end
@@ -23,7 +23,7 @@ module OmniAI
       # @param model [String] required
       # @param temperature [Float, nil] optional
       # @param stream [Proc, nil] optional
-      # @param format [Symbol, nil] optional - :text or :json
+      # @param format [Symbol, nil] optional - :json
       def initialize(client:, messages:, model:, temperature: nil, stream: nil, format: nil)
         @client = client
         @messages = messages
@@ -43,39 +43,54 @@ module OmniAI
 
       protected
 
+      # @return [Hash]
+      def payload
+        raise NotImplementedError, "#{self.class.name}#payload undefined"
+      end
+
+      # @return [String]
+      def path
+        raise NotImplementedError, "#{self.class.name}#path undefined"
+      end
+
       # @param response [HTTP::Response]
-      # @return [OmniAI::Chat::Completion::Response]
+      # @return [OmniAI::Chat::Completion]
       def parse!(response:)
         if @stream
           stream!(response:)
         else
-          OmniAI::OpenAI::Chat::Completion.new(data: response.parse)
+          complete!(response:)
         end
       end
 
+      # @param response [OmniAI::Chat::Completion]
+      def complete!(response:)
+        Completion.new(data: response.parse)
+      end
+
       # @param response [HTTP::Response]
       # @return [OmniAI::Chat::Chunk]
       def stream!(response:)
-        parser = EventStreamParser::Parser.new
+        raise Error, "#{self.class.name}#stream! unstreamable" unless @stream
 
-        response.body.each do |chunk|
-          parser.feed(chunk) do |_, data|
-            break if data.eql?('[DONE]')
+        Stream.new(response:).stream! { |chunk| @stream.call(chunk) }
+      end
 
-            chunk = OmniAI::OpenAI::Chat::Chunk.new(data: JSON.parse(data))
-            @stream.call(chunk)
+      # @return [Array<Hash>]
+      def messages
+        arrayify(@messages).map do |content|
+          case content
+          when String then { role: OmniAI::Chat::Role::USER, content: }
+          when Hash then content
+          else raise Error, "Unsupported content=#{content.inspect}"
           end
         end
       end
 
-      # @return [Hash]
-      def payload
-        raise NotImplementedError, "#{self.class.name}#payload undefined"
-      end
-
-      # @return [String]
-      def path
-        raise NotImplementedError, "#{self.class.name}#path undefined"
+      # @param value [Object, Array<Object>]
+      # @return [Array<Object>]
+      def arrayify(value)
+        value.is_a?(Array) ? value : [value]
       end
 
       # @return [HTTP::Response]

data/lib/omniai/chat/stream.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # A stream given when streaming.
+    class Stream
+      # @param response [HTTP::Response]
+      # @param block [Proc]
+      def initialize(response:)
+        @response = response
+        @parser = EventStreamParser::Parser.new
+      end
+
+      # @yield [OmniAI::Chat::Chunk]
+      def stream!
+        @response.body.each do |chunk|
+          @parser.feed(chunk) do |_, data|
+            next if data.eql?('[DONE]')
+
+            yield(Chunk.new(data: JSON.parse(data)))
+          end
+        end
+      end
+    end
+  end
+end

data/lib/omniai/chat/usage.rb

@@ -11,20 +11,30 @@ module OmniAI
         @data = data
       end
 
-      # @return [Integer]
+      # @return [Integer, nil]
       def completion_tokens
         @data['completion_tokens']
       end
 
-      # @return [Integer]
+      # @return [Integer, nil]
       def prompt_tokens
         @data['prompt_tokens']
       end
 
-      # @return [Integer]
+      # @return [Integer, nil]
       def total_tokens
         @data['total_tokens']
       end
+
+      # @return [Integer, nil]
+      def input_tokens
+        @data['input_tokens']
+      end
+
+      # @return [Integer, nil]
+      def output_tokens
+        @data['output_tokens']
+      end
     end
   end
 end

data/lib/omniai/chat.rb

@@ -17,6 +17,12 @@ module OmniAI
   #
   # @param client [OmniAI::Client] the client
   class Chat
+    module Role
+      ASSISTANT = 'assistant'
+      USER = 'user'
+      SYSTEM = 'system'
+    end
+
     def initialize(client:)
       @client = client
     end

data/lib/omniai/config.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module OmniAI
   # A configuration for each agent w/ `api_key` / `host` / `logger`.
   class Config

data/lib/omniai/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module OmniAI
-  VERSION = '0.1.4'
+  VERSION = '0.2.0'
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 0.1.4
+  version: 0.2.0
 platform: ruby
 authors:
 - Kevin Sylvestre
@@ -73,6 +73,7 @@ files:
 - lib/omniai/chat/delta.rb
 - lib/omniai/chat/message.rb
 - lib/omniai/chat/request.rb
+- lib/omniai/chat/stream.rb
 - lib/omniai/chat/usage.rb
 - lib/omniai/client.rb
 - lib/omniai/config.rb