はじめに
以前の記事で、生成AIを使ってレポートを作成するLambda関数を実装しました。
miyohide.hatenablog.com
当初は生成結果をMarkdown形式で出力していました。これは人が読むには非常に便利なのですが、実際にシステムを作り始めると少し問題が見えてきました。
今回作っている仕組みでは、AIが生成した内容をそのまま画面へ表示するだけではありません。今後は例えば
- 指摘事項だけを一覧表示したい
- 重要度で並び替えたい
- 指摘件数を集計したい
- PDFやPowerPointへ変換したい
- 後続のAIに入力として渡したい
といった処理を考えると、Markdown出力ではこれらを実現するたびにパース処理を書く必要があります。そこで、出力結果をJSONにすることにしました。
プロンプトだけでJSONを強制するのは難しい
プロンプトで JSONのみを出力してください とお願いすることはできますが、余計な文章が付いたり、JSONの途中で改行が崩れたり、キー名が微妙に変わったりすることがあります。人が読む分には問題ありませんが、プログラムで扱うには非常に困ります。JSONとしてパースできなければ、その後の処理がすべて止まってしまいます。
Structured Outputを使う
Amazon BedrockのStructured Outputを利用するようにしました。Structured Outputでは、あらかじめJSON Schemaを定義しておきます。生成AIはそのSchemaに従ったJSONだけを返してくれます。
ただし、Amazon Nova Lite 2ではStructured Outputには対応していません。公式ドキュメントの日本語版では「構造化出力」と表現されているものです。
docs.aws.amazon.com
そこでサポートしているGemma 3 27B PTを使うことにしました。
docs.aws.amazon.com
実装
具体的な実装は以下のとおり。
require "json"
require "aws-sdk-s3"
require "aws-sdk-dynamodb"
require "aws-sdk-bedrockruntime"
def lambda_handler(event:, context:)
s3 = Aws::S3::Client.new
dynamodb = Aws::DynamoDB::Client.new
input_bucket = ENV.fetch("INPUT_BUCKET")
output_bucket = ENV.fetch("OUTPUT_BUCKET")
table_name = ENV.fetch("TABLE_NAME")
job_id = event["jobId"]
begin
item = dynamodb.get_item(
table_name: table_name,
key: { "jobId" => job_id }
).item
raise "Job not found: #{job_id}" if item.nil?
template_key = item["templateKey"]
checkov_key = item["checkovKey"]
template = s3.get_object(
bucket: input_bucket,
key: template_key
).body.read
checkov = s3.get_object(
bucket: output_bucket,
key: checkov_key
).body.read
prompt = build_prompt(
template: template,
checkov: checkov
)
update_status(dynamodb, table_name, job_id, "GENERATING_REPORT")
report = generate_report(prompt)
update_status(dynamodb, table_name, job_id, "COMPLETED")
report_key = "jobs/#{job_id}/report/report.json"
s3.put_object(
bucket: output_bucket,
key: report_key,
body: report,
content_type: "application/json"
)
dynamodb.update_item(
table_name: table_name,
key: { "jobId" => job_id },
update_expression: "SET reportJSONKey = :key, updatedAt = :updated_at",
expression_attribute_values: {
":key" => report_key,
":updated_at" => Time.now.utc.iso8601
}
)
{
reportKey: report_key
}
rescue => e
update_status(dynamodb, table_name, job_id, "FAILED") if job_id
raise e
end
end
def update_status(dynamodb, table_name, job_id, status)
dynamodb.update_item(
table_name: table_name,
key: { "jobId" => job_id },
update_expression: "SET jobStatus = :status, updatedAt = :updated_at",
expression_attribute_values: {
":status" => status,
":updated_at" => Time.now.utc.iso8601
}
)
end
def summarize_checkov(json_text)
data = JSON.parse(json_text)
failed =
data
.dig("results", "failed_checks") || []
failed.map do |c|
{
id: c["check_id"],
name: c["check_name"],
resource: c["resource"],
file_line_range: c["file_line_range"]
}
end
end
def build_prompt(template:, checkov:)
failed_checks =
JSON.pretty_generate(
summarize_checkov(checkov)
)
<<~PROMPT
あなたはAWSセキュリティレビュー担当者です。
CloudFormationとCheckovの結果をレビューし結果を日本語で出力してください。
各findingには、問題が発生したテンプレートの行範囲(line_range)を含めてください。
Checkov結果のfile_line_rangeを参照してください。
# CloudFormation
#{template}
# Checkov結果
#{failed_checks}
PROMPT
end
def report_schema
{
type: "object",
properties: {
summary: { type: "string", description: "全体のサマリ" },
findings: {
type: "array",
items: {
type: "object",
properties: {
severity: { type: "string", enum: ["Critical", "High", "Medium", "Low"] },
issue: { type: "string", description: "問題点" },
reason: { type: "string", description: "なぜ危険か" },
attack_scenario: { type: "string", description: "攻撃シナリオ" },
recommendation: { type: "string", description: "推奨修正" },
line_range: {
type: "array",
items: { type: "integer" },
description: "問題が発生したテンプレートの行範囲 [開始行, 終了行]"
}
},
required: ["severity", "issue", "reason", "attack_scenario", "recommendation", "line_range"],
additionalProperties: false
}
},
overall_assessment: { type: "string", description: "総評" }
},
required: ["summary", "findings", "overall_assessment"],
additionalProperties: false
}.to_json
end
def generate_report(prompt)
client =
Aws::BedrockRuntime::Client.new
response =
client.converse(
model_id: ENV.fetch("MODEL_ID"),
messages: [
{
role: "user",
content: [
{ text: prompt }
]
}
],
output_config: {
text_format: {
type: "json_schema",
structure: {
json_schema: {
schema: report_schema,
name: "security_review_report",
description: "IaCセキュリティレビューレポート"
}
}
}
}
)
raw_text = response.output.message.content[0].text
parsed = JSON.parse(raw_text)
JSON.pretty_generate(parsed)
end
JSONの定義は report_schema メソッドで実装しています。これを使ってconverseメソッドにてoutput_configを指定することでstructured outputを実装することができます。
まとめ
生成AIを使ったシステムでは、最終的に人が読むからといってMarkdownを出力することが最適とは限りません。むしろ、後続処理まで考えるとJSONで受け取り、画面表示の直前でMarkdownやHTMLへ変換する方が扱いやすいケースが多いです。
今回Structured Outputへ変更したことで、今後の実装に幅が広がると考えています。生成AIを「チャットツール」として使うだけでなく、「システムの部品」として組み込む場合には、Structured Outputは非常に有効な機能だと感じています。
