Claudeのsearch_result引用を回帰テストする — 引用元の取り違えを検出

# citecheck.py
def collect_results(messages):
    found = []
    for m in messages:
        if isinstance(m["content"], str):
            continue
        for b in m["content"]:
            if b["type"] == "search_result":
                found.append(b)
            elif b["type"] == "tool_result" and isinstance(b["content"], list):
                found += [x for x in b["content"] if x["type"] == "search_result"]
    return found
 
 
def check_citations(messages, response_content):
    """引用1件ごとの違反を文字列で返す。空なら合格。"""
    results = collect_results(messages)
    errors = []
    for block in response_content:
        for c in block.get("citations") or []:
            if c["type"] != "search_result_location":
                continue
            i = c["search_result_index"]
            if not 0 <= i < len(results):
                errors.append(f"index {i} は範囲外(全{len(results)}件)")
                continue
            r = results[i]
            if c["source"] != r["source"]:
                errors.append(f"index {i}: source {c['source']} != {r['source']}")
            if c["title"] != r["title"]:
                errors.append(f"index {i}: title 不一致")
            s, e = c["start_block_index"], c["end_block_index"]
            if not 0 <= s < e <= len(r["content"]):
                errors.append(f"index {i}: ブロック範囲 [{s}:{e}] が不正")
                continue
            text = "".join(x["text"] for x in r["content"][s:e])
            if c["cited_text"] != text:
                errors.append(f"index {i}: cited_text が content[{s}:{e}] と不一致")
    return errors
 
 
def cited_sources(response_content):
    return {c["source"] for b in response_content
            for c in (b.get("citations") or [])}

# test_offline.py
from citecheck import check_citations, cited_sources
 
def sr(source, title, *texts):
    return {"type": "search_result", "source": source, "title": title,
            "content": [{"type": "text", "text": t} for t in texts],
            "citations": {"enabled": True}}
 
MESSAGES = [{"role": "user", "content": [
    sr("kb://billing-2024", "請求ガイド", "返金は30日以内。"),
    sr("kb://billing-2025", "請求ガイド", "返金は14日以内。", "分割払いは不可。"),
    {"type": "text", "text": "返金期限は?"},
]}]
 
def cite(i, s, e, src, title, text):
    return {"type": "search_result_location", "cited_text": text, "source": src,
            "title": title, "search_result_index": i,
            "start_block_index": s, "end_block_index": e}
 
GOOD = [{"type": "text", "text": "14日以内です。",
         "citations": [cite(1, 0, 1, "kb://billing-2025", "請求ガイド", "返金は14日以内。")]}]
SWAPPED = [{"type": "text", "text": "14日以内です。",
            "citations": [cite(0, 0, 1, "kb://billing-2025", "請求ガイド", "返金は14日以内。")]}]
 
def test_good():
    assert check_citations(MESSAGES, GOOD) == []
    assert cited_sources(GOOD) == {"kb://billing-2025"}
 
def test_swapped_index_is_caught():
    errs = check_citations(MESSAGES, SWAPPED)
    assert any("source" in e for e in errs)
    assert any("cited_text" in e for e in errs)

# test_live.py
import itertools, os
import anthropic
from citecheck import check_citations, cited_sources
from test_offline import sr
 
MODEL = os.environ["CITE_TEST_MODEL"]  # 検査対象のモデルIDを環境変数で渡す
client = anthropic.Anthropic()
 
RESULTS = [
    sr("kb://billing-2024", "請求ガイド(旧)", "返金は30日以内。"),
    sr("kb://billing-2025", "請求ガイド(新)", "返金は14日以内。"),
    sr("kb://shipping", "配送ガイド", "通常配送は3営業日。"),
]
CASES = [("現行の返金期限は?", {"kb://billing-2025"})]
 
def ask(results, question):
    messages = [{"role": "user", "content": results
                 + [{"type": "text", "text": question}]}]
    resp = client.messages.create(model=MODEL, max_tokens=512, messages=messages)
    return messages, [b.model_dump() for b in resp.content]
 
def test_expected_source_for_every_order():
    for question, expected in CASES:
        for perm in itertools.permutations(RESULTS):
            messages, content = ask(list(perm), question)
            assert check_citations(messages, content) == []
            got = cited_sources(content)
            assert got and got <= expected, (question, [r["source"] for r in perm], got)

import json, pathlib
 
def ask_and_record(messages, tag):
    resp = client.messages.create(model=MODEL, max_tokens=512, messages=messages)
    content = [b.model_dump() for b in resp.content]
    pathlib.Path(f"fixtures/{tag}.json").write_text(
        json.dumps({"messages": messages, "content": content},
                   ensure_ascii=False, indent=2))
    return content