From 2ccb466598a387b233e2c6566bc3335f7c5dbdd1 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Tue, 21 Jul 2026 12:04:41 +0530 Subject: [PATCH] feat(captain): track document usage for v2 answers --- ...21100000_create_captain_message_sources.rb | 17 +++++++ db/schema.rb | 19 +++++++- .../accounts/captain/documents_controller.rb | 21 ++++++--- .../conversation/response_builder_job.rb | 9 ++-- enterprise/app/models/captain/document.rb | 1 + .../app/models/captain/message_source.rb | 33 ++++++++++++++ .../app/models/enterprise/concerns/account.rb | 1 + .../app/models/enterprise/concerns/message.rb | 1 + .../assistant/session_capture_service.rb | 45 +++++++++++++++++-- .../v1/models/captain/_document.json.jbuilder | 4 ++ .../lib/captain/tools/faq_lookup_tool.rb | 11 ++++- 11 files changed, 147 insertions(+), 15 deletions(-) create mode 100644 db/migrate/20260721100000_create_captain_message_sources.rb create mode 100644 enterprise/app/models/captain/message_source.rb diff --git a/db/migrate/20260721100000_create_captain_message_sources.rb b/db/migrate/20260721100000_create_captain_message_sources.rb new file mode 100644 index 000000000..b9272ff8e --- /dev/null +++ b/db/migrate/20260721100000_create_captain_message_sources.rb @@ -0,0 +1,17 @@ +class CreateCaptainMessageSources < ActiveRecord::Migration[7.1] + def change + create_table :captain_message_sources do |t| + t.references :account, null: false, index: true + t.references :assistant, null: false, index: true + t.references :conversation, null: false, index: true + t.references :message, null: false, index: true + t.references :document, null: false, index: true + t.bigint :assistant_response_id, null: false + + t.timestamps + end + + add_index :captain_message_sources, [:message_id, :assistant_response_id], + unique: true, name: 'idx_captain_message_sources_on_message_and_response' + end +end diff --git a/db/schema.rb b/db/schema.rb index a3dfbda46..da96de484 100644 --- a/db/schema.rb +++ b/db/schema.rb @@ -10,7 +10,7 @@ # # It's strongly recommended that you check this file into your version control system. -ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do +ActiveRecord::Schema[7.1].define(version: 2026_07_21_100000) do # These extensions should be enabled to support this database enable_extension "pg_stat_statements" enable_extension "pg_trgm" @@ -477,6 +477,23 @@ ActiveRecord::Schema[7.1].define(version: 2026_07_13_184351) do t.index ["user_id"], name: "index_captain_message_reports_on_user_id" end + create_table "captain_message_sources", force: :cascade do |t| + t.bigint "account_id", null: false + t.bigint "assistant_id", null: false + t.bigint "conversation_id", null: false + t.bigint "message_id", null: false + t.bigint "document_id", null: false + t.bigint "assistant_response_id", null: false + t.datetime "created_at", null: false + t.datetime "updated_at", null: false + t.index ["account_id"], name: "index_captain_message_sources_on_account_id" + t.index ["assistant_id"], name: "index_captain_message_sources_on_assistant_id" + t.index ["conversation_id"], name: "index_captain_message_sources_on_conversation_id" + t.index ["document_id"], name: "index_captain_message_sources_on_document_id" + t.index ["message_id", "assistant_response_id"], name: "idx_captain_message_sources_on_message_and_response", unique: true + t.index ["message_id"], name: "index_captain_message_sources_on_message_id" + end + create_table "captain_scenarios", force: :cascade do |t| t.string "title" t.text "description" diff --git a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb index cb6fb782d..420dfb6f6 100644 --- a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb +++ b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb @@ -11,7 +11,7 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC @documents = filtered_documents @documents_count = @documents.count @sync_interval_hours = current_sync_interval&.in_hours&.to_i - @documents = with_responses_count(@documents).page(@current_page).per(RESULTS_PER_PAGE) + @documents = with_document_usage(@documents).page(@current_page).per(RESULTS_PER_PAGE) end def show; end @@ -61,14 +61,23 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC apply_sort(documents, permitted_params[:sort]) end - def with_responses_count(scope) - scope.left_joins(:responses) - .select('captain_documents.*, COUNT(captain_assistant_responses.id) AS responses_count') - .group('captain_documents.id') + def with_document_usage(scope) + response_counts = Captain::AssistantResponse.where(documentable_type: 'Captain::Document') + .group(:documentable_id) + .select('documentable_id AS document_id, COUNT(*) AS responses_count') + source_counts = Captain::MessageSource.group(:document_id).select( + 'document_id, COUNT(DISTINCT message_id) AS used_in_answers_count, COUNT(DISTINCT conversation_id) AS used_in_conversations_count' + ) + + scope.joins("LEFT JOIN (#{response_counts.to_sql}) response_counts ON response_counts.document_id = captain_documents.id") + .joins("LEFT JOIN (#{source_counts.to_sql}) source_counts ON source_counts.document_id = captain_documents.id") + .select('captain_documents.*, COALESCE(response_counts.responses_count, 0) AS responses_count, ' \ + 'COALESCE(source_counts.used_in_answers_count, 0) AS used_in_answers_count, ' \ + 'COALESCE(source_counts.used_in_conversations_count, 0) AS used_in_conversations_count') end def set_document - @document = @documents.find(permitted_params[:id]) + @document = with_document_usage(@documents).find(permitted_params[:id]) end def set_assistant diff --git a/enterprise/app/jobs/captain/conversation/response_builder_job.rb b/enterprise/app/jobs/captain/conversation/response_builder_job.rb index fb0e72721..6e44070bc 100644 --- a/enterprise/app/jobs/captain/conversation/response_builder_job.rb +++ b/enterprise/app/jobs/captain/conversation/response_builder_job.rb @@ -68,7 +68,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob # left is the customer-facing follow-up message. process_v2_handoff end - capture_assistant_session(result_message: @handoff_message, credits_consumed: 0.0) + capture_assistant_session(result_message: @handoff_message, credits_consumed: 0.0, capture_message_sources: false) elsif v1_handoff_requested? # V1 only signals via the response string — no state has been touched yet. If # the conversation isn't pending anymore, a human took over mid-run; bail out @@ -83,7 +83,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob Rails.logger.info("[CAPTAIN][ResponseBuilderJob] Incrementing response usage for #{account.id}") account.increment_response_usage end - capture_assistant_session(result_message: message, credits_consumed: 1.0) + capture_assistant_session(result_message: message, credits_consumed: 1.0, capture_message_sources: captain_v2_enabled?) end end @@ -142,9 +142,10 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob # Capture runs outside the delivery transaction and never raises (the service # swallows its own failures): a session-logging bug must never roll back the # customer reply or trigger the top-level handle_error handoff on top of it. - def capture_assistant_session(result_message:, credits_consumed:) + def capture_assistant_session(result_message:, credits_consumed:, capture_message_sources:) Captain::Assistant::SessionCaptureService.new(assistant: @assistant, conversation: @conversation, run_result: @run_result, - result_message: result_message, credits_consumed: credits_consumed).capture + result_message: result_message, credits_consumed: credits_consumed, + capture_message_sources: capture_message_sources).capture end def handle_error(error) diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb index 6eda0eb5a..a3f8d176e 100644 --- a/enterprise/app/models/captain/document.rb +++ b/enterprise/app/models/captain/document.rb @@ -34,6 +34,7 @@ class Captain::Document < ApplicationRecord belongs_to :assistant, class_name: 'Captain::Assistant' has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy, as: :documentable + has_many :message_sources, class_name: 'Captain::MessageSource', dependent: :destroy_async belongs_to :account has_one_attached :pdf_file store_accessor :metadata, :content_fingerprint, :last_sync_error_code, :sync_step, :openai_file_id diff --git a/enterprise/app/models/captain/message_source.rb b/enterprise/app/models/captain/message_source.rb new file mode 100644 index 000000000..e29c2d8f0 --- /dev/null +++ b/enterprise/app/models/captain/message_source.rb @@ -0,0 +1,33 @@ +# == Schema Information +# +# Table name: captain_message_sources +# +# id :bigint not null, primary key +# created_at :datetime not null +# updated_at :datetime not null +# account_id :bigint not null +# assistant_id :bigint not null +# assistant_response_id :bigint not null +# conversation_id :bigint not null +# document_id :bigint not null +# message_id :bigint not null +# +# Indexes +# +# idx_captain_message_sources_on_message_and_response (message_id,assistant_response_id) UNIQUE +# index_captain_message_sources_on_account_id (account_id) +# index_captain_message_sources_on_assistant_id (assistant_id) +# index_captain_message_sources_on_conversation_id (conversation_id) +# index_captain_message_sources_on_document_id (document_id) +# index_captain_message_sources_on_message_id (message_id) +# +class Captain::MessageSource < ApplicationRecord + self.table_name = 'captain_message_sources' + + belongs_to :account + belongs_to :assistant, class_name: 'Captain::Assistant' + belongs_to :conversation, class_name: '::Conversation' + belongs_to :message + belongs_to :document, class_name: 'Captain::Document' + belongs_to :assistant_response, class_name: 'Captain::AssistantResponse', optional: true +end diff --git a/enterprise/app/models/enterprise/concerns/account.rb b/enterprise/app/models/enterprise/concerns/account.rb index 427b1e1af..30ebfca16 100644 --- a/enterprise/app/models/enterprise/concerns/account.rb +++ b/enterprise/app/models/enterprise/concerns/account.rb @@ -16,6 +16,7 @@ module Enterprise::Concerns::Account has_many :captain_documents, dependent: :destroy_async, class_name: 'Captain::Document' has_many :captain_custom_tools, dependent: :destroy_async, class_name: 'Captain::CustomTool' has_many :captain_agent_sessions, dependent: :destroy_async, class_name: 'Captain::AgentSession' + has_many :captain_message_sources, dependent: :destroy_async, class_name: 'Captain::MessageSource' has_many :copilot_threads, dependent: :destroy_async has_many :companies, dependent: :destroy_async diff --git a/enterprise/app/models/enterprise/concerns/message.rb b/enterprise/app/models/enterprise/concerns/message.rb index 3c7cdc3f7..3f2e27c98 100644 --- a/enterprise/app/models/enterprise/concerns/message.rb +++ b/enterprise/app/models/enterprise/concerns/message.rb @@ -4,5 +4,6 @@ module Enterprise::Concerns::Message included do has_one :call, dependent: :nullify has_many :message_reports, class_name: 'Captain::MessageReport', dependent: :destroy_async + has_many :captain_message_sources, class_name: 'Captain::MessageSource', dependent: :destroy_async end end diff --git a/enterprise/app/services/captain/assistant/session_capture_service.rb b/enterprise/app/services/captain/assistant/session_capture_service.rb index 36af50308..b825dbe55 100644 --- a/enterprise/app/services/captain/assistant/session_capture_service.rb +++ b/enterprise/app/services/captain/assistant/session_capture_service.rb @@ -1,12 +1,13 @@ class Captain::Assistant::SessionCaptureService SCENARIO_AGENT_REGEX = /\A#{Captain::Scenario::HANDOFF_KEY_PREFIX}_(\d+)_/ - def initialize(assistant:, conversation:, run_result:, result_message:, credits_consumed:) + def initialize(assistant:, conversation:, run_result:, result_message:, **options) @assistant = assistant @conversation = conversation @run_result = run_result @result_message = result_message - @credits_consumed = credits_consumed + @credits_consumed = options.fetch(:credits_consumed) + @capture_message_sources = options.fetch(:capture_message_sources, false) end def capture @@ -24,7 +25,7 @@ class Captain::Assistant::SessionCaptureService model = @assistant.agent_model metadata = context.dig(:state, :cw_metadata) || {} - Captain::AgentSession.create!( + session = Captain::AgentSession.create!( assistant: @assistant, session_type: :assistant, subject: @conversation, @@ -36,6 +37,8 @@ class Captain::Assistant::SessionCaptureService scenario_ids: scenario_ids, run_context: current_turn_history ) + capture_message_sources(metadata) if @capture_message_sources + session end private @@ -54,6 +57,42 @@ class Captain::Assistant::SessionCaptureService ids & @assistant.scenarios.where(id: ids).pluck(:id) end + def capture_message_sources(metadata) + sources = Array(metadata[:message_sources]) + return if sources.empty? + + document_ids = @assistant.documents.where(id: sources.pluck(:document_id)).pluck(:id) + return if document_ids.empty? + + insert_message_sources(message_source_rows(sources, document_ids)) + end + + def message_source_rows(sources, document_ids) + timestamp = Time.current + sources.filter_map do |source| + next unless document_ids.include?(source[:document_id]) + + { + account_id: @assistant.account_id, + assistant_id: @assistant.id, + conversation_id: @conversation.id, + message_id: @result_message.id, + document_id: source[:document_id], + assistant_response_id: source[:assistant_response_id], + created_at: timestamp, + updated_at: timestamp + } + end + end + + def insert_message_sources(rows) + return if rows.empty? + + # rubocop:disable Rails/SkipsModelValidations + Captain::MessageSource.insert_all(rows, unique_by: :idx_captain_message_sources_on_message_and_response) + # rubocop:enable Rails/SkipsModelValidations + end + # Trim to the current turn: the last user message and everything after it # (assistant replies, tool calls/results, handoff hops). def current_turn_history diff --git a/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder b/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder index 0ab031dbf..5090bf1a9 100644 --- a/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder +++ b/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder @@ -11,6 +11,10 @@ json.file_size resource.file_size json.pdf_document resource.pdf_document? responses_count = resource.respond_to?(:responses_count) ? resource.responses_count : resource.responses.count json.responses_count responses_count.to_i +used_in_answers_count = resource.respond_to?(:used_in_answers_count) ? resource.used_in_answers_count : resource.message_sources.distinct.count(:message_id) +json.used_in_answers_count used_in_answers_count.to_i +used_in_conversations_count = resource.respond_to?(:used_in_conversations_count) ? resource.used_in_conversations_count : resource.message_sources.distinct.count(:conversation_id) +json.used_in_conversations_count used_in_conversations_count.to_i json.id resource.id json.name resource.name json.status resource.status diff --git a/enterprise/lib/captain/tools/faq_lookup_tool.rb b/enterprise/lib/captain/tools/faq_lookup_tool.rb index 2a16e7e99..bb8c0af95 100644 --- a/enterprise/lib/captain/tools/faq_lookup_tool.rb +++ b/enterprise/lib/captain/tools/faq_lookup_tool.rb @@ -26,8 +26,17 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool metadata = tool_context.state[:cw_metadata] ||= {} metadata[:faq_ids] = Array(metadata[:faq_ids]) | responses.map(&:id) - document_ids = responses.filter_map { |response| response.documentable_id if response.documentable_type == 'Captain::Document' } + document_ids = document_responses(responses).map(&:documentable_id) metadata[:document_ids] = Array(metadata[:document_ids]) | document_ids + metadata[:message_sources] = Array(metadata[:message_sources]) | message_sources(document_responses(responses)) + end + + def document_responses(responses) + responses.select { |response| response.documentable_type == 'Captain::Document' } + end + + def message_sources(responses) + responses.map { |response| { assistant_response_id: response.id, document_id: response.documentable_id } } end def format_responses(responses)