omniai 0.1.5 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4e387dddb26068bab8b4aa35a8e990d6db7aed7f034ae4593105b680282e5d43
-  data.tar.gz: 20f9b49d315dea539937521507058a16b8d6f8b0aaeea26f1b6be2cacf07122c
+  metadata.gz: 1694fba08fa9ef9bd8dc3496ce1acbc3c321357d082303869df74ee667a83a65
+  data.tar.gz: c0d98caa29e034c610d20654cc175e42eeb69cef414ba79a5fda61d883d2d9f3
 SHA512:
-  metadata.gz: 8c3f503fe81921c2d690f169bc22d3d55b00b590920dc31c47c625a93dd142c995e369a9ffa3675c82bf7d04d3861e3b6e09b98b6e3afbd72bfed91b606e8473
-  data.tar.gz: 443d264b0f4c39058fd759b379e604cb631435334fff5fe7528b06cf8502af40c2fbfdca86300e5135029e90ab6104bb0c2b86e5198bd69fcd9c8071e6fbc07a
+  metadata.gz: ddfcd3c15811a63b401290577b36170573e4e170064a8cf64686ceacfc1ddb2c3e5af10ab2a7bb1057db9fe691b8359a94da24c0ce774f71ed5df50063b3925b
+  data.tar.gz: 66dcd111279e7de048473789cfae9020a0c6bb9b812bdd10a1bb2c2c12ee246bd9cebf6b7890322fdcc4a72743033746bd9e5a3f8d283b43ae0bbea5315de05e

data/lib/omniai/chat/request.rb

@@ -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]
       def parse!(response:)
         if @stream
           stream!(response:)
         else
-          OmniAI::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::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/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  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