diff --git a/.rubocop.yml b/.rubocop.yml index bd3e51e0f..40003348c 100644 --- a/.rubocop.yml +++ b/.rubocop.yml @@ -49,6 +49,7 @@ Metrics/ClassLength: - 'method_call' Exclude: - 'app/controllers/**/*' + - 'app/models/**/*' - 'lib/**/*_form.rb' Metrics/ModuleLength: CountAsOne: diff --git a/app/models/conversation.rb b/app/models/conversation.rb index edc1f0505..0c2a85fe9 100644 --- a/app/models/conversation.rb +++ b/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, diff --git a/app/models/generate_text_request.rb b/app/models/generate_text_request.rb index 59daf6c58..e445b2e19 100644 --- a/app/models/generate_text_request.rb +++ b/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 diff --git a/app/views/components/prompt_form_component/prompt_form_component.html.haml b/app/views/components/prompt_form_component/prompt_form_component.html.haml index 50bc30229..0a676db85 100644 --- a/app/views/components/prompt_form_component/prompt_form_component.html.haml +++ b/app/views/components/prompt_form_component/prompt_form_component.html.haml @@ -76,9 +76,13 @@ .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' }} @@ -86,6 +90,7 @@ %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' } diff --git a/lib/anthropic/invoke_model_response.rb b/lib/anthropic/invoke_model_response.rb index ff8725561..d8c1cb332 100644 --- a/lib/anthropic/invoke_model_response.rb +++ b/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 diff --git a/lib/generative_text/aws/invoke_model_response.rb b/lib/generative_text/aws/invoke_model_response.rb index 9c25c406b..a35ba8a79 100644 --- a/lib/generative_text/aws/invoke_model_response.rb +++ b/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 diff --git a/spec/models/generate_text_request_spec.rb b/spec/models/generate_text_request_spec.rb index 257f1ec62..767d3dbd1 100644 --- a/spec/models/generate_text_request_spec.rb +++ b/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) diff --git a/spec/views/components/prompt_form_component_spec.rb b/spec/views/components/prompt_form_component_spec.rb index 90d872c16..064583f5e 100644 --- a/spec/views/components/prompt_form_component_spec.rb +++ b/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