From 02f288520ee9738eee2d16b07fbc537a821d1957 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Tue, 21 Jul 2026 12:10:40 +0530 Subject: [PATCH] fix(captain): resolve v2 FAQ citations from documents --- .../captain/conversation/message_builder.rb | 67 ++++++++++++++++++- enterprise/app/models/captain/assistant.rb | 1 + .../lib/captain/prompts/assistant.liquid | 5 ++ .../lib/captain/tools/faq_lookup_tool.rb | 27 ++++---- .../lib/captain/tools/faq_lookup_tool_spec.rb | 5 +- 5 files changed, 89 insertions(+), 16 deletions(-) diff --git a/enterprise/app/jobs/captain/conversation/message_builder.rb b/enterprise/app/jobs/captain/conversation/message_builder.rb index 7954e9465..eca692fb6 100644 --- a/enterprise/app/jobs/captain/conversation/message_builder.rb +++ b/enterprise/app/jobs/captain/conversation/message_builder.rb @@ -1,4 +1,8 @@ module Captain::Conversation::MessageBuilder + SOURCE_MARKER_PATTERN = /\[\[source:(\d+)\]\]/ + MALFORMED_SOURCE_MARKER_PATTERN = /\[\[source:[^\]\n]*\]\]?/ + CUSTOMER_CITATION_PATTERN = /\[\[(\d+)\]\([^\)\n]*\)\]/ + private def collect_previous_messages @@ -28,8 +32,67 @@ module Captain::Conversation::MessageBuilder end def create_messages - validate_message_content!(@response['response']) - create_outgoing_message(@response['response'], agent_name: @response['agent_name']) + message_content = resolve_citations(@response['response']) + validate_message_content!(message_content) + create_outgoing_message(message_content, agent_name: @response['agent_name']) + end + + def resolve_citations(content) + return content if content.blank? + + citation_urls = trusted_citation_urls + content = content.gsub(CUSTOMER_CITATION_PATTERN) do + citation_for(Regexp.last_match(1), citation_urls) + end + content = content.gsub(SOURCE_MARKER_PATTERN) do + citation_for(Regexp.last_match(1), citation_urls) + end + content.gsub(MALFORMED_SOURCE_MARKER_PATTERN, '') + end + + def citation_for(source_number, citation_urls) + url = citation_urls[source_number.to_s] + return '' if url.blank? + + "[[#{source_number}](#{url})]" + end + + def trusted_citation_urls + return {} unless @assistant.config['feature_citation'] + + source_map = citation_source_map + return {} if source_map.empty? + + document_links = Captain::Document.where( + id: source_map.values, + account_id: account.id, + assistant_id: @assistant.id + ).pluck(:id, :external_link).to_h + + source_map.each_with_object({}) do |(source_number, document_id), urls| + external_link = document_links[document_id] + urls[source_number] = external_link if customer_safe_document_link?(external_link) + end + end + + def citation_source_map + raw_citation_source_map.filter_map do |source_number, document_id| + [source_number.to_s, document_id.to_i] if valid_citation_source?(source_number, document_id) + end.to_h + end + + def raw_citation_source_map + context = @run_result&.context || {} + state = context[:state] || context['state'] || {} + state[:captain_v2_citation_source_map] || state['captain_v2_citation_source_map'] || {} + end + + def valid_citation_source?(source_number, document_id) + source_number.to_s.match?(/\A[1-9]\d*\z/) && document_id.to_s.match?(/\A\d+\z/) + end + + def customer_safe_document_link?(external_link) + external_link.present? && !external_link.start_with?('PDF:') end def validate_message_content!(content) diff --git a/enterprise/app/models/captain/assistant.rb b/enterprise/app/models/captain/assistant.rb index dc0969cd4..f7cf22bd4 100644 --- a/enterprise/app/models/captain/assistant.rb +++ b/enterprise/app/models/captain/assistant.rb @@ -108,6 +108,7 @@ class Captain::Assistant < ApplicationRecord name: name, description: description, product_name: config['product_name'] || 'this product', + citation_enabled: config['feature_citation'], scenarios: scenarios.enabled.map do |scenario| { title: scenario.title, diff --git a/enterprise/lib/captain/prompts/assistant.liquid b/enterprise/lib/captain/prompts/assistant.liquid index a8f1dada3..78b4de6bd 100644 --- a/enterprise/lib/captain/prompts/assistant.liquid +++ b/enterprise/lib/captain/prompts/assistant.liquid @@ -10,6 +10,11 @@ You are {{name}}, a helpful, friendly, and knowledgeable assistant for the produ Don't digress away from your instructions, and use all the available tools at your disposal for solving customer issues. If you are to state something factual about {{product_name}}, use the `captain--tools--faq_lookup` tool to check the available information first. +{% if citation_enabled -%} +# Citations +When you use FAQ content that includes a source marker, add that exact marker at the end of the supported sentence or paragraph. Source markers have the form `[[source:n]]`. Do not create URLs or markdown links. Reuse the same source marker when the same document supports more than one statement. Do not cite conversation context. +{% endif -%} + {% render 'current_time', current_time: current_time %} {% render 'core_rules' %} diff --git a/enterprise/lib/captain/tools/faq_lookup_tool.rb b/enterprise/lib/captain/tools/faq_lookup_tool.rb index bb8c0af95..4a081b431 100644 --- a/enterprise/lib/captain/tools/faq_lookup_tool.rb +++ b/enterprise/lib/captain/tools/faq_lookup_tool.rb @@ -14,7 +14,7 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool "No relevant FAQs found for: #{query}" else log_tool_usage('found_results', { query: query, count: responses.size }) - format_responses(responses) + format_responses(tool_context, responses) end end @@ -29,6 +29,8 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool 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)) + + document_ids.each { |document_id| citation_source_number(tool_context, document_id) } end def document_responses(responses) @@ -39,30 +41,31 @@ class Captain::Tools::FaqLookupTool < Captain::Tools::BasePublicTool responses.map { |response| { assistant_response_id: response.id, document_id: response.documentable_id } } end - def format_responses(responses) - responses.map { |response| format_response(response) }.join + def format_responses(tool_context, responses) + responses.map { |response| format_response(tool_context, response) }.join end - def format_response(response) + def format_response(tool_context, response) formatted_response = " Question: #{response.question} Answer: #{response.answer} " - if should_show_source?(response) + if response.documentable_type == 'Captain::Document' formatted_response += " - Source: #{response.documentable.external_link} + Source marker: [[source:#{citation_source_number(tool_context, response.documentable_id)}]] " end formatted_response end - def should_show_source?(response) - return false if response.documentable.blank? - return false unless response.documentable.try(:external_link) + def citation_source_number(tool_context, document_id) + source_map = tool_context.state[:captain_v2_citation_source_map] ||= {} + existing_number = source_map.find { |_number, id| id.to_i == document_id }&.first + return existing_number if existing_number.present? - # Don't show source if it's a PDF placeholder - external_link = response.documentable.external_link - !external_link.start_with?('PDF:') + source_number = (source_map.keys.map(&:to_i).max || 0) + 1 + source_map[source_number.to_s] = document_id + source_number.to_s end end diff --git a/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb b/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb index dc2b9f6c6..bd76ef58e 100644 --- a/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb +++ b/spec/enterprise/lib/captain/tools/faq_lookup_tool_spec.rb @@ -66,12 +66,13 @@ RSpec.describe Captain::Tools::FaqLookupTool, type: :model do expect(result).to include('Answer: Go to settings and update email') end - it 'includes source link when document has external_link' do + it 'includes a source marker instead of the document link' do document.update!(external_link: 'https://help.example.com/password') result = tool.perform(tool_context, query: 'password') - expect(result).to include('Source: https://help.example.com/password') + expect(result).to include('Source marker: [[source:1]]') + expect(result).not_to include('https://help.example.com/password') end it 'logs tool usage for search' do