omniai 1.0.3 → 1.0.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9ef66e140869ad45643072513852d7ed06439fb6a7fecda4f552ecb34c14ddb1
-  data.tar.gz: 55740c3bdc98c9d83330b1a45b1b83d760b279d388830e83fce8a671f635d26d
+  metadata.gz: dddfc4d4ac08e936e1b2443e37ff93b4256428579aa1280b6ede769a01365f04
+  data.tar.gz: 235c5f109e0e3526d2abfa2cf7d6c3a3d1e00749da70beaf884a2a1bcb9ba8c6
 SHA512:
-  metadata.gz: f335ee44a2460431d3f7bf83991448569f451930102ba1d3714c79af225db08d22467269c1543434a987de2d0f70bbe33faece99c4f291cc67032e7a5121721b
-  data.tar.gz: 33ca9dbf52c1a7af1c60683db4fd31680f4e3230da6c139591c6c62796616767b67109862963d8b1b7b93c6af46143e140471a606b20f0a6457b3f1989c87030
+  metadata.gz: fae66992c4f0a399ecbe795a39901e73caea1554861e23894fd005daab9b26e23cd6a110cca9cde9b9c9ce73af174037626b3505074bb4bd23477677a286b3ef
+  data.tar.gz: 5f604bcdc4ca4b748e523eff6cb4b383cd9211188673c768d2295a7f9e43c1f043ac9c8b14ff86b9dabb85a99b50a005ce40b920cdbcdf19f9f1ef05b4f33ca5

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

@@ -35,15 +35,7 @@ module OmniAI
       def 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:
-        )
+        @usage ||= Usage.for(data: @data['usage'])
       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

@@ -6,6 +6,20 @@ module OmniAI
     class Usage
       attr_accessor :input_tokens, :output_tokens, :total_tokens
 
+      # @param data [Hash]
+      # @return [OmniAI::Chat::Usage]
+      def self.for(data:)
+        input_tokens = data['input_tokens'] || data['prompt_tokens']
+        output_tokens = data['output_tokens'] || data['completion_tokens']
+        total_tokens = data['total_tokens'] || (input_tokens + output_tokens)
+
+        new(
+          input_tokens:,
+          output_tokens:,
+          total_tokens:
+        )
+      end
+
       # @param input_tokens [Integer]
       # @param output_tokens [Integer]
       # @param total_tokens [Integer]

data/lib/omniai/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 1.0.3
+  version: 1.0.5
 platform: ruby
 authors:
 - Kevin Sylvestre