lib/omniai/chat/usage.rb in omniai-0.1.5 vs lib/omniai/chat/usage.rb in omniai-0.2.0

@@ -9,22 +9,32 @@
       # @param data [Hash]
       def initialize(data:)
         @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