Merge pull request #1050 from alphagov/473-spike-using-structured-response-for-guardrails

[CHAT-473][DO NOT MERGE] Use structured output for non-sonnet 4.0 guardrail requests

Author: davidgisbey

lib/answer_composition/multiple_guardrail/checker.rb

@@ -27,14 +27,36 @@ def initialize(input, llm_prompt_name)
     end
 
     def call
-      response = anthropic_bedrock_client.messages.create(
+      shared_config = {
         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)
+      }
+
+      if self.class.bedrock_model == :claude_sonnet_4_0
+        response = anthropic_bedrock_client.messages.create(**shared_config)
+        parse_response(response)
+      else
+        response = anthropic_bedrock_client.messages.create(**shared_config.merge(
+          output_config: {
+            format: json_schema,
+          },
+        ))
+
+        llm_guardrail_result = JSON.parse(response.content.first.text)
+
+        Result.new(
+          llm_response: response.to_h,
+          llm_guardrail_result: llm_guardrail_result.to_s,
+          triggered: llm_guardrail_result.present?,
+          guardrails: to_guardrail_hash(llm_guardrail_result),
+          llm_prompt_tokens: response[:usage][:input_tokens],
+          llm_completion_tokens: response[:usage][:output_tokens],
+          llm_cached_tokens: response[:usage][:cache_read_input_tokens],
+          model: response.model,
+        )
+      end
     end
 
   private
@@ -67,7 +89,8 @@ def parse_response(response)
 
       parts = llm_guardrail_result.split(" | ")
       triggered = parts.first.chomp == "True"
-      guardrails = to_guardrail_hash(parts.second)
+      triggered_guardrail_numbers = parts.second.scan(/\d+/).map(&:to_i)
+      guardrails = to_guardrail_hash(triggered_guardrail_numbers)
 
       Result.new(
         llm_response:,
@@ -96,12 +119,22 @@ def response_pattern
       end
     end
 
-    def to_guardrail_hash(parts)
-      triggered_guardrail_numbers = parts.scan(/\d+/).map(&:to_i)
-
+    def to_guardrail_hash(triggered_guardrail_numbers)
       prompt.guardrails.each_with_object({}) do |guardrail, guardrails_hash|
         guardrails_hash[guardrail.name.to_sym] = triggered_guardrail_numbers.include?(guardrail.key)
       end
     end
+
+    def json_schema
+      guardrail_keys = prompt.guardrails.map(&:key)
+      {
+        "type" => "json_schema",
+        "schema" => {
+          "description" => "Array of triggered guardrail numbers. Returns [] if none triggered.",
+          "type" => "array",
+          "items" => { "type" => "integer", "enum" => guardrail_keys },
+        },
+      }
+    end
   end
 end

lib/answer_composition/multiple_guardrail/prompt.rb

@@ -16,7 +16,13 @@ def system_prompt
       guardrails_content = guardrails.map { |g| "#{g.key}. #{g.content}" }
                                      .join("\n")
 
-      prompts.fetch(:system_prompt)
+      system_prompt_key = if Checker.bedrock_model == :claude_sonnet_4_0
+                            :system_prompt
+                          else
+                            :system_prompt_structured
+                          end
+
+      prompts.fetch(system_prompt_key)
              .sub("{guardrails}", guardrails_content)
              .sub("{date}", Date.current.strftime("%A %d %B %Y"))
     end