Breakout input vs output tokens

.rubocop.yml

@@ -49,6 +49,7 @@ Metrics/ClassLength:
     - 'method_call'
   Exclude:
     - 'app/controllers/**/*'
+    - 'app/models/**/*'
     - 'lib/**/*_form.rb'
 Metrics/ModuleLength:
   CountAsOne:

app/models/conversation.rb

@@ -69,6 +69,14 @@ def token_count
     generate_text_requests.sum(&:response_token_count)
   end
 
+  def input_token_count
+    generate_text_requests.sum(&:response_input_token_count)
+  end
+
+  def output_token_count
+    generate_text_requests.sum(&:response_output_token_count)
+  end
+
   def blobify
     [
       title,

app/models/generate_text_request.rb

@@ -55,8 +55,20 @@ def blobify
   end
 
   def response_token_count
+    response_input_token_count + response_output_token_count
+  end
+
+  def response_input_token_count
+    if completed?
+      response.input_token_count
+    else
+      0
+    end
+  end
+
+  def response_output_token_count
     if completed?
-      response.token_count
+      response.output_token_count
     else
       0
     end

app/views/components/prompt_form_component/prompt_form_component.html.haml

@@ -76,16 +76,21 @@
 
           .control-buttons-section
             .d-flex.align-items-center.justify-content-end.gap-2
-              .token-display.d-flex.align-items-center
-                %span.label.me-1.small Tokens:
-                %span.value.badge.bg-secondary= conversation.token_count
+              .token-display.d-flex.align-items-center.gap-2
+                #input-tokens.token-item.d-flex.align-items-center{ data: { 'bs-toggle': 'tooltip', 'bs-title': 'Input tokens' }}
+                  %i.bi.bi-arrow-down-circle.text-primary
+                  %span.value.badge.bg-primary= conversation.input_token_count
+                #output-tokens.token-item.d-flex.align-items-center{ data: { 'bs-toggle': 'tooltip', 'bs-title': 'Output tokens' }}
+                  %i.bi.bi-arrow-up-circle.text-success
+                  %span.value.badge.bg-success= conversation.output_token_count
               .button-group.d-flex.gap-1
                 - if conversation.persisted?
                   %button.context-btn.btn.icon-btn.icon-btn-sm{ type: 'button', data: { 'bs-toggle': 'modal', 'bs-target': '#convo-context-modal' }}
                     %i.bi.bi-folder-plus{ data: { 'bs-toggle': 'tooltip', 'bs-title': 'Conversation Context' }}
                 %button.settings-btn.btn.icon-btn.icon-btn-sm{ type: 'button', data: { 'bs-toggle': 'modal', 'bs-target': '#convo-settings-modal' }}
                   %i.bi.bi-gear{ data: { 'bs-toggle': 'tooltip', 'bs-title': 'Conversation Settings' }}
 
+
   / CONVERSATION SETTINGS MODAL
 
   .modal.settings-modal#convo-settings-modal{ tabindex: '-1', 'aria-labelledby': 'conversationSettingsModalLabel', 'aria-hidden': 'true' }

lib/anthropic/invoke_model_response.rb

@@ -65,10 +65,18 @@ def tool_use?
 
     # Cache token keys were added later. Older responses won't have that key.
     def token_count
-      usage.fetch('input_tokens') + usage.fetch('output_tokens') +
+      input_token_count + output_token_count
+    end
+
+    def input_token_count
+      usage.fetch('input_tokens') +
         usage.fetch('cache_creation_input_tokens', 0) + usage.fetch('cache_read_input_tokens', 0)
     end
 
+    def output_token_count
+      usage.fetch('output_tokens')
+    end
+
     def usage
       data.fetch('usage')
     end

lib/generative_text/aws/invoke_model_response.rb

@@ -33,7 +33,15 @@ def completion_reason
       end
 
       def token_count
-        (data['inputTextTokenCount'] || 0) + (results['tokenCount'] || 0)
+        input_token_count + output_token_count
+      end
+
+      def input_token_count
+        data['inputTextTokenCount'] || 0
+      end
+
+      def output_token_count
+        results['tokenCount'] || 0
       end
     end
   end

spec/models/generate_text_request_spec.rb

@@ -125,7 +125,9 @@
         build(:generate_text_request, :with_response, :completed)
       end
 
-      let(:response_obj) { instance_double(Anthropic::InvokeModelResponse, token_count: 100) }
+      let(:response_obj) do
+        instance_double(Anthropic::InvokeModelResponse, input_token_count: 50, output_token_count: 50)
+      end
 
       before do
         allow(Anthropic::InvokeModelResponse).to receive(:new).and_return(response_obj)

spec/views/components/prompt_form_component_spec.rb

@@ -9,16 +9,23 @@
   let(:setting) { build_stubbed(:setting) }
   let(:conversation) { build_stubbed(:conversation, user:) }
   let(:conversation_form) { ConversationForm.new conversation:, user: }
-  let(:token_count) { 99 }
+  let(:input_token_count) { 99 }
+  let(:output_token_count) { 101 }
   let(:opts) { {} }
 
   before do
-    allow(conversation).to receive(:token_count).and_return token_count
+    allow(conversation).to receive(:input_token_count).and_return input_token_count
+    allow(conversation).to receive(:output_token_count).and_return output_token_count
   end
 
-  it 'shows the token count' do
+  it 'shows the input token count' do
     with_current_user(user) { render_inline component }
-    expect(page).to have_css '.token-display', text: /Tokens:\s+#{token_count}/
+    expect(page).to have_css '#input-tokens', text: /#{input_token_count}/
+  end
+
+  it 'shows the output token count' do
+    with_current_user(user) { render_inline component }
+    expect(page).to have_css '#output-tokens', text: /#{output_token_count}/
   end
 
   describe '#id' do