fix(captain): resolve v2 FAQ citations from documents

2026-07-21 12:10:40 +05:30
parent 2ccb466598
commit 02f288520e
5 changed files with 89 additions and 16 deletions
@@ -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)
@@ -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,
@@ -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' %}
@@ -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
@@ -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