omniai 1.0.1 → 1.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b730453eea6e211031174bde7e43485b5402b9f0590bb832924762f8c2f6345c
-  data.tar.gz: cad17682fc35243454f58b28a9bf0e2d48f9857415bfe3b72012f370b7e9d4ab
+  metadata.gz: 9ef66e140869ad45643072513852d7ed06439fb6a7fecda4f552ecb34c14ddb1
+  data.tar.gz: 55740c3bdc98c9d83330b1a45b1b83d760b279d388830e83fce8a671f635d26d
 SHA512:
-  metadata.gz: 0ebd9716858486dc399504b8caefd3d4186cdc97da05715e8519fcbc2238b38672f6a760e5134b3231bfc4683662f69098edc3eb6bd22808e91658bb4207b1b4
-  data.tar.gz: b7623a59effdc4714cbf78f8dd950cde0d3746da0dfc39ec7c9a6caf745ffedbd841d0b72097b1fdd30dd73f3f8ba35e7736a4253b32b7d2fcd7367c732172b0
+  metadata.gz: f335ee44a2460431d3f7bf83991448569f451930102ba1d3714c79af225db08d22467269c1543434a987de2d0f70bbe33faece99c4f291cc67032e7a5121721b
+  data.tar.gz: 33ca9dbf52c1a7af1c60683db4fd31680f4e3230da6c139591c6c62796616767b67109862963d8b1b7b93c6af46143e140471a606b20f0a6457b3f1989c87030

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/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.1'
+  VERSION = '1.0.3'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 1.0.1
+  version: 1.0.3
 platform: ruby
 authors:
 - Kevin Sylvestre
-autorequire:
+autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-06-14 00:00:00.000000000 Z
+date: 2024-06-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: event_stream_parser
@@ -83,7 +83,7 @@ metadata:
   homepage_uri: https://github.com/ksylvest/omniai
   changelog_uri: https://github.com/ksylvest/omniai/releases
   rubygems_mfa_required: 'true'
-post_install_message:
+post_install_message: 
 rdoc_options: []
 require_paths:
 - lib
@@ -98,8 +98,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.11
-signing_key:
+rubygems_version: 3.5.3
+signing_key: 
 specification_version: 4
 summary: A generalized framework for interacting with many AI services
 test_files: []