[アップデート] Amazon Bedrock Managed Knowledge Baseでドキュメントのアクセス権を調べるAPIが追加されたので試してみた

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)

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 setup():
    if STATE.exists():
        raise RuntimeError("state.json already exists")
    account = boto3.client("sts").get_caller_identity()["Account"]
    name = "kb-acl-debug-" + 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))[
        "Role"
    ]
    state = {"roleName": name}
    STATE.write_text(json.dumps(state))
    time.sleep(15)
    kb = agent.create_knowledge_base(
        name=name,
        roleArn=role["Arn"],
        knowledgeBaseConfiguration={
            "type": "MANAGED",
            "managedKnowledgeBaseConfiguration": {"embeddingModelType": "MANAGED"},
        },
    )["knowledgeBase"]
    state["knowledgeBaseId"] = kb["knowledgeBaseId"]
    STATE.write_text(json.dumps(state))
    save("created-kb", kb)
    while (
        agent.get_knowledge_base(knowledgeBaseId=state["knowledgeBaseId"])[
            "knowledgeBase"
        ]["status"]
        != "ACTIVE"
    ):
        time.sleep(10)
    ds = agent.create_data_source(
        knowledgeBaseId=state["knowledgeBaseId"],
        name="custom-acl",
        dataSourceConfiguration={
            "type": "MANAGED_KNOWLEDGE_BASE_CONNECTOR",
            "managedKnowledgeBaseConnectorConfiguration": {
                "connectorParameters": {
                    "type": "CUSTOM",
                    "version": "1",
                    "aclEnabled": True,
                }
            },
        },
    )["dataSource"]
    state["dataSourceId"] = ds["dataSourceId"]
    STATE.write_text(json.dumps(state))
    save("created-ds", ds)
    print(json.dumps(state))

def ingest():
    state = json.loads(STATE.read_text())
    acl = [
        {"name": u + "@example.com", "type": "USER", "access": a}
        for u, a in [
            ("tanaka", "ALLOW"),
            ("sato", "DENY"),
            ("suzuki", "ALLOW"),
            ("suzuki", "DENY"),
        ]
    ]
    document = {
        "content": {
            "dataSourceType": "CUSTOM",
            "custom": {
                "customDocumentIdentifier": {"id": "vpn-guide"},
                "sourceType": "IN_LINE",
                "inlineContent": {
                    "type": "TEXT",
                    "textContent": {
                        "data": "ハナミズキ社のVPN接続手順。VPNエラーHANA-042が発生したら、社内IT窓口の内線8420へ連絡してください。受付時間は平日9時から17時です。"
                    },
                },
            },
        },
        "metadata": {
            "type": "IN_LINE_ATTRIBUTE",
            "inlineAttributes": [
                {"key": "department", "value": {"type": "STRING", "stringValue": "IT"}}
            ],
            "accessControlList": acl,
        },
    }
    params = {k: state[k] for k in ["knowledgeBaseId", "dataSourceId"]}
    save("ingest-input", document)
    result = agent.ingest_knowledge_base_documents(**params, documents=[document])
    save("ingest", result)
    print(json.dumps(result, default=str))

def check():
    state = json.loads(STATE.read_text())
    params = {k: state[k] for k in ["knowledgeBaseId", "dataSourceId"]}
    for _ in range(60):
        status = agent.get_knowledge_base_documents(
            **params,
            documentIdentifiers=[
                {"dataSourceType": "CUSTOM", "custom": {"id": "vpn-guide"}}
            ],
        )
        save("document-status", status)
        current = status["documentDetails"][0]["status"]
        print("status", current, flush=True)
        if current == "INDEXED":
            break
        if current in (
            "FAILED",
            "IGNORED",
            "PARTIALLY_INDEXED",
            "METADATA_UPDATE_FAILED",
        ):
            raise RuntimeError(status["documentDetails"])
        time.sleep(10)
    else:
        raise RuntimeError("Document indexing timed out")
    params["documentId"] = "vpn-guide"
    acl = runtime.get_ingested_document_acl(**params)
    save("acl", acl)
    print("acl", json.dumps(acl["documentAcl"]))
    for user in ["tanaka", "sato", "suzuki", "takahashi"]:
        context = {"userId": user + "@example.com"}
        result = runtime.check_ingested_document_acl(**params, userContext=context)
        save("check-" + user, result)
        retrieval = runtime.retrieve(
            knowledgeBaseId=state["knowledgeBaseId"],
            retrievalQuery={"text": "VPNエラーHANA-042の問い合わせ先は？"},
            userContext=context,
        )
        save("retrieve-" + user, retrieval)
        print(
            user,
            "hasAccess=",
            result["hasAccess"],
            "results=",
            len(retrieval["retrievalResults"]),
        )

def cleanup():
    state = json.loads(STATE.read_text())
    if "knowledgeBaseId" in state:
        agent.delete_knowledge_base(knowledgeBaseId=state["knowledgeBaseId"])
        for _ in range(180):
            try:
                agent.get_knowledge_base(knowledgeBaseId=state["knowledgeBaseId"])
            except agent.exceptions.ResourceNotFoundException:
                break
            time.sleep(5)
        else:
            raise RuntimeError("Knowledge Base deletion timed out")
    iam.delete_role(RoleName=state["roleName"])
    save("cleanup", {"knowledgeBaseDeleted": True, "roleDeleted": True})
    STATE.unlink()
    print("Knowledge Base and IAM role deleted")

if __name__ == "__main__":
    {"setup": setup, "ingest": ingest, "check": check, "cleanup": cleanup}[
        sys.argv[1]
    ]()

status INDEXED
acl {"allowList": {"conditions": [{"conditionOperator": "OR", "users": [{"id": "tanaka@example.com", "type": "KNOWLEDGE_BASE"}]}], "memberRelation": "AND"}, "denyList": {"conditions": [{"conditionOperator": "OR", "users": [{"id": "sato@example.com", "type": "KNOWLEDGE_BASE"}, {"id": "suzuki@example.com", "type": "KNOWLEDGE_BASE"}]}], "memberRelation": "AND"}}
tanaka hasAccess= True results= 1
sato hasAccess= False results= 0
suzuki hasAccess= False results= 0
takahashi hasAccess= False results= 0