lib/omniai/chat/usage.rb in omniai-1.0.2 vs lib/omniai/chat/usage.rb in omniai-1.0.3

@@ -1,40 +1,34 @@
 # frozen_string_literal: true
 
 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
 end