Merge pull request #1020 from alphagov/472-combine-multiple-checker

[CHAT-472] Merge multiple checker into single class

Author: davidgisbey

lib/answer_composition/composer.rb

@@ -40,10 +40,10 @@ def compose_answer
           Pipeline::JailbreakGuardrails,
           Pipeline::QuestionRephraser,
           Pipeline::QuestionRouter,
-          Pipeline::QuestionRoutingGuardrails.new(llm_provider: :claude),
+          Pipeline::QuestionRoutingGuardrails,
           Pipeline::SearchResultFetcher,
           Pipeline::StructuredAnswerComposer,
-          Pipeline::AnswerGuardrails.new(llm_provider: :claude),
+          Pipeline::AnswerGuardrails,
         ])
       else
         raise "Answer strategy #{answer_strategy} not configured"

lib/answer_composition/pipeline/answer_guardrails.rb

@@ -1,7 +1,7 @@
 module AnswerComposition
   module Pipeline
     class AnswerGuardrails < OutputGuardrails
-      def call(context)
+      def call
         start_time = Clock.monotonic_time
         response = generate_response(context)
 

lib/answer_composition/pipeline/output_guardrails.rb

@@ -1,14 +1,16 @@
 module AnswerComposition
   module Pipeline
     class OutputGuardrails
-      def initialize(llm_provider: :claude)
-        @llm_provider = llm_provider
+      attr_reader :context
+
+      def self.call(...) = new(...).call
+
+      def initialize(context)
+        @context = context
       end
 
     protected
 
-      attr_reader :llm_provider
-
       def build_metrics(start_time, response_or_error)
         {
           duration: Clock.monotonic_time - start_time,
@@ -20,7 +22,7 @@ def build_metrics(start_time, response_or_error)
       end
 
       def generate_response(context)
-        result = ::Guardrails::MultipleChecker.call(context.answer.message, guardrail_name, llm_provider)
+        result = ::Guardrails::MultipleChecker.call(context.answer.message, guardrail_name)
         context.answer.assign_llm_response(guardrail_name, result.llm_response)
         result
       end

lib/answer_composition/pipeline/question_routing_guardrails.rb

@@ -1,7 +1,7 @@
 module AnswerComposition
   module Pipeline
     class QuestionRoutingGuardrails < OutputGuardrails
-      def call(context)
+      def call
         return if context.answer.question_routing_label == "genuine_rag"
 
         start_time = Clock.monotonic_time

lib/guardrails/claude/multiple_checker.rb

@@ -9,6 +9,10 @@ def triggered_guardrails
       end
     end
 
+    MAX_TOKENS = 100
+    SUPPORTED_MODELS = %i[claude_sonnet_4_0 claude_haiku_4_5].freeze
+    DEFAULT_MODEL = :claude_sonnet_4_0
+
     class ResponseError < StandardError
       attr_reader :llm_response, :llm_guardrail_result, :llm_prompt_tokens,
                   :llm_completion_tokens, :llm_cached_tokens, :model
@@ -35,12 +39,10 @@ class Prompt
 
       Guardrail = Data.define(:key, :name, :content)
 
-      def initialize(prompt_name, llm_provider = :claude)
-        prompts = if llm_provider == :claude
-                    AnswerComposition::Pipeline::Prompts.config(prompt_name, Claude::MultipleChecker.bedrock_model)
-                  else
-                    Rails.configuration.govuk_chat_private.llm_prompts[llm_provider][prompt_name]
-                  end
+      def initialize(prompt_name)
+        prompts = AnswerComposition::Pipeline::Prompts.config(
+          prompt_name, Guardrails::MultipleChecker.bedrock_model
+        )
 
         raise "No LLM prompts found for #{prompt_name}" unless prompts
 
@@ -72,8 +74,12 @@ def guardrails
 
     def self.call(...) = new(...).call
 
-    def self.collated_prompts(llm_prompt_name, llm_provider)
-      prompt = Prompt.new(llm_prompt_name, llm_provider)
+    def self.bedrock_model
+      BedrockModels.determine_model(ENV["BEDROCK_CLAUDE_GUARDRAILS_MODEL"], DEFAULT_MODEL, SUPPORTED_MODELS).last
+    end
+
+    def self.collated_prompts(llm_prompt_name)
+      prompt = Prompt.new(llm_prompt_name)
 
       <<~PROMPT
         # System prompt
@@ -85,38 +91,46 @@ def self.collated_prompts(llm_prompt_name, llm_provider)
       PROMPT
     end
 
-    def initialize(input, llm_prompt_name, llm_provider)
+    def initialize(input, llm_prompt_name)
       @input = input
       @llm_prompt_name = llm_prompt_name
-      @llm_provider = llm_provider
     end
 
     def call
-      case llm_provider
-      when :claude
-        response = Claude::MultipleChecker.call(input, prompt)
-      else
-        raise "Unexpected provider #{llm_provider}"
-      end
-      parse_response(**response)
+      response = anthropic_bedrock_client.messages.create(
+        system: [{ type: "text", text: prompt.system_prompt, cache_control: { type: "ephemeral" } }],
+        model: BedrockModels.model_id(self.class.bedrock_model),
+        messages: [{ role: "user", content: prompt.user_prompt(input) }],
+        max_tokens: MAX_TOKENS,
+      )
+
+      parse_response(response)
     end
 
   private
 
-    def parse_response(llm_response:,
-                       llm_guardrail_result:,
-                       llm_prompt_tokens:,
-                       llm_completion_tokens:,
-                       llm_cached_tokens:,
-                       model:)
+    def anthropic_bedrock_client
+      @anthropic_bedrock_client ||= Anthropic::BedrockClient.new(
+        aws_region: ENV["CLAUDE_AWS_REGION"],
+      )
+    end
+
+    def parse_response(response)
+      llm_response = response.to_h
+      llm_guardrail_result = response[:content][0][:text]
+      input_tokens = response[:usage][:input_tokens]
+      output_tokens = response[:usage][:output_tokens]
+      cache_read_input_tokens = response[:usage][:cache_read_input_tokens]
+      model = response[:model]
+
       unless response_pattern =~ llm_guardrail_result
         raise ResponseError.new(
           "Error parsing guardrail response",
           llm_response,
           llm_guardrail_result,
-          llm_prompt_tokens,
-          llm_completion_tokens,
-          llm_cached_tokens,
+          input_tokens,
+          output_tokens,
+          cache_read_input_tokens,
           model,
         )
       end
@@ -126,19 +140,19 @@ def parse_response(llm_response:,
       guardrails = to_guardrail_hash(parts.second)
 
       Result.new(
-        llm_response: llm_response,
-        llm_guardrail_result: llm_guardrail_result,
-        triggered: triggered,
-        guardrails: guardrails,
-        llm_prompt_tokens: llm_prompt_tokens,
-        llm_completion_tokens: llm_completion_tokens,
-        llm_cached_tokens: llm_cached_tokens,
+        llm_response:,
+        llm_guardrail_result:,
+        triggered:,
+        guardrails:,
+        llm_prompt_tokens: input_tokens,
+        llm_completion_tokens: output_tokens,
+        llm_cached_tokens: cache_read_input_tokens,
         model:,
       )
     end
 
     def prompt
-      @prompt ||= Prompt.new(llm_prompt_name, llm_provider)
+      @prompt ||= Prompt.new(llm_prompt_name)
     end
 
     def guardrail_numbers

lib/guardrails/multiple_checker.rb

@@ -32,7 +32,7 @@ namespace :evaluation do
     raise "Requires an INPUT env var" if ENV["INPUT"].blank?
     raise "Requires a guardrail type" if args[:guardrail_type].blank?
 
-    response = Guardrails::MultipleChecker.call(ENV["INPUT"], args[:guardrail_type].to_sym, :claude)
+    response = Guardrails::MultipleChecker.call(ENV["INPUT"], args[:guardrail_type].to_sym)
 
     puts(response.to_json)
   end

lib/tasks/evaluation.rake

@@ -7,7 +7,7 @@ namespace "guardrails" do
       abort("Invalid guardrail type. Valid guardrail types are #{valid_guardrail_types.to_sentence}")
     end
 
-    prompt = Guardrails::MultipleChecker.collated_prompts(guardrail_type, :claude)
+    prompt = Guardrails::MultipleChecker.collated_prompts(guardrail_type)
     puts prompt
   end
 end

lib/tasks/guardrails.rake

@@ -5,7 +5,7 @@
     llm_prompt_tokens { 13 }
     llm_completion_tokens { 7 }
     llm_cached_tokens { 10 }
-    model { BedrockModels.model_id(Guardrails::Claude::MultipleChecker::DEFAULT_MODEL) }
+    model { BedrockModels.model_id(Guardrails::MultipleChecker::DEFAULT_MODEL) }
 
     llm_response do
       content = Anthropic::Models::TextBlock.new(

spec/factories/output_guardrail_result_factory.rb

@@ -24,17 +24,14 @@ def stub_pipeline_initialize(klass, *args, **kwargs)
       let(:question) { create :question, answer_strategy: :claude_structured_answer }
 
       it "calls PipelineRunner with the correct pipeline" do
-        stub_pipeline_initialize(AnswerComposition::Pipeline::QuestionRoutingGuardrails, llm_provider: :claude)
-        stub_pipeline_initialize(AnswerComposition::Pipeline::AnswerGuardrails, llm_provider: :claude)
-
         expected_pipeline = [
           AnswerComposition::Pipeline::JailbreakGuardrails,
           AnswerComposition::Pipeline::QuestionRephraser,
           AnswerComposition::Pipeline::QuestionRouter,
-          AnswerComposition::Pipeline::QuestionRoutingGuardrails.new(llm_provider: :claude),
+          AnswerComposition::Pipeline::QuestionRoutingGuardrails,
           AnswerComposition::Pipeline::SearchResultFetcher,
           AnswerComposition::Pipeline::StructuredAnswerComposer,
-          AnswerComposition::Pipeline::AnswerGuardrails.new(llm_provider: :claude),
+          AnswerComposition::Pipeline::AnswerGuardrails,
         ]
         expected_pipeline.each do |pipeline|
           allow(pipeline).to receive(:call) { it }