[アップデート] Amazon Bedrock Managed Knowledge Basesでメタデータの項目を検索の対象にできるようになったので試してみた

import json
import sys
import time
from pathlib import Path

import boto3

REGION = "us-east-1"
STATE = Path("state.json")
agent = boto3.client("bedrock-agent", region_name=REGION)
runtime = boto3.client("bedrock-agent-runtime", region_name=REGION)
iam = boto3.client("iam", region_name=REGION)

# 本文には作成者もプロジェクト名も書かず、メタデータにだけ持たせる
REPORTS = [
    {
        "id": "report-01",
        "author": "山田花子",
        "projects": ["コハク", "ルリ"],
        "text": "10月の月次報告です。問い合わせ対応の件数は前月より12%減りました。FAQページを改修した影響と見ています。来月はチャットボットの回答精度の改善に取り組みます。",
    },
    {
        "id": "report-02",
        "author": "佐藤健",
        "projects": ["ヒスイ"],
        "text": "10月の月次報告です。問い合わせ対応の件数は前月より5%増えました。新機能のリリース直後に質問が集中したためです。来月はリリースノートの書き方を見直します。",
    },
    {
        "id": "report-03",
        "author": "鈴木一郎",
        "projects": ["メノウ"],
        "text": "10月の月次報告です。問い合わせ対応の件数は前月とほぼ同じでした。夜間の一次回答をチャットボットに任せる準備を進めています。来月は回答テンプレートを整理します。",
    },
    {
        "id": "report-04",
        "author": "高橋美咲",
        "projects": ["サンゴ"],
        "text": "10月の月次報告です。問い合わせ対応の件数は前月より8%減りました。よくある質問を社内ポータルに掲載したためです。来月は担当者向けの研修資料を作ります。",
    },
]

QUERIES = ["山田花子さんが書いた月次報告", "コハクの進捗"]

def save(name, data):
    Path("evidence").mkdir(exist_ok=True)
    Path("evidence", name + ".json").write_text(
        json.dumps(data, ensure_ascii=False, indent=2, default=str)
    )

def state():
    return json.loads(STATE.read_text())

def wait_kb(kb_id):
    for _ in range(120):
        kb = agent.get_knowledge_base(knowledgeBaseId=kb_id)["knowledgeBase"]
        print("kb status", kb["status"], flush=True)
        if kb["status"] == "ACTIVE":
            return kb
        if kb["status"] in ("FAILED", "UPDATE_UNSUCCESSFUL"):
            raise RuntimeError(kb.get("failureReasons"))
        time.sleep(10)
    raise RuntimeError("timed out")

def setup():
    if STATE.exists():
        raise RuntimeError("state.json already exists")
    account = boto3.client("sts").get_caller_identity()["Account"]
    name = "kb-meta-search-" + str(int(time.time()))
    trust = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"Service": "bedrock.amazonaws.com"},
                "Action": "sts:AssumeRole",
                "Condition": {
                    "StringEquals": {"aws:SourceAccount": account},
                    "ArnLike": {
                        "aws:SourceArn": f"arn:aws:bedrock:{REGION}:{account}:knowledge-base/*"
                    },
                },
            }
        ],
    }
    role = iam.create_role(RoleName=name, AssumeRolePolicyDocument=json.dumps(trust))
    s = {"name": name, "roleArn": role["Role"]["Arn"]}
    STATE.write_text(json.dumps(s))
    time.sleep(15)
    kb = agent.create_knowledge_base(
        name=name,
        roleArn=s["roleArn"],
        knowledgeBaseConfiguration={
            "type": "MANAGED",
            "managedKnowledgeBaseConfiguration": {"embeddingModelType": "MANAGED"},
        },
    )["knowledgeBase"]
    s["knowledgeBaseId"] = kb["knowledgeBaseId"]
    STATE.write_text(json.dumps(s))
    save("created-kb", kb)
    wait_kb(s["knowledgeBaseId"])
    ds = agent.create_data_source(
        knowledgeBaseId=s["knowledgeBaseId"],
        name="custom",
        dataSourceConfiguration={
            "type": "MANAGED_KNOWLEDGE_BASE_CONNECTOR",
            "managedKnowledgeBaseConnectorConfiguration": {
                "connectorParameters": {"type": "CUSTOM", "version": "1"}
            },
        },
    )["dataSource"]
    s["dataSourceId"] = ds["dataSourceId"]
    STATE.write_text(json.dumps(s))
    print(json.dumps(s))

def ingest():
    s = state()
    documents = [
        {
            "content": {
                "dataSourceType": "CUSTOM",
                "custom": {
                    "customDocumentIdentifier": {"id": r["id"]},
                    "sourceType": "IN_LINE",
                    "inlineContent": {
                        "type": "TEXT",
                        "textContent": {"data": r["text"]},
                    },
                },
            },
            "metadata": {
                "type": "IN_LINE_ATTRIBUTE",
                "inlineAttributes": [
                    {
                        "key": "author",
                        "value": {"type": "STRING", "stringValue": r["author"]},
                    },
                    {
                        "key": "projects",
                        "value": {
                            "type": "STRING_LIST",
                            "stringListValue": r["projects"],
                        },
                    },
                ],
            },
        }
        for r in REPORTS
    ]
    params = {k: s[k] for k in ["knowledgeBaseId", "dataSourceId"]}
    agent.ingest_knowledge_base_documents(**params, documents=documents)
    ids = [{"dataSourceType": "CUSTOM", "custom": {"id": r["id"]}} for r in REPORTS]
    for _ in range(60):
        details = agent.get_knowledge_base_documents(
            **params, documentIdentifiers=ids
        )["documentDetails"]
        statuses = [d["status"] for d in details]
        print("documents", statuses, flush=True)
        if all(x == "INDEXED" for x in statuses):
            return
        time.sleep(10)
    raise RuntimeError("timed out")

def search(label="search"):
    s = state()
    out = {}
    for q in QUERIES:
        resp = runtime.retrieve(
            knowledgeBaseId=s["knowledgeBaseId"],
            retrievalQuery={"text": q},
            retrievalConfiguration={
                "managedSearchConfiguration": {"numberOfResults": 4}
            },
        )
        out[q] = resp["retrievalResults"]
        print("Q:", q)
        for i, r in enumerate(resp["retrievalResults"], 1):
            m = r.get("metadata", {})
            print(
                f"  [{i}] score={r.get('score', 0):.4f} author={m.get('author')} projects={m.get('projects')}"
            )
    save(label, out)

def enable():
    s = state()
    kb = agent.get_knowledge_base(knowledgeBaseId=s["knowledgeBaseId"])[
        "knowledgeBase"
    ]
    resp = agent.update_knowledge_base(
        knowledgeBaseId=s["knowledgeBaseId"],
        name=kb["name"],
        roleArn=kb["roleArn"],
        knowledgeBaseConfiguration=kb["knowledgeBaseConfiguration"],
        metadataConfigurations=[
            {"name": "author", "type": "STRING", "searchability": "ENABLED"},
            {"name": "projects", "type": "STRING_LIST", "searchability": "ENABLED"},
        ],
    )
    save("update", resp)
    kb = wait_kb(s["knowledgeBaseId"])
    save("kb-after-update", kb)
    print(json.dumps(kb.get("metadataConfigurations"), ensure_ascii=False))

def cleanup():
    s = state()
    agent.delete_knowledge_base(knowledgeBaseId=s["knowledgeBaseId"])
    for _ in range(180):
        try:
            agent.get_knowledge_base(knowledgeBaseId=s["knowledgeBaseId"])
        except agent.exceptions.ResourceNotFoundException:
            break
        time.sleep(5)
    iam.delete_role(RoleName=s["name"])
    STATE.unlink()
    print("Knowledge Base and IAM role deleted")

if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "search":
        search(sys.argv[2] if len(sys.argv) > 2 else "search")
    else:
        {"setup": setup, "ingest": ingest, "enable": enable, "cleanup": cleanup}[cmd]()

# REPORTS の各報告に部署を追加（report-01から順に人事部・経理部・法務部・情報システム部）
"department": "経理部",

# ingest() の inlineAttributes に追加
{"key": "department", "value": {"type": "STRING", "stringValue": r["department"]}},

# 検索するクエリ
QUERIES = [
    "経理部の月次報告",
    "お金の出入りを管理する部署の月次報告",
    "契約書をチェックする部署の報告",
    "琥珀の進捗",
    "Hanako Yamada report",
]

# enable() の metadataConfigurations に追加
{"name": "department", "type": "STRING", "searchability": "ENABLED"},

# search() の表示に部署を追加
print(
    f"  [{i}] score={r.get('score', 0):.4f} author={m.get('author')} dept={m.get('department')} projects={m.get('projects')}"
)

uv run demo.py setup
uv run demo.py ingest
uv run demo.py search before
uv run demo.py enable
uv run demo.py search after
uv run demo.py cleanup

Q: 経理部の月次報告
  [3] score=0.1592 author=佐藤健 dept=経理部 projects=['ヒスイ']
Q: お金の出入りを管理する部署の月次報告
  [3] score=0.1820 author=佐藤健 dept=経理部 projects=['ヒスイ']
Q: 契約書をチェックする部署の報告
  [3] score=0.1780 author=鈴木一郎 dept=法務部 projects=['メノウ']

Q: 経理部の月次報告
  [1] score=0.4642 author=佐藤健 dept=経理部 projects=['ヒスイ']
  [2] score=0.2342 author=高橋美咲 dept=情報システム部 projects=['サンゴ']
Q: お金の出入りを管理する部署の月次報告
  [1] score=0.3474 author=佐藤健 dept=経理部 projects=['ヒスイ']
  [2] score=0.3278 author=高橋美咲 dept=情報システム部 projects=['サンゴ']
Q: 契約書をチェックする部署の報告
  [3] score=0.1749 author=鈴木一郎 dept=法務部 projects=['メノウ']
Q: 琥珀の進捗
  [1] score=0.3938 author=高橋美咲 dept=情報システム部 projects=['サンゴ']
  [2] score=0.2020 author=佐藤健 dept=経理部 projects=['ヒスイ']
  [3] score=0.1685 author=鈴木一郎 dept=法務部 projects=['メノウ']
Q: Hanako Yamada report
  [1] score=0.3558 author=高橋美咲 dept=情報システム部 projects=['サンゴ']
  [2] score=0.2760 author=鈴木一郎 dept=法務部 projects=['メノウ']
  [3] score=0.1510 author=佐藤健 dept=経理部 projects=['ヒスイ']