omniai 1.0.2 → 1.0.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ffc27287a50c39e3aab00622ea1e389f2e2da0828e5ec5b01f0e62bf7ec63c58
-  data.tar.gz: 3979a4a92abd12706cc36ff66d4befab252a48a2deb7d5b3db3a81162af41b50
+  metadata.gz: 6d821530b3b40043754c26a175c2b5d9c70fe3f04bd62bf61c5506963e56dae3
+  data.tar.gz: 6f0cf3e3292b11492bd25526391ec3a396d5a9262050dcf4d5aba35726c2ed04
 SHA512:
-  metadata.gz: 32f6f70655128b1eb60bcc1add45eb53cd312550a521fd3a7a7663f09f4084809d139ecb2a617464266b12199f590db26bfa240be9c5fb90b6d1e264b1bd57c9
-  data.tar.gz: d7b170ae465c6e72e4faa5fde86a216bd8317205c9c1630753f527424aa9e59b60f6f1dca32fcecd242824485f7fc328368d9fd8bcff92658fdfd078a5d4080f
+  metadata.gz: 9f2d2fa897c94b741435df368ca0952fe00e4dafd0f199756008749a1267ac2d4ea325338b1da359c7c84944b76ec3a875d678039b9bd88ee819cacb470a6fdd
+  data.tar.gz: 698f4ec0159ac6a359f5bd93c02db3e4acd0cde6cfd42c17a30c341d6ad674720aa442040c0890cd3c600068c52717b011b28601335721fbdd7a55c1d74ae61d

data/lib/omniai/chat/choice.rb

@@ -23,7 +23,9 @@ module OmniAI
 
       # @return [OmniAI::Chat::Message]
       def message
-        Message.new(data: @data['message']) if @data['message']
+        return unless @data['message']
+
+        @message ||= Message.for(role: @data['message'])
       end
     end
   end

data/lib/omniai/chat/completion.rb

@@ -33,7 +33,17 @@ module OmniAI
 
       # @return [OmniAI::Chat::Usage]
       def usage
-        @usage ||= Usage.new(data: @data['usage']) if @data['usage']
+        return unless @data['usage']
+
+        input_tokens = @data['usage']['input_tokens'] || @data['usage']['prompt_tokens']
+        output_tokens = @data['usage']['output_tokens'] || @data['usage']['completion_tokens']
+        total_tokens = @data['usage']['total_tokens'] || (input_tokens + output_tokens)
+
+        @usage ||= Usage.new(
+          input_tokens:,
+          output_tokens:,
+          total_tokens:
+        )
       end
 
       # @return [Array<OmniAI::Chat::Choice>]

data/lib/omniai/chat/message.rb

@@ -4,22 +4,27 @@ module OmniAI
   class Chat
     # A message returned by the API.
     class Message
-      attr_accessor :data
+      attr_accessor :role, :content
 
-      # @param content [Integer]
-      # @param role [String]
-      def initialize(data:)
-        @data = data
+      # @param data [Hash]
+      # @return [OmniAI::Chat::Message]
+      def self.for(data:)
+        content = data['content'] || data[:content]
+        role = data['role'] || data[:role]
+
+        new(content:, role: role || Role::USER)
       end
 
-      # @return [String]
-      def role
-        @data['role']
+      # @param content [String]
+      # @param role [String] optional (default to "user") e.g. "assistant" / "user" / "system"
+      def initialize(content:, role: OmniAI::Chat::Role::USER)
+        @role = role
+        @content = content
       end
 
       # @return [String]
-      def content
-        @data['content']
+      def inspect
+        "#<#{self.class.name} role=#{role.inspect} content=#{content.inspect}>"
       end
     end
   end

data/lib/omniai/chat/usage.rb

@@ -3,37 +3,31 @@
 module OmniAI
   class Chat
     # A usage returned by the API.
-    class Completion
-      attr_accessor :data
+    class Usage
+      attr_accessor :input_tokens, :output_tokens, :total_tokens
 
-      # @param data [Hash]
-      def initialize(data:)
-        @data = data
+      # @param input_tokens [Integer]
+      # @param output_tokens [Integer]
+      # @param total_tokens [Integer]
+      def initialize(input_tokens:, output_tokens:, total_tokens:)
+        @input_tokens = input_tokens
+        @output_tokens = output_tokens
+        @total_tokens = total_tokens
       end
 
-      # @return [Integer, nil]
+      # @return [Integer]
       def completion_tokens
-        @data['completion_tokens']
+        @output_tokens
       end
 
-      # @return [Integer, nil]
+      # @return [Integer]
       def prompt_tokens
-        @data['prompt_tokens']
+        @input_tokens
       end
 
-      # @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']
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} input_tokens=#{input_tokens} output_tokens=#{output_tokens} total_tokens=#{total_tokens}>"
       end
     end
   end

data/lib/omniai/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 1.0.2
+  version: 1.0.4
 platform: ruby
 authors:
 - Kevin Sylvestre
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-06-15 00:00:00.000000000 Z
+date: 2024-06-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: event_stream_parser