Move jailbreak guardrails into single pipeline class

We've got 3 classes for Jailbreak guardrails. It's a fairly simple
class and can just be moved into a single class in the answer composition
pipeline.

I did consider leaving the JailbreakChecker class as a separate class, but
it just didn't seem like there was a lot of value.

One thing i noticed while merging this, is that we create a ResponseError
class in the JailbreakChecker, but we don't actually raise it anywhere
so it was redundant. I've removed it, but if we want we can check
that the llm response either returns the pass or fail value and
create the object in the db with the correct status if it doesn't.

We'd need to add some additional prompt config for the fail value though
in order to do that so i've avoided it for now.

We will need to do some follow up work to update the evaluation repo
since we just return the serialised answer now so it'll need to grab
the info it wants from the serialised answer instead of the result that
was previously returned from the JailbreakChecker.

Author: davidgisbey

lib/answer_composition/composer.rb

@@ -37,7 +37,7 @@ def compose_answer
       case answer_strategy
       when "claude_structured_answer"
         PipelineRunner.call(question:, pipeline: [
-          Pipeline::JailbreakGuardrails.new(llm_provider: :claude),
+          Pipeline::JailbreakGuardrails,
           Pipeline::QuestionRephraser,
           Pipeline::QuestionRouter,
           Pipeline::QuestionRoutingGuardrails.new(llm_provider: :claude),

lib/answer_composition/pipeline/jailbreak_guardrails.rb

@@ -1,45 +1,96 @@
 module AnswerComposition
   module Pipeline
     class JailbreakGuardrails
-      def initialize(llm_provider: :claude)
-        @llm_provider = llm_provider
+      SUPPORTED_MODELS = %i[claude_sonnet_4_0 claude_haiku_4_5].freeze
+      DEFAULT_MODEL = :claude_sonnet_4_0
+
+      def self.call(...) = new(...).call
+
+      def initialize(context)
+        @context = context
+        @model_id, @model_name = BedrockModels.determine_model(
+          ENV["BEDROCK_CLAUDE_JAILBREAK_GUARDRAILS_MODEL"],
+          DEFAULT_MODEL,
+          SUPPORTED_MODELS,
+        )
       end
 
-      def call(context)
+      def call
         start_time = Clock.monotonic_time
+        response = anthropic_bedrock_client.messages.create(
+          system: [{ type: "text", text: system_prompt }],
+          model: model_id,
+          messages:,
+          **inference_config,
+        )
 
-        response = Guardrails::JailbreakChecker.call(context.question.message, llm_provider)
-        context.answer.assign_attributes(jailbreak_guardrails_status: response.triggered ? :fail : :pass)
-        context.answer.assign_llm_response("jailbreak_guardrails", response.llm_response)
+        jailbreak_guardrails_status = response[:content][0][:text] == pass_value ? :pass : :fail
+
+        context.answer.assign_attributes(jailbreak_guardrails_status:)
+        context.answer.assign_llm_response("jailbreak_guardrails", response.to_h)
         context.answer.assign_metrics("jailbreak_guardrails", build_metrics(start_time, response))
 
-        if response.triggered
+        if jailbreak_guardrails_status == :fail
           context.abort_pipeline!(
             message: Answer::CannedResponses::JAILBREAK_GUARDRAILS_FAILED_MESSAGE,
             status: "guardrails_jailbreak",
           )
         end
-      rescue Guardrails::JailbreakChecker::ResponseError => e
-        context.abort_pipeline!(
-          message: Answer::CannedResponses::JAILBREAK_GUARDRAILS_FAILED_MESSAGE,
-          status: "error_jailbreak_guardrails",
-          jailbreak_guardrails_status: :error,
-          metrics: { "jailbreak_guardrails" => build_metrics(start_time, e) },
-          llm_response: { "jailbreak_guardrails" => e.llm_response },
-        )
       end
 
     private
 
-      attr_reader :llm_provider
+      attr_reader :context, :model_id, :model_name
+
+      def anthropic_bedrock_client
+        @anthropic_bedrock_client ||= Anthropic::BedrockClient.new(
+          aws_region: ENV["CLAUDE_AWS_REGION"],
+        )
+      end
+
+      def guardrails_llm_prompts
+        AnswerComposition::Pipeline::Claude.prompt_config(:jailbreak_guardrails, model_name)
+      end
+
+      # TODO: Move the common prompts into the claude config and use one set of prompts here.
+      def common_guardrails_llm_prompts
+        Rails.configuration.govuk_chat_private.llm_prompts.common.jailbreak_guardrails
+      end
+
+      def pass_value
+        common_guardrails_llm_prompts.fetch(:pass_value)
+      end
+
+      def max_tokens
+        guardrails_llm_prompts.fetch(:max_tokens)
+      end
+
+      def inference_config
+        {
+          max_tokens: max_tokens,
+          temperature: 0.0,
+        }
+      end
+
+      def messages
+        [{ role: "user", content: user_prompt }]
+      end
+
+      def user_prompt
+        guardrails_llm_prompts[:user_prompt].sub("{input}", context.question.message)
+      end
+
+      def system_prompt
+        guardrails_llm_prompts[:system_prompt]
+      end
 
-      def build_metrics(start_time, response_or_error)
+      def build_metrics(start_time, response)
         {
           duration: Clock.monotonic_time - start_time,
-          llm_prompt_tokens: response_or_error.llm_prompt_tokens,
-          llm_completion_tokens: response_or_error.llm_completion_tokens,
-          llm_cached_tokens: response_or_error.llm_cached_tokens,
-          model: response_or_error.model,
+          llm_prompt_tokens: response[:usage][:input_tokens],
+          llm_completion_tokens: response[:usage][:output_tokens],
+          llm_cached_tokens: nil,
+          model: response[:model],
         }
       end
     end

lib/guardrails/claude/jailbreak_checker.rb

@@ -21,13 +21,10 @@ namespace :evaluation do
   task generate_jailbreak_guardrail_response: :environment do
     raise "Requires an INPUT env var" if ENV["INPUT"].blank?
 
-    begin
-      response = Guardrails::JailbreakChecker.call(ENV["INPUT"], :claude)
+    question = Question.new(message: ENV["INPUT"], conversation: Conversation.new)
+    answer = AnswerComposition::PipelineRunner.call(question:, pipeline: [AnswerComposition::Pipeline::JailbreakGuardrails])
 
-      puts({ success: response }.to_json)
-    rescue Guardrails::JailbreakChecker::ResponseError => e
-      puts({ response_error: e }.to_json)
-    end
+    puts(answer.serialize_for_evaluation.to_json)
   end
 
   desc "Produce the output guardrails response for a user input"

lib/guardrails/jailbreak_checker.rb

@@ -26,10 +26,9 @@ def stub_pipeline_initialize(klass, *args, **kwargs)
       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)
-        stub_pipeline_initialize(AnswerComposition::Pipeline::JailbreakGuardrails, llm_provider: :claude)
 
         expected_pipeline = [
-          AnswerComposition::Pipeline::JailbreakGuardrails.new(llm_provider: :claude),
+          AnswerComposition::Pipeline::JailbreakGuardrails,
           AnswerComposition::Pipeline::QuestionRephraser,
           AnswerComposition::Pipeline::QuestionRouter,
           AnswerComposition::Pipeline::QuestionRoutingGuardrails.new(llm_provider: :claude),