Amazon BedrockのClaudeモデルでPDFチャート画像の認識精度を検証する

"""Bedrock-Claude PDFチャート認識検証の共通ユーティリティ。"""
import json
import time
from pathlib import Path

import boto3

REGION = "us-east-1"
PROFILE = "xxxxxxxxx"

MODELS = {
    "sonnet-4-6": "us.anthropic.claude-sonnet-4-6",
    "opus-4-6":   "us.anthropic.claude-opus-4-6-v1",
    "opus-4-7":   "us.anthropic.claude-opus-4-7",
}

EMBED_MODEL_ID = "amazon.titan-embed-text-v2:0"
EMBED_DIM = 1024
CHUNK_SIZE = 1024
CHUNK_OVERLAP = CHUNK_SIZE // 10  # 10%オーバーラップでチャンク境界の文脈を保持

ROOT = Path(__file__).resolve().parent.parent
PDF_PATH = ROOT / "data" / "n1120000.pdf"
QUESTIONS_PATH = ROOT / "scripts" / "questions.json"
RESULTS_DIR = ROOT / "results"


def session():
    """AWS SSO プロファイルを使ったセッションを返す。"""
    return boto3.Session(profile_name=PROFILE, region_name=REGION)


def bedrock_runtime():
    """Bedrock Runtime クライアントを返す。"""
    return session().client("bedrock-runtime")


def s3vectors():
    """S3 Vectors クライアントを返す。"""
    return session().client("s3vectors")


def load_questions():
    """questions.json から評価質問リストを読み込む。"""
    return json.loads(QUESTIONS_PATH.read_text(encoding="utf-8"))


def load_pdf_bytes():
    """検証対象 PDF をバイト列で返す。"""
    return PDF_PATH.read_bytes()


def converse(client, model_id, content, max_tokens=2048):
    """Bedrock Converse API を呼び出し、テキスト回答とトークン使用量を返す。

    Opus 4.7 は temperature が廃止されているため additionalModelRequestFields で
    thinking と effort を設定する。他モデルは temperature=0 で出力を固定する。
    失敗時は指数バックオフで最大3回リトライする。
    """
    inference_config = {"maxTokens": max_tokens}
    extra = {}
    if "opus-4-7" in model_id:
        # Opus 4.7: temperature は Breaking Change で廃止。thinking と effort で制御する
        extra["additionalModelRequestFields"] = {
            "thinking": {"type": "adaptive"},
            "output_config": {"effort": "high"},
        }
    else:
        # Opus 4.7 以外: temperature=0 で出力のランダム性を抑制する
        inference_config["temperature"] = 0
    started = time.monotonic()
    for attempt in range(3):
        try:
            resp = client.converse(
                modelId=model_id,
                messages=[{"role": "user", "content": content}],
                inferenceConfig=inference_config,
                **extra,
            )
            break
        except Exception:
            if attempt == 2:
                raise
            time.sleep(2 ** attempt)  # 指数バックオフ: 1s → 2s
    blocks = resp["output"]["message"]["content"]
    text = next(b["text"] for b in blocks if "text" in b)
    usage = resp.get("usage", {})
    return text, {
        "input_tokens": usage.get("inputTokens", 0),
        "output_tokens": usage.get("outputTokens", 0),
        "latency_sec": round(time.monotonic() - started, 2),
    }


def embed(client, text):
    """Titan Embed Text v2 でテキストを 1024 次元のベクトルに変換する。"""
    resp = client.invoke_model(
        modelId=EMBED_MODEL_ID,
        body=json.dumps({"inputText": text, "dimensions": EMBED_DIM, "normalize": True}),
    )
    return json.loads(resp["body"].read())["embedding"]


def chunk_text(text):
    """テキストを CHUNK_SIZE 文字ごとに分割し、CHUNK_OVERLAP 文字のオーバーラップを持たせる。"""
    step = CHUNK_SIZE - CHUNK_OVERLAP
    return [text[i:i + CHUNK_SIZE] for i in range(0, len(text), step)]


def write_results(name, md, payload):
    """結果を Markdown と JSON の2形式で results/ ディレクトリに書き出す。"""
    RESULTS_DIR.mkdir(exist_ok=True)
    (RESULTS_DIR / f"{name}.md").write_text(md, encoding="utf-8")
    (RESULTS_DIR / f"{name}.json").write_text(
        json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
    )

"""Approach 1: Bedrock Converse API の document ブロックでPDFを直接読ませる。"""
import common


def ask(client, model_id, pdf_bytes, question):
    """PDF バイナリを document ブロックに渡してモデルに質問する。"""
    content = [
        # PDF を document ブロックとして直接渡す（Bedrock は最大 4.5MB をサポート、Claude モデルは最大 100 ページをサポート）
        {"document": {"format": "pdf", "name": "target", "source": {"bytes": pdf_bytes}}},
        {"text": question},
    ]
    return common.converse(client, model_id, content)


def to_markdown(results):
    """全モデルの回答結果を1つの Markdown レポートにまとめる。"""
    lines = ["# Approach 1: PDF直接読み取り", ""]
    any_key = next(iter(results))
    for i, row in enumerate(results[any_key]):
        q = row["question"]
        lines += [f"## {q['id']}: {q['question']}", "", f"観点: {q.get('notes', '')}", ""]
        for key, rows in results.items():
            r = rows[i]
            header = f"### {key} (in/out={r['input_tokens']}/{r['output_tokens']}, {r['latency_sec']}s)"
            lines += [header, "", r["answer"], ""]
    return "\n".join(lines)


def main():
    """3モデル × 全質問で PDF 直接読み取りを実行し、結果を保存する。"""
    client = common.bedrock_runtime()
    pdf_bytes = common.load_pdf_bytes()
    questions = common.load_questions()

    results = {}
    for key, model_id in common.MODELS.items():
        print(f"== {key} ==")
        rows = []
        for q in questions:
            print(f"  {q['id']}")
            answer, m = ask(client, model_id, pdf_bytes, q["question"])
            rows.append({"question": q, "answer": answer, **m})
        results[key] = rows

    common.write_results("approach1-pdf-direct", to_markdown(results), results)
    print("results/approach1-pdf-direct.md")


if __name__ == "__main__":
    main()

"""Approach 2 / Step 1: PDFをベクトル化してS3 Vectorsへ登録する。"""
import argparse
import io
import json
import uuid

from pdf2image import convert_from_bytes

import common

IMAGE_PROMPT = """\
Extract the content from an image page and output in Markdown syntax. Enclose the content in the <markdown></markdown> tag and do not use code blocks. If the image is empty then output a <markdown></markdown> without anything in it.

Follow these steps:

1. Examine the provided page carefully.

2. Identify all elements present in the page, including headers, body text, footnotes, tables, images, captions, and page numbers, etc.

3. Use markdown syntax to format your output:
    - Headings: # for main, ## for sections, ### for subsections, etc.
    - Lists: * or - for bulleted, 1. 2. 3. for numbered
    - Do not repeat yourself

4. If the element is an image (not table)
    - If the information in the image can be represented by a table, generate the table containing the information of the image
    - Otherwise provide a detailed description about the information in image
    - For charts specifically, work through the following steps in order before writing the final table:
        - Step 1 - Structure: Identify the chart type (bar, stacked bar, line, pie, etc.), read the title, axis labels with units, and list all legend categories with their colors or patterns
        - Step 2 - Scale: Note the numerical range of each axis (min, max, major gridlines) so you can derive estimates
        - Step 3 - Per-item extraction: For every item (row/country/data point), read or estimate the value of EACH series separately. Do not skip any item
        - Step 4 - Estimate unlabeled values: When no label is printed, estimate by comparing bar length to the axis scale. Mark each estimate with a leading ~. Provide numeric estimates (e.g., ~85) rather than qualitative labels (e.g., "large")
        - Step 5 - Table: Output a table with one row per item and one column per series. Add a total column when totals are visible. Do not omit any item or series
    - Classify the element as one of: Chart, Diagram, Logo, Icon, Natural Image, Screenshot, Other. Enclose the class in <figure_type></figure_type>
    - Enclose <figure_type></figure_type>, the table or description, and the figure title or caption (if available), in <figure></figure> tags
    - Do not transcribe text in the image after providing the table or description

5. If the element is a table
    - Create a markdown table, ensuring every row has the same number of columns
    - Maintain cell alignment as closely as possible
    - Do not split a table into multiple tables
    - If a merged cell spans multiple rows or columns, place the text in the top-left cell and output ' ' for other
    - Use | for column separators, |-|-| for header row separators
    - If a cell has multiple items, list them in separate rows
    - If the table contains sub-headers, separate the sub-headers from the headers in another row

6. If the element is a paragraph
    - Transcribe each text element precisely as it appears

7. If the element is a header, footer, footnote, page number
    - Transcribe each text element precisely as it appears

Output Example: (省略)
"""


def render_pages(pdf_bytes, dpi=300):
    """PDF をページごとに PNG 画像に変換して返す。DPI が高いほど細部まで読める。"""
    pages = []
    for img in convert_from_bytes(pdf_bytes, dpi=dpi):
        buf = io.BytesIO()
        img.convert("RGB").save(buf, format="PNG")
        pages.append(buf.getvalue())
    return pages


def image_to_markdown(client, model_id, png_bytes):
    """PNG 画像を Claude に渡して Markdown テキストに変換する。トークン使用量も返す。"""
    content = [
        {"image": {"format": "png", "source": {"bytes": png_bytes}}},
        {"text": IMAGE_PROMPT},
    ]
    md, usage = common.converse(client, model_id, content, max_tokens=3500)
    return md, usage


def main():
    """PDF の全ページを画像化して Markdown に変換し、S3 Vectors インデックスに登録する。"""
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, choices=list(common.MODELS))
    ap.add_argument("--bucket", required=True)
    args = ap.parse_args()

    bedrock = common.bedrock_runtime()
    vectors = common.s3vectors()
    pdf_bytes = common.load_pdf_bytes()

    # 全ページを画像化し、Claude で Markdown に変換してチャンク分割する
    chunks = []
    page_usages = []
    for i, png in enumerate(render_pages(pdf_bytes)):
        print(f"  page {i}: image->markdown")
        md, usage = image_to_markdown(bedrock, common.MODELS[args.model], png)
        page_usages.append(usage)
        chunks.extend(common.chunk_text(md))

    # S3 Vectors は create_index が冪等でないため、再実行時は先に削除してから作成する
    index_name = f"chart-rag-{args.model}"
    try:
        vectors.delete_index(vectorBucketName=args.bucket, indexName=index_name)
    except vectors.exceptions.NotFoundException:
        pass
    vectors.create_index(
        vectorBucketName=args.bucket,
        indexName=index_name,
        dataType="float32",
        dimension=common.EMBED_DIM,
        distanceMetric="cosine",
        metadataConfiguration={"nonFilterableMetadataKeys": ["source_text"]},
    )

    # チャンクを埋め込みベクトルに変換し、500件ずつバッチ投入する（API の上限）
    records = [
        {
            "key": str(uuid.uuid4()),
            "data": {"float32": common.embed(bedrock, c)},
            "metadata": {"source_text": c[:4000]},
        }
        for c in chunks
    ]
    for i in range(0, len(records), 500):
        vectors.put_vectors(
            vectorBucketName=args.bucket,
            indexName=index_name,
            vectors=records[i:i + 500],
        )
    print(f"indexed {len(records)} vectors into {index_name}")

    # ページごとのトークン使用量を集計して JSON に保存する
    total_input = sum(u["input_tokens"] for u in page_usages)
    total_output = sum(u["output_tokens"] for u in page_usages)
    n = len(page_usages)
    ingest_stats = {
        "model": args.model,
        "pages": n,
        "total_input_tokens": total_input,
        "total_output_tokens": total_output,
        "avg_input_tokens_per_page": round(total_input / n),
        "avg_output_tokens_per_page": round(total_output / n),
        "per_page": page_usages,
    }
    out_path = common.RESULTS_DIR / f"approach2-rag_ingest-{args.model}.json"
    common.RESULTS_DIR.mkdir(exist_ok=True)
    out_path.write_text(json.dumps(ingest_stats, ensure_ascii=False, indent=2))
    print(f"ingest stats -> {out_path}")
    print(f"total input: {total_input}, total output: {total_output}, avg/page: {total_input//n}/{total_output//n}")


if __name__ == "__main__":
    main()

"""Approach 2 / Step 2: S3 Vectorsへクエリし、Claudeで回答する。"""
import argparse

import common

ANSWER_PROMPT = """\
あなたはPDFドキュメントのアナリストです。以下の検索結果のみを根拠に質問に回答してください。
検索結果に情報がない場合は「資料内に該当する情報がありません」と答えてください。

## 検索結果
{context}

## 質問
{question}
"""


def retrieve(vectors, bucket, index, query_vec, top_k=5):
    """質問ベクトルで S3 Vectors を検索し、上位 top_k チャンクのテキストを返す。"""
    resp = vectors.query_vectors(
        vectorBucketName=bucket,
        indexName=index,
        topK=top_k,
        queryVector={"float32": query_vec},
        returnMetadata=True,
    )
    return [v["metadata"]["source_text"] for v in resp.get("vectors", [])]


def answer(client, model_id, question, contexts):
    """検索コンテキストを組み込んだプロンプトで Claude に回答させる。"""
    # 複数チャンクを "---" で区切って1つのコンテキストブロックにまとめる
    prompt = ANSWER_PROMPT.format(context="\n---\n".join(contexts), question=question)
    return common.converse(client, model_id, [{"text": prompt}])


def to_markdown(args, rows):
    """回答結果を Markdown レポート形式に変換する。"""
    lines = [f"# Approach 2: RAG (ingest={args.ingest_model}, answer={args.answer_model})", ""]
    for r in rows:
        q = r["question"]
        lines += [
            f"## {q['id']}: {q['question']}",
            "",
            f"観点: {q.get('notes', '')}",
            "",
            f"- in/out tokens: {r['input_tokens']}/{r['output_tokens']}",
            f"- latency: {r['latency_sec']}s",
            "",
            "### 回答",
            r["answer"],
            "",
            "### 検索コンテキスト（先頭120文字）",
        ]
        for i, c in enumerate(r["contexts"]):
            lines.append(f"- [{i}] {c[:120]}...")
        lines.append("")
    return "\n".join(lines)


def main():
    """質問を埋め込み→ベクトル検索→Claude 回答の流れで RAG クエリを実行する。"""
    ap = argparse.ArgumentParser()
    ap.add_argument("--ingest-model", required=True, choices=list(common.MODELS))
    ap.add_argument("--answer-model", required=True, choices=list(common.MODELS))
    ap.add_argument("--bucket", required=True)
    args = ap.parse_args()

    bedrock = common.bedrock_runtime()
    vectors = common.s3vectors()
    # ingest 時に使ったモデルと対応するインデックスを参照する
    index_name = f"chart-rag-{args.ingest_model}"

    rows = []
    for q in common.load_questions():
        print(f"  {q['id']}")
        qv = common.embed(bedrock, q["question"])          # 質問をベクトル化
        contexts = retrieve(vectors, args.bucket, index_name, qv)  # 類似チャンクを取得
        text, m = answer(bedrock, common.MODELS[args.answer_model], q["question"], contexts)
        rows.append({"question": q, "contexts": contexts, "answer": text, **m})

    name = f"approach2-rag_ingest-{args.ingest_model}_answer-{args.answer_model}"
    common.write_results(name, to_markdown(args, rows), rows)
    print(f"results/{name}.md")


if __name__ == "__main__":
    main()

"""結果JSONから比較サマリーMarkdownを作成する。"""
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
RES = ROOT / "results"
MODELS = ("sonnet-4-6", "opus-4-6", "opus-4-7")

APP1 = json.loads((RES / "approach1-pdf-direct.json").read_text(encoding="utf-8"))
APP2 = {
    m: json.loads((RES / f"approach2-rag_ingest-{m}_answer-{m}.json").read_text(encoding="utf-8"))
    for m in MODELS
}

QIDS = [r["question"]["id"] for r in APP1[MODELS[0]]]


def rows_of(source):
    """質問 ID をキーにした辞書に変換する。"""
    return {r["question"]["id"]: r for r in source}


def fmt(r):
    """トークン数とレイテンシを 'input/output (Xs)' 形式の文字列にフォーマットする。"""
    return f"{r['input_tokens']}/{r['output_tokens']} ({r['latency_sec']}s)"


header = "| 実行 | " + " | ".join(QIDS) + " |"
sep = "| --- | " + " | ".join("---" for _ in QIDS) + " |"
lines = ["# 検証結果サマリー", "", "## 入力/出力トークンとレイテンシ", "", header, sep]
for m in MODELS:
    r = rows_of(APP1[m])
    lines.append(f"| A1 / {m} | " + " | ".join(fmt(r[q]) for q in QIDS) + " |")
for m in MODELS:
    r = rows_of(APP2[m])
    lines.append(f"| A2 / {m} | " + " | ".join(fmt(r[q]) for q in QIDS) + " |")

lines += ["",
          "A1 = Approach 1 (PDF直接読み取り), A2 = Approach 2 (RAG)",
          "フォーマット: `input_tokens/output_tokens (latency)`"]

(RES / "summary.md").write_text("\n".join(lines) + "\n", encoding="utf-8")
print("results/summary.md")