Merge pull request #602 from alphagov/2831-add-link-token-mapping-to-llm-response

Add link token mapping to LLM response

Author: davidgisbey

lib/answer_composition/link_token_mapper.rb

@@ -1,6 +1,7 @@
 module AnswerComposition
   class LinkTokenMapper
     TOKEN_PREFIX = "link_".freeze
+    attr_reader :mapping
 
     def initialize
       @mapping = {}
@@ -50,8 +51,6 @@ def replace_tokens_with_links(markdown)
 
   private
 
-    attr_reader :mapping
-
     def rewrite_links(element)
       if element.type == :a
         rewrite_link(element)

lib/answer_composition/pipeline/claude/structured_answer_composer.rb

@@ -21,8 +21,12 @@ def call
           tool_choice: { type: "tool", name: "output_schema" },
           **inference_config,
         )
+        llm_response_with_link_token_mapping = {
+          "response" => response.to_h.stringify_keys,
+          "link_token_mapping" => link_token_mapper.mapping.invert,
+        }
 
-        context.answer.assign_llm_response("structured_answer", response.to_h)
+        context.answer.assign_llm_response("structured_answer", llm_response_with_link_token_mapping)
         tool_output = response[:content][0][:input]
 
         unless tool_output[:answered]

lib/answer_composition/pipeline/openai/structured_answer_composer.rb

@@ -11,7 +11,11 @@ def initialize(context)
 
     def call
       start_time = Clock.monotonic_time
-      context.answer.assign_llm_response("structured_answer", openai_response_choice)
+      llm_response_with_link_token_mapping = {
+        "response" => openai_response_choice,
+        "link_token_mapping" => link_token_mapper.mapping.invert,
+      }
+      context.answer.assign_llm_response("structured_answer", llm_response_with_link_token_mapping)
 
       unless parsed_structured_response["answered"]
         return context.abort_pipeline!(

spec/lib/answer_composition/link_token_mapper_spec.rb

@@ -17,6 +17,13 @@
     HTML
   end
 
+  describe "#initialize" do
+    it "initializes mapping as an empty hash" do
+      mapper = described_class.new
+      expect(mapper.mapping).to eq({})
+    end
+  end
+
   describe "#map_links_to_tokens" do
     it "replaces href attributes with tokens" do
       amended_html = described_class.new.map_links_to_tokens(html, "/exact-path")
@@ -48,6 +55,20 @@
 
       expect(result).to eq(html)
     end
+
+    it "updates the mapping attribute with the link mappings" do
+      mapper = described_class.new
+      mapper.map_links_to_tokens(html, "/exact-path")
+
+      expect(mapper.mapping).to eq(
+        {
+          "https://www.test.gov.uk/tax-returns" => "link_1",
+          "https://www.test.gov.uk/tax-news" => "link_2",
+          "https://www.gov.uk/national-insurance/what-national-insurance-is" => "link_3",
+          "https://www.test.gov.uk/exact-path#some-heading" => "link_4",
+        },
+      )
+    end
   end
 
   describe "#map_link_to_token" do

spec/lib/answer_composition/pipeline/claude/structured_answer_composer_spec.rb

@@ -43,11 +43,17 @@
         input: { answer:, answered: true, sources_used: %w[link_1], answer_completeness: "complete" },
         name: "output_schema",
       )
-      expected_llm_response = claude_messages_response(
-        content: [expected_content],
-        usage: { cache_read_input_tokens: 20 },
-        stop_reason: :tool_use,
-      ).to_h
+      expected_llm_response = {
+        "response" => claude_messages_response(
+          content: [expected_content],
+          usage: { cache_read_input_tokens: 20 },
+          stop_reason: :tool_use,
+        ).to_h.stringify_keys,
+        "link_token_mapping" => {
+          "link_1" => "https://www.test.gov.uk/vat-rates#vat-basics",
+          "link_2" => "https://www.test.gov.uk/what-is-tax",
+        },
+      }
       expect(context.answer.llm_responses["structured_answer"]).to eq(expected_llm_response)
     end
 
@@ -139,11 +145,17 @@
           },
           name: "output_schema",
         )
-        expected_llm_response = claude_messages_response(
-          content: [expected_content],
-          usage: { cache_read_input_tokens: 20 },
-          stop_reason: :tool_use,
-        ).to_h
+        expected_llm_response = {
+          "response" => claude_messages_response(
+            content: [expected_content],
+            usage: { cache_read_input_tokens: 20 },
+            stop_reason: :tool_use,
+          ).to_h.stringify_keys,
+          "link_token_mapping" => {
+            "link_1" => "https://www.test.gov.uk/vat-rates#vat-basics",
+            "link_2" => "https://www.test.gov.uk/what-is-tax",
+          },
+        }
         expect(context.answer.llm_responses["structured_answer"]).to eq(expected_llm_response)
       end
     end

spec/lib/answer_composition/pipeline/openai/structured_answer_composer_spec.rb

@@ -59,13 +59,18 @@
 
         described_class.call(context)
 
+        expected_structured_answer = hash_including(
+          "response" => hash_including_openai_response_with_tool_call("generate_answer_using_retrieved_contexts"),
+          "link_token_mapping" => {
+            "link_1" => "https://www.test.gov.uk/vat-rates#vat-basics",
+            "link_2" => "https://www.test.gov.uk/what-is-tax",
+          },
+        )
+        expect(context.answer.llm_responses["structured_answer"]).to match(expected_structured_answer)
         expect(context.answer.message.squish).to eq(
           "VAT (Value Added Tax) is a [tax][1] applied to most goods and services in the UK. [1]: https://www.test.gov.uk/what-is-tax",
         )
         expect(context.answer.status).to eq("answered")
-        expect(context.answer.llm_responses["structured_answer"]).to match(
-          hash_including_openai_response_with_tool_call("generate_answer_using_retrieved_contexts"),
-        )
       end
 
       it "sets the 'used' boolean to false for unused sources" do