Merge pull request #1002 from alphagov/472-use-claude-as-default

[CHAT-472] Use Claude as default LLM provider for answer composition

Author: davidgisbey

lib/answer_composition/pipeline/jailbreak_guardrails.rb

@@ -1,7 +1,7 @@
 module AnswerComposition
   module Pipeline
     class JailbreakGuardrails
-      def initialize(llm_provider: :openai)
+      def initialize(llm_provider: :claude)
         @llm_provider = llm_provider
       end
 

lib/answer_composition/pipeline/output_guardrails.rb

@@ -1,7 +1,7 @@
 module AnswerComposition
   module Pipeline
     class OutputGuardrails
-      def initialize(llm_provider: :openai)
+      def initialize(llm_provider: :claude)
         @llm_provider = llm_provider
       end
 

lib/guardrails/jailbreak_checker.rb

@@ -50,7 +50,7 @@ def self.guardrails_llm_prompts
 
     def self.call(...) = new(...).call
 
-    def initialize(input, llm_provider = :openai)
+    def initialize(input, llm_provider = :claude)
       @input = input
       @llm_provider = llm_provider
     end

spec/factories/output_guardrail_result_factory.rb

@@ -5,16 +5,29 @@
     llm_prompt_tokens { 13 }
     llm_completion_tokens { 7 }
     llm_cached_tokens { 10 }
-    model { "gpt-4o-mini-2024-07-18" }
+    model { BedrockModels.model_id(Guardrails::Claude::MultipleChecker::DEFAULT_MODEL) }
 
     llm_response do
-      {
-        "message": {
-          "role": "assistant",
-          "content": llm_guardrail_result,
-        },
-        "finish_reason": "stop",
-      }
+      content = Anthropic::Models::TextBlock.new(
+        type: :text,
+        text: llm_guardrail_result,
+      )
+
+      usage = Anthropic::Models::Usage.new(
+        input_tokens: llm_prompt_tokens,
+        output_tokens: llm_completion_tokens,
+        cache_read_input_tokens: llm_cached_tokens,
+      )
+
+      Anthropic::Models::Message.new(
+        id: "msg-id",
+        model:,
+        role: :assistant,
+        content:,
+        stop_reason: :end_turn,
+        usage:,
+        type: :message,
+      ).to_h
     end
 
     trait :pass do

spec/lib/answer_composition/composer_spec.rb

@@ -106,7 +106,7 @@ def stub_pipeline_initialize(klass, *args, **kwargs)
     end
 
     context "when an error is returned during answer generation" do
-      let(:question) { create :question, answer_strategy: :openai_structured_answer }
+      let(:question) { create :question, answer_strategy: :claude_structured_answer }
       let(:result) { described_class.call(question) }
 
       before do

spec/lib/answer_composition/pipeline/answer_guardrails_spec.rb

@@ -33,7 +33,7 @@
     it_behaves_like "a passing guardrail pipeline step", "answer_guardrails"
 
     it "does not abort the pipeline" do
-      described_class.new(llm_provider: :openai).call(context)
+      described_class.new(llm_provider: :claude).call(context)
       expect(context.aborted?).to be false
     end
   end
@@ -43,7 +43,7 @@
 
     it "aborts the pipeline and updates the answer's status and message attributes" do
       expect {
-        described_class.new(llm_provider: :openai).call(context)
+        described_class.new(llm_provider: :claude).call(context)
       }.to throw_symbol(:abort)
 
       expect(context.answer).to have_attributes(
@@ -55,21 +55,21 @@
     end
 
     it "assigns the llm response to the answer" do
-      expect { described_class.new(llm_provider: :openai).call(context) }.to throw_symbol(:abort)
+      expect { described_class.new(llm_provider: :claude).call(context) }.to throw_symbol(:abort)
       expect(context.answer.llm_responses["answer_guardrails"]).to eq(guardrail_response.llm_response)
     end
 
     it "assigns metrics to the answer" do
       allow(Clock).to receive(:monotonic_time).and_return(100.0, 101.5)
 
-      expect { described_class.new(llm_provider: :openai).call(context) }.to throw_symbol(:abort)
+      expect { described_class.new(llm_provider: :claude).call(context) }.to throw_symbol(:abort)
 
       expect(context.answer.metrics["answer_guardrails"]).to eq({
         duration: 1.5,
         llm_prompt_tokens: 13,
         llm_completion_tokens: 7,
         llm_cached_tokens: 10,
-        model: "gpt-4o-mini-2024-07-18",
+        model: BedrockModels.model_id(Guardrails::Claude::MultipleChecker::DEFAULT_MODEL),
       })
     end
   end

spec/lib/answer_composition/pipeline/jailbreak_guardrails_spec.rb

@@ -1,6 +1,6 @@
 RSpec.describe AnswerComposition::Pipeline::JailbreakGuardrails do
   let(:context) { build(:answer_pipeline_context) }
-  let(:default_provider) { :openai }
+  let(:default_provider) { :claude }
 
   let(:llm_response) do
     {
@@ -17,7 +17,7 @@
   let(:llm_prompt_tokens) { 10 }
   let(:llm_completion_tokens) { 5 }
   let(:llm_cached_tokens) { 0 }
-  let(:model) { Guardrails::OpenAI::JailbreakChecker::OPENAI_MODEL }
+  let(:model) { Guardrails::Claude::MultipleChecker.bedrock_model }
 
   context "when the guardrails are not triggered" do
     before do

spec/lib/answer_composition/pipeline/question_routing_guardrails_spec.rb

@@ -37,11 +37,11 @@
 
       context.answer.question_routing_label = "genuine_rag"
 
-      described_class.new(llm_provider: :openai).call(context)
+      described_class.new(llm_provider: :claude).call(context)
     end
 
     it "aborts the pipeline" do
-      described_class.new(llm_provider: :openai).call(context)
+      described_class.new(llm_provider: :claude).call(context)
       expect(context.aborted?).to be true
     end
   end
@@ -50,7 +50,7 @@
     let(:guardrail_response) { build(:guardrails_multiple_checker_result, :fail) }
 
     it "sets the attributes on the answer" do
-      described_class.new(llm_provider: :openai).call(context)
+      described_class.new(llm_provider: :claude).call(context)
 
       expect(context.answer).to have_attributes({
         message: Answer::CannedResponses::QUESTION_ROUTING_GUARDRAILS_FAILED_MESSAGE,
@@ -60,7 +60,7 @@
     end
 
     it "aborts the pipeline and assigns the right attributes" do
-      described_class.new(llm_provider: :openai).call(context)
+      described_class.new(llm_provider: :claude).call(context)
 
       expect(context.aborted?).to be true
       expect(context.answer.question_routing_guardrails_status).to eq("fail")

spec/lib/guardrails/jailbreak_checker_spec.rb

@@ -1,12 +1,12 @@
-RSpec.describe Guardrails::JailbreakChecker do
+RSpec.describe Guardrails::JailbreakChecker, :aws_credentials_stubbed do
   let(:input) { "User question" }
   let(:pass_value) { "PassValue" }
 
   before do
     allow(described_class).to receive_messages(pass_value:)
   end
 
-  it "calls the OpenAI jailbreak checker by default" do
+  it "calls the Claude jailbreak checker by default" do
     result = {
       llm_guardrail_result: pass_value,
       llm_response: {
@@ -21,12 +21,12 @@
       },
     }
 
-    allow(Guardrails::OpenAI::JailbreakChecker).to receive(:call).and_return(result)
+    allow(Guardrails::Claude::JailbreakChecker).to receive(:call).and_return(result)
     described_class.call(input)
-    expect(Guardrails::OpenAI::JailbreakChecker).to have_received(:call).with(input)
+    expect(Guardrails::Claude::JailbreakChecker).to have_received(:call).with(input)
   end
 
-  it "calls the Claude jailbreak checker when the provider is specified as :claude" do
+  it "calls the OpenAI jailbreak checker when the provider is specified as :openai" do
     result = {
       llm_guardrail_result: pass_value,
       llm_response: {
@@ -41,31 +41,33 @@
       },
     }
 
-    allow(Guardrails::Claude::JailbreakChecker).to receive(:call).and_return(result)
-    described_class.call(input, :claude)
-    expect(Guardrails::Claude::JailbreakChecker).to have_received(:call).with(input)
+    allow(Guardrails::OpenAI::JailbreakChecker).to receive(:call).and_return(result)
+    described_class.call(input, :openai)
+    expect(Guardrails::OpenAI::JailbreakChecker).to have_received(:call).with(input)
   end
 
   it "returns a result object" do
-    stub_openai_jailbreak_guardrails(input)
+    stub_claude_jailbreak_guardrails(input)
 
     result = described_class.call(input)
+
     expect(result)
       .to be_an_instance_of(described_class::Result)
       .and have_attributes(
         triggered: boolean,
-        llm_response: hash_including("message", "finish_reason", "index"),
+        llm_response: hash_including(:content, :stop_reason),
         llm_prompt_tokens: be_a(Integer),
         llm_completion_tokens: be_a(Integer),
         llm_cached_tokens: be_a(Integer).or(be_nil),
+        model: be_a(String),
       )
   end
 
   it "returns a result object with triggered true when guardrails fail" do
-    allow(Guardrails::OpenAI::JailbreakChecker).to receive(:call).with(input).and_call_original
-    stub_openai_jailbreak_guardrails(input, triggered: true)
+    allow(Guardrails::Claude::JailbreakChecker).to receive(:call).with(input).and_call_original
+    stub_claude_jailbreak_guardrails(input, triggered: true)
     second_input = "Second user question"
-    allow(Guardrails::OpenAI::JailbreakChecker)
+    allow(Guardrails::Claude::JailbreakChecker)
       .to receive(:call)
       .with(second_input)
       .and_return({
@@ -88,7 +90,7 @@
   end
 
   it "returns a result object with triggered false when guardrails pass" do
-    stub_openai_jailbreak_guardrails(input, triggered: false)
+    stub_claude_jailbreak_guardrails(input, triggered: false)
     expect(described_class.call(input)).to have_attributes(triggered: false)
   end
 end

spec/lib/guardrails/multiple_checker_spec.rb

@@ -7,18 +7,12 @@
     let(:llm_prompt_name) { :answer_guardrails }
     let(:guardrail_response_hash) do
       {
-        llm_response: {
-          message: {
-            role: "assistant",
-            content: "False | None",
-          },
-          finish_reason: "stop",
-        },
+        llm_response: guardrail_result.llm_response,
         llm_guardrail_result: "False | None",
         llm_prompt_tokens: 13,
         llm_completion_tokens: 7,
         llm_cached_tokens: 10,
-        model: "gpt-4o-mini-2024-07-18",
+        model: BedrockModels.model_id(Guardrails::Claude::MultipleChecker::DEFAULT_MODEL),
       }
     end
     let(:guardrail_result) { build(:guardrails_multiple_checker_result, :pass) }