エクサウィザーズ Tech Blog
🏢

AgentCore Memoryにおけるテナント分離

伊賀 裕展
に公開
AI
AWS
SaaS
memory
AgentCore
tech

先端技術開発グループ（WAND）の伊賀です。
本記事はエクサウィザーズの Advent Calendar 22 日目の記事です。

本記事では Amazon Bedrock AgentCore Memory のテナント分離に関して説明します。

Amazon Bedrock AgentCore Memory

Amazon Bedrock AgentCore Memory（以下、AgentCore Memory）は、AI エージェント向けのメモリ管理を提供するフルマネージドサービスです。AgentCore Memory は、ユーザーとエージェントとの対話を記録し、過去の会話内容を活用した応答生成を可能にします。

AgentCore Memory は短期記憶と長期記憶の両方をサポートしています。
短期記憶では同一のセッション内でのやりとりやコンテキストのデータを保存します。長期記憶ではセッションを超えた情報を保存し、エージェントが過去の対話履歴や学習内容を参照できるようにします。なお、短期記憶から長期記憶への変換は設定している記憶戦略に基づいて非同期に実行されます。
記憶戦略は複数あり、例えば会話の要約を作成するSummaryStrategyやユーザーの嗜好、選択、スタイルを自動的に識別・抽出するUserPreferenceStrategyがあります。

https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/memory.html

データパーティショニング

AgentCore Memory では短期記憶と長期記憶でそれぞれ異なった記憶領域が使用されます。
短期記憶では actorIdとsessionIdの組み合わせでデータがパーティショニングされます。
長期記憶ではnamespaceによって分離され、namespace には長期記憶の保存のされ方を決めるstrategyの識別子であるstrategyIdやactorIdおよびsessionIdを使用することができます。
長期記憶のnamespaceは / 区切りで設定することが推奨されており、例えば以下のように設定することができます。

namespace = strategy/{strategyId}/actor/actorId

マルチテナント環境においては、actorIdにテナントを識別する情報を含めることで、テナントごとにデータを分離することが可能です。

テナント分離

actorIdにテナント識別子を含めることで、AgentCore Memory に保存されるデータをテナントごとに分離することができますが、データを分離したとしても、テナント間のデータアクセスを防止できているわけではありません。
例えば、あるテナントのエージェントが別のテナントのactorIdを指定してメモリデータにアクセスすることが可能です。

そのため、マルチテナント環境で AgentCore Memory を使用する場合は、テナント間のデータアクセスを防止する仕組みを実装する必要があります。例えば、IAM Role を定義して、STS のsessionTagでテナントコンテキストを指定したassumeRoleで AI エージェントに IAM Role を引き受けさせ、AgentCore Memory へのアクセスを制御することが考えられます。
全体のアーキテクチャは以下のようになります。

以下ではこの考え方に則り、具体的な実装例を示します。

IAM Role の Permission Policy の作成

AI エージェントに引き受けさせる IAM Role に対して、AgentCore Memory へのアクセスを制御する Permission Policy を作成します。以下は、特定のactorIdにのみアクセスを許可するポリシーの例です。STS のセッションタグを使用して、引き受けたロールにactorIdタグを付与し、そのタグに基づいてアクセス制御を行います。

Short Term Memory Permission ポリシー

Short Term Memory では Condition にbedrock-agentcore:actorIdを使用することができます。

{
  "Sid": "ShortTermMemoryPolicy",
  "Effect": "Allow",
  "Action": [
    "bedrock-agentcore:CreateEvent",
    "bedrock-agentcore:GetEvent"
    // ... other short term memory actions ...
  ],
  "Resource": [
    "arn:aws:bedrock-agentcore:{REGION}:{account_id}:memory/{memory_id}"
  ],
  "Condition": {
    "StringEquals": {
      "bedrock-agentcore:actorId": "${aws:PrincipalTag/actorId}"
    }
  }
}

Long Term Memory Permission ポリシー

Long Term Memory では Condition にbedrock-agentcore:namespaceを使用することができます。

{
  "Sid": "LongTermMemoryPolicy",
  "Effect": "Allow",
  "Action": [
    "bedrock-agentcore:GetMemoryRecord",
    "bedrock-agentcore:ListMemoryRecords"
    // ... other long term memory actions ...
  ],
  "Resource": [
    "arn:aws:bedrock-agentcore:{REGION}:{account_id}:memory/{memory_id}"
  ],
  "Condition": {
    "StringEquals": {
      "bedrock-agentcore:namespace": "/strategies/{strategy_id}/actor/${aws:PrincipalTag/actorId}"
    }
  }
}

IAM Role の Trust Policy の作成

IAM Role の Trust Policy では、STS のassumeRoleでセッションタグを付与できるように設定します。また、AgentCore Runtime からのみ引き受け可能にするために、Principal に AgentCore Runtime のロールを指定します。

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AllowAssumeFromAgentRuntime",
      "Effect": "Allow",
      "Principal": {
        "AWS": "arn:aws:iam::{ACCOUNT_ID}:role/AgentRuntimeRole"
      },
      "Action": "sts:AssumeRole"
    },
    {
      "Sid": "AllowTagSessionOnlyActorId",
      "Effect": "Allow",
      "Principal": {
        "AWS": "arn:aws:iam::{ACCOUNT_ID}:role/AgentRuntimeRole"
      },
      "Action": "sts:TagSession",
      "Condition": {
        "ForAllValues:StringEquals": {
          "aws:TagKeys": ["actorId"]
        }
      }
    }
  ]
}

AI エージェントでの assumeRole の実装

AI エージェントのコード内で、STS のassumeRoleを呼び出し、テナントごとに異なる IAM Role を引き受けます。引き受ける際に、セッションタグとしてactorIdを指定します。

def assume_role_for_tenant(role_arn, actor_id):
    sts_client = boto3.client('sts')
    response = sts_client.assume_role(
        RoleArn=role_arn,
        RoleSessionName='AgentCoreMemorySession',
        DurationSeconds=900,
        Tags=[{ 'Key': 'actorId', 'Value': actor_id }]
    )
    credentials = response['Credentials']
    return credentials

AgentCore Memory へのアクセス

引き受けた IAM Role の一時的な認証情報を使用して、AgentCore Memory にアクセスします。以下は、引き受けた認証情報を使用して短期記憶と長期記憶にアクセスする例です。

Short Term Memory へのアクセス例は下記となります。

def access_short_term_memory(memory_id, actor_id, session_id, credentials):
    agentcore_client = boto3.client(
        'bedrock-agentcore',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
    )
    response = agentcore_client.list_events(
        memoryId=memory_id,
        actorId=actor_id,
        sessionId=session_id,
        maxResults=10
    )
    return response

Long Term Memory へのアクセス例は下記となります。

def access_long_term_memory(memory_id, namespace, credentials):
    agentcore_client = boto3.client(
        'bedrock-agentcore',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
    )
    response = agentcore_client.list_memory_records(
        memoryId=memory_id,
        namespace=namespace,
        maxResults=10
    )
    return response

まとめ

本記事では、Amazon Bedrock AgentCore Memory におけるテナント分離の考え方と実装例を紹介しました。AgentCore Memory は強力なメモリ管理機能を提供しますが、マルチテナント環境でのデータ分離とアクセス制御はアプリケーション側で適切に設計・実装する必要があります。
適切な IAM Role とポリシーを使用して、テナント間のデータアクセスを防止することが重要です。

エクサウィザーズ Tech Blog
エクサウィザーズ Tech Blog

株式会社エクサウィザーズのテックブログ

Discussion