omniai 0.1.5 → 1.0.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4e387dddb26068bab8b4aa35a8e990d6db7aed7f034ae4593105b680282e5d43
-  data.tar.gz: 20f9b49d315dea539937521507058a16b8d6f8b0aaeea26f1b6be2cacf07122c
+  metadata.gz: a227888b197b576509bec7208b79f125e2c55baac317ce32b3fed8e2101f11cc
+  data.tar.gz: 16fe842b15ef266556c9887b1625c32a62edbbdba21dea0a8b21e418d76cd435
 SHA512:
-  metadata.gz: 8c3f503fe81921c2d690f169bc22d3d55b00b590920dc31c47c625a93dd142c995e369a9ffa3675c82bf7d04d3861e3b6e09b98b6e3afbd72bfed91b606e8473
-  data.tar.gz: 443d264b0f4c39058fd759b379e604cb631435334fff5fe7528b06cf8502af40c2fbfdca86300e5135029e90ab6104bb0c2b86e5198bd69fcd9c8071e6fbc07a
+  metadata.gz: a4bc067e862d142036dcc15721fb7f3cf5084112145676b3884977e5d3e7025f34637eab3520ad2a1f2d524909f5e8a55b7a2bee35c139062cb481ffc64b3071
+  data.tar.gz: c95602154261e59d7e42fbc94d7afb42f836f54287c4842c0c982a82ced52271ac8c41afbb6cb620865ca948d55971c789c94967f5f18e2af8305271e1c2885b

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]
+      # @return [OmniAI::Chat::Stream]
       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,14 @@ module OmniAI
   #
   # @param client [OmniAI::Client] the client
   class Chat
+    JSON_PROMPT = 'Respond with valid JSON. Do not include any non-JSON in the response.'
+
+    module Role
+      ASSISTANT = 'assistant'
+      USER = 'user'
+      SYSTEM = 'system'
+    end
+
     def initialize(client:)
       @client = client
     end

data/lib/omniai/client.rb

@@ -27,6 +27,12 @@ module OmniAI
       @logger = logger
     end
 
+    # @return [String]
+    def inspect
+      masked_api_key = "#{api_key[..2]}***" if api_key
+      "#<#{self.class.name} api_key=#{masked_api_key.inspect}>"
+    end
+
     # @return [HTTP::Client]
     def connection
       raise NotImplementedError, "#{self.class.name}#connection undefined"

data/lib/omniai/config.rb

@@ -4,5 +4,11 @@ module OmniAI
   # A configuration for each agent w/ `api_key` / `host` / `logger`.
   class Config
     attr_accessor :api_key, :host, :logger
+
+    # @return [String]
+    def inspect
+      masked_api_key = "#{api_key[..2]}***" if api_key
+      "#<#{self.class.name} api_key=#{masked_api_key.inspect} host=#{host.inspect}>"
+    end
   end
 end

data/lib/omniai/version.rb

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

metadata

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