Add support for cross-account agents (#1053)

Author: toshikwa

docs/en/DEPLOY_OPTION.md

@@ -1768,7 +1768,7 @@ const envs: Record<string, Partial<StackInput>> = {
 ## Using Bedrock from a Different AWS Account
 
 > [!NOTE]
-> Agent-related tasks (Agent, Flow, Prompt Optimization Tool) do not support using a different AWS account and may result in errors during execution.
+> Flow Chat use case and Prompt Optimization Tool do not support using a different AWS account and may result in errors during execution.
 
 You can use Bedrock from a different AWS account. As a prerequisite, the initial deployment of GenU must be completed.
 
@@ -1827,11 +1827,13 @@ For details on how to specify Principals, refer to: [AWS JSON Policy Elements: P
       "Sid": "AllowBedrockInvokeModel",
       "Effect": "Allow",
       "Action": [
-        "bedrock:InvokeModel*",
+        "bedrock:Invoke*",
         "bedrock:Rerank",
         "bedrock:GetInferenceProfile",
         "bedrock:GetAsyncInvoke",
-        "bedrock:ListAsyncInvokes"
+        "bedrock:ListAsyncInvokes",
+        "bedrock:GetAgent*",
+        "bedrock:ListAgent*"
       ],
       "Resource": ["*"]
     },
@@ -1871,6 +1873,13 @@ When using Knowledge Base, you'll need to include these additional parameters:
 - `ragKnowledgeBaseId` ... Knowledge Base ID created in advance in the different account
   - Knowledge Base must exist in the `modelRegion`
 
+When using Agent Chat use case, you'll need to include these additional parameters:
+
+- `agents` ... a list of Bedrock Agent configurations, which has following properties:
+  - `displayName` ... Display name of the agent
+  - `agentId` ... Agent ID created in advance in the different account
+  - `aliasId` ... Agent Alias ID created in advance in the different account
+
 **Edit [parameter.ts](/packages/cdk/parameter.ts)**
 
 ```typescript
@@ -1879,8 +1888,17 @@ const envs: Record<string, Partial<StackInput>> = {
   dev: {
     crossAccountBedrockRoleArn:
       'arn:aws:iam::AccountID:role/PreCreatedRoleName',
-    ragKnowledgeBaseEnabled: true, // Only when using Knowledge Base
-    ragKnowledgeBaseId: 'XXXXXXXXXX', // Only when using Knowledge Base
+    // Only when using Knowledge Base
+    ragKnowledgeBaseEnabled: true,
+    ragKnowledgeBaseId: 'YOUR_KNOWLEDGE_BASE_ID',
+    // Only when using agents
+    agents: [
+      {
+        displayName: 'YOUR AGENT NAME',
+        agentId: 'YOUR_AGENT_ID',
+        aliasId: 'YOUR_AGENT_ALIAS_ID',
+      },
+    ],
   },
 };
 ```
@@ -1892,8 +1910,17 @@ const envs: Record<string, Partial<StackInput>> = {
 {
   "context": {
     "crossAccountBedrockRoleArn": "arn:aws:iam::AccountID:role/PreCreatedRoleName",
-    "ragKnowledgeBaseEnabled": true, // Only when using Knowledge Base
-    "ragKnowledgeBaseId": "XXXXXXXXXX" // Only when using Knowledge Base
+    // Only when using Knowledge Base
+    "ragKnowledgeBaseEnabled": true,
+    "ragKnowledgeBaseId": "YOUR_KNOWLEDGE_BASE_ID",
+    // Only when using agents
+    "agents": [
+      {
+        "displayName": "YOUR AGENT NAME",
+        "agentId": "YOUR_AGENT_ID",
+        "aliasId": "YOUR_AGENT_ALIAS_ID"
+      }
+    ]
   }
 }
 ```

docs/ja/DEPLOY_OPTION.md

@@ -1775,7 +1775,7 @@ const envs: Record<string, Partial<StackInput>> = {
 ## 別 AWS アカウントの Bedrock を利用したい場合
 
 > [!NOTE]
-> Agent 系のタスク (Agent, Flow, プロンプト最適化ツール) に関しては別 AWS アカウントの利用をサポートしていないため、実行時にエラーになる可能性があります。
+> 「Flow チャットユースケース」および「プロンプト最適化ツール」は別 AWS アカウントの利用をサポートしていないため、実行時にエラーになる可能性があります。
 
 別 AWS アカウントの Bedrock を利用することができます。前提条件として、GenU の初回デプロイは完了済みとします。
 
@@ -1834,11 +1834,13 @@ Principal の指定方法について詳細を確認したい場合はこちら
       "Sid": "AllowBedrockInvokeModel",
       "Effect": "Allow",
       "Action": [
-        "bedrock:InvokeModel*",
+        "bedrock:Invoke*",
         "bedrock:Rerank",
         "bedrock:GetInferenceProfile",
         "bedrock:GetAsyncInvoke",
-        "bedrock:ListAsyncInvokes"
+        "bedrock:ListAsyncInvokes",
+        "bedrock:GetAgent*",
+        "bedrock:ListAgent*"
       ],
       "Resource": ["*"]
     },
@@ -1878,6 +1880,13 @@ Knowledge Base を利用する場合は、下記パラメーターも指定し
 - `ragKnowledgeBaseId` ... 別アカウントに事前構築した Knowledge Base の ID です
   - Knowledge Base は `modelRegion` に存在する必要があります
 
+Agent Chat ユースケースを使用する場合、下記パラメーターも指定します。
+
+- `agents` ... 以下の属性を持つ Bedrock Agent の設定のリストです
+  - `displayName` ... エージェントの表示名
+  - `agentId` ... 別アカウントに事前構築したエージェントの ID
+  - `aliasId` ... 別アカウントに事前構築したエージェントのエイリアス ID
+
 **[parameter.ts](/packages/cdk/parameter.ts) を編集**
 
 ```typescript
@@ -1886,8 +1895,17 @@ const envs: Record<string, Partial<StackInput>> = {
   dev: {
     crossAccountBedrockRoleArn:
       'arn:aws:iam::アカウントID:role/事前に作成したロール名',
-    ragKnowledgeBaseEnabled: true, // Knowledge Base を利用する場合のみ
-    ragKnowledgeBaseId: 'XXXXXXXXXX', // Knowledge Base を利用する場合のみ
+    // Knowledge Base を利用する場合のみ
+    ragKnowledgeBaseEnabled: true,
+    ragKnowledgeBaseId: 'YOUR_KNOWLEDGE_BASE_ID',
+    // Bedrock エージェントを利用する場合のみ
+    agents: [
+      {
+        displayName: 'YOUR AGENT NAME',
+        agentId: 'YOUR_AGENT_ID',
+        aliasId: 'YOUR_AGENT_ALIAS_ID',
+      },
+    ],
   },
 };
 ```
@@ -1899,8 +1917,17 @@ const envs: Record<string, Partial<StackInput>> = {
 {
   "context": {
     "crossAccountBedrockRoleArn": "arn:aws:iam::アカウントID:role/事前に作成したロール名",
-    "ragKnowledgeBaseEnabled": true, // Knowledge Base を利用する場合のみ
-    "ragKnowledgeBaseId": "XXXXXXXXXX" // Knowledge Base を利用する場合のみ
+    // Knowledge Base を利用する場合のみ
+    "ragKnowledgeBaseEnabled": true,
+    "ragKnowledgeBaseId": "YOUR_KNOWLEDGE_BASE_ID",
+    // Bedrock エージェントを利用する場合のみ
+    "agents": [
+      {
+        "displayName": "YOUR AGENT NAME",
+        "agentId": "YOUR_AGENT_ID",
+        "aliasId": "YOUR_AGENT_ALIAS_ID"
+      }
+    ]
   }
 }
 ```

packages/cdk/lambda/utils/bedrockAgentApi.ts

@@ -1,10 +1,8 @@
 import {
-  BedrockAgentClient,
   GetAgentAliasCommand,
   ListAgentActionGroupsCommand,
 } from '@aws-sdk/client-bedrock-agent';
 import {
-  BedrockAgentRuntimeClient,
   DependencyFailedException,
   InvokeAgentCommand,
   Parameter,
@@ -22,13 +20,12 @@ import {
   BraveSearchResult,
 } from 'generative-ai-use-cases';
 import { streamingChunk } from './streamingChunk';
+import {
+  initBedrockAgentClient,
+  initBedrockAgentRuntimeClient,
+} from './bedrockClient';
 
-const agentClient = new BedrockAgentClient({
-  region: process.env.MODEL_REGION,
-});
-const agentRuntimeClient = new BedrockAgentRuntimeClient({
-  region: process.env.MODEL_REGION,
-});
+const MODEL_REGION = process.env.MODEL_REGION as string;
 const s3Client = new S3Client({});
 
 // Agent information
@@ -64,8 +61,11 @@ const encodeUrlString = (str: string): string => {
 const getAgentInfo = async (agentId: string, agentAliasId: string) => {
   // Get Agent Info if not cached
   if (!agentInfoMap[agentAliasId]) {
+    const bedrockAgentClient = await initBedrockAgentClient({
+      region: MODEL_REGION,
+    });
     // Get Agent Version
-    const agentAliasInfoRes = await agentClient.send(
+    const agentAliasInfoRes = await bedrockAgentClient.send(
       new GetAgentAliasCommand({
         agentId: agentId,
         agentAliasId: agentAliasId,
@@ -75,7 +75,7 @@ const getAgentInfo = async (agentId: string, agentAliasId: string) => {
       agentAliasInfoRes.agentAlias?.routingConfiguration?.pop()?.agentVersion ??
       '1';
     // List Action Group
-    const actionGroups = await agentClient.send(
+    const actionGroups = await bedrockAgentClient.send(
       new ListAgentActionGroupsCommand({
         agentId: agentId,
         agentVersion: agentVersion,
@@ -129,7 +129,11 @@ const bedrockAgentApi: ApiInterface = {
         enableTrace: true,
         inputText: messages[messages.length - 1].content,
       });
-      const res = await agentRuntimeClient.send(command);
+
+      const bedrockAgentRuntimeClient = await initBedrockAgentRuntimeClient({
+        region: MODEL_REGION,
+      });
+      const res = await bedrockAgentRuntimeClient.send(command);
 
       if (!res.completion) {
         return;

packages/cdk/lambda/utils/bedrockClient.ts

@@ -8,13 +8,18 @@ import {
   BedrockAgentRuntimeClientConfig,
 } from '@aws-sdk/client-bedrock-agent-runtime';
 import { S3Client, S3ClientConfig } from '@aws-sdk/client-s3';
+import {
+  BedrockAgentClient,
+  BedrockAgentClientConfig,
+} from '@aws-sdk/client-bedrock-agent';
 
 // Temporary credentials for cross-account access
 const stsClient = new STSClient();
 let temporaryCredentials: Credentials | undefined;
 
 // Store Bedrock clients
 let bedrockRuntimeClient: BedrockRuntimeClient | undefined;
+let bedrockAgentClient: BedrockAgentClient | undefined;
 let bedrockAgentRuntimeClient: BedrockAgentRuntimeClient | undefined;
 let knowledgeBaseS3Client: S3Client | undefined;
 
@@ -91,6 +96,25 @@ export const initBedrockRuntimeClient = async (
   return bedrockRuntimeClient;
 };
 
+export const initBedrockAgentClient = async (
+  config: BedrockAgentClientConfig & { region: string }
+) => {
+  // Use cross-account role
+  if (process.env.CROSS_ACCOUNT_BEDROCK_ROLE_ARN) {
+    return new BedrockAgentClient({
+      ...(await getCrossAccountCredentials(
+        process.env.CROSS_ACCOUNT_BEDROCK_ROLE_ARN
+      )),
+      ...config,
+    });
+  }
+  // Use Lambda execution role
+  if (!bedrockAgentClient) {
+    bedrockAgentClient = new BedrockAgentClient(config);
+  }
+  return bedrockAgentClient;
+};
+
 export const initBedrockAgentRuntimeClient = async (
   config: BedrockAgentRuntimeClientConfig & { region: string }
 ) => {

packages/cdk/lib/create-stacks.ts

@@ -52,6 +52,13 @@ export const createStacks = (app: cdk.App, params: ProcessedStackInput) => {
       : null;
 
   // Agent
+  if (params.crossAccountBedrockRoleArn) {
+    if (params.agentEnabled || params.searchApiKey) {
+      throw new Error(
+        'When `crossAccountBedrockRoleArn` is specified, the `agentEnabled` and `searchApiKey` parameters are not supported. Please create agents in the other account and specify them in the `agents` parameter.'
+      );
+    }
+  }
   const agentStack = params.agentEnabled
     ? new AgentStack(app, `WebSearchAgentStack${params.env}`, {
         env: {