別アカウントの Bedrock を利用するためのクロスアカウント設定を追加

README.md

@@ -154,6 +154,7 @@ npm run cdk:deploy
   - [SAML 認証](/docs/DEPLOY_OPTION.md#SAML-認証)
 - [モニタリング用のダッシュボードの有効化](/docs/DEPLOY_OPTION.md#モニタリング用のダッシュボードの有効化)
 - [ファイルアップロード機能の有効化](/docs/DEPLOY_OPTION.md#ファイルアップロード機能の有効化)
+- [別 AWS アカウントの Bedrock を利用したい場合](/docs/DEPLOY_OPTION.md#別-AWS-アカウントの-Bedrock-を利用したい場合)
 
 ## その他
  - [アップデート方法](/docs/UPDATE.md)

docs/DEPLOY_OPTION.md

@@ -433,4 +433,66 @@ cdk.json には以下の値を設定します。
     "hostedZoneId": "XXXXXXXXXXXXXXXXXXXX"
   }
 }
-```
+```
+
+## 別 AWS アカウントの Bedrock を利用したい場合
+
+別 AWS アカウントの Bedrock を利用することができます。前提条件として、本サンプルアプリケーションによるデプロイは完了済みとします。
+
+別 AWS アカウントの Bedrock を利用するためには、別 AWS アカウントに IAM ロールを 1 つ作成する必要があります。作成する IAM ロール名は任意ですが、本サンプルアプリケーションでデプロイされた以下の名前で始まる IAM ロール名を、別アカウントで作成した IAM ロールの Principal に指定します。
+
+- `GenerativeAiUseCasesStack-APIPredictTitleService`
+- `GenerativeAiUseCasesStack-APIPredictService`
+- `GenerativeAiUseCasesStack-APIPredictStreamService`
+- `GenerativeAiUseCasesStack-APIGenerateImageService`
+
+Principal の指定方法について詳細を確認したい場合はこちらを参照ください: [AWS JSON ポリシーの要素: Principal](https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/reference_policies_elements_principal.html)
+
+Principal 設定例 (別アカウントにて設定)
+
+```json
+{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Principal": {
+                "AWS": [
+                    "arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictTitleServiceXXX-XXXXXXXXXXXX",
+                    "arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictServiceXXXXXXXX-XXXXXXXXXXXX",
+                    "arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictStreamServiceXX-XXXXXXXXXXXX",
+                    "arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateImageServiceXX-XXXXXXXXXXXX"
+                ]
+            },
+            "Action": "sts:AssumeRole",
+            "Condition": {}
+        }
+    ]
+}
+```
+
+cdk.json には以下の値を設定します。
+
+- `crossAccountBedrockRoleArn` ... 別アカウントで事前に作成した IAM ロールの ARN です
+
+**[packages/cdk/cdk.json](/packages/cdk/cdk.json) を編集**
+
+```json
+{
+  "context": {
+    "crossAccountBedrockRoleArn": "arn:aws:iam::アカウントID:role/事前に作成したロール名"
+  }
+}
+```
+
+cdk.json 設定例
+
+```json
+{
+  "context": {
+    "crossAccountBedrockRoleArn": "arn:aws:iam::222222222222:role/YYYYYYYYYYYYYYYYYYYYY"
+  }
+}
+```
+
+設定変更後に `npm run cdk:deploy` を実行して変更内容を反映させます。

packages/cdk/cdk.json

@@ -48,6 +48,7 @@
     "anonymousUsageTracking": true,
     "recognizeFileEnabled": false,
     "vpcId": null,
+    "crossAccountBedrockRoleArn": "",
     "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
     "@aws-cdk/core:checkSecretUsage": true,
     "@aws-cdk/core:target-partitions": [

packages/cdk/lambda/utils/bedrockApi.ts

@@ -13,10 +13,69 @@ import {
   UnrecordedMessage,
 } from 'generative-ai-use-cases-jp';
 import { BEDROCK_MODELS, BEDROCK_IMAGE_GEN_MODELS } from './models';
+import { STSClient, AssumeRoleCommand } from '@aws-sdk/client-sts';
 
-const client = new BedrockRuntimeClient({
-  region: process.env.MODEL_REGION,
-});
+// STSから一時的な認証情報を取得する関数
+const assumeRole = async (crossAccountBedrockRoleArn: string) => {
+  const stsClient = new STSClient({ region: process.env.MODEL_REGION });
+  const command = new AssumeRoleCommand({
+    RoleArn: crossAccountBedrockRoleArn,
+    RoleSessionName: 'BedrockApiAccess',
+  });
+
+  try {
+    const response = await stsClient.send(command);
+    if (response.Credentials) {
+      return {
+        accessKeyId: response.Credentials?.AccessKeyId,
+        secretAccessKey: response.Credentials?.SecretAccessKey,
+        sessionToken: response.Credentials?.SessionToken,
+      };
+    } else {
+      throw new Error('認証情報を取得できませんでした。');
+    }
+  } catch (error) {
+    console.error('Error assuming role: ', error);
+    throw error;
+  }
+};
+
+// BedrockRuntimeClient を初期化するこの関数は、通常では単純に BedrockRuntimeClient を環境変数で指定されたリージョンで初期化します。
+// 特別なケースとして、異なる AWS アカウントに存在する Bedrock リソースを利用したい場合があります。
+// そのような場合、CROSS_ACCOUNT_BEDROCK_ROLE_ARN 環境変数が設定されているかをチェックします。(cdk.json で crossAccountBedrockRoleArn が設定されている場合に環境変数として設定される)
+// 設定されている場合、指定されたロールを AssumeRole 操作によって引き受け、取得した一時的な認証情報を用いて BedrockRuntimeClient を初期化します。
+// これにより、別の AWS アカウントの Bedrock リソースへのアクセスが可能になります。
+const initBedrockClient = async () => {
+  // CROSS_ACCOUNT_BEDROCK_ROLE_ARN が設定されているかチェック
+  if (process.env.CROSS_ACCOUNT_BEDROCK_ROLE_ARN) {
+    // STS から一時的な認証情報を取得してクライアントを初期化
+    const tempCredentials = await assumeRole(
+      process.env.CROSS_ACCOUNT_BEDROCK_ROLE_ARN
+    );
+
+    if (
+      !tempCredentials.accessKeyId ||
+      !tempCredentials.secretAccessKey ||
+      !tempCredentials.sessionToken
+    ) {
+      throw new Error('STSからの認証情報が不完全です。');
+    }
+
+    return new BedrockRuntimeClient({
+      region: process.env.MODEL_REGION,
+      credentials: {
+        accessKeyId: tempCredentials.accessKeyId,
+        secretAccessKey: tempCredentials.secretAccessKey,
+        sessionToken: tempCredentials.sessionToken,
+      },
+    });
+  } else {
+    // STSを使用しない場合のクライアント初期化
+    return new BedrockRuntimeClient({
+      region: process.env.MODEL_REGION,
+    });
+  }
+};
 
 const createBodyText = (
   model: string,
@@ -49,6 +108,7 @@ const extractOutputImage = (
 
 const bedrockApi: ApiInterface = {
   invoke: async (model, messages) => {
+    const client = await initBedrockClient();
     const command = new InvokeModelCommand({
       modelId: model.modelId,
       body: createBodyText(model.modelId, messages),
@@ -59,6 +119,7 @@ const bedrockApi: ApiInterface = {
     return extractOutputText(model.modelId, body);
   },
   invokeStream: async function* (model, messages) {
+    const client = await initBedrockClient();
     try {
       const command = new InvokeModelWithResponseStreamCommand({
         modelId: model.modelId,
@@ -98,6 +159,7 @@ const bedrockApi: ApiInterface = {
     }
   },
   generateImage: async (model, params) => {
+    const client = await initBedrockClient();
     const command = new InvokeModelCommand({
       modelId: model.modelId,
       body: createBodyImage(model.modelId, params),

packages/cdk/lib/construct/api.ts

@@ -31,6 +31,7 @@ export class Api extends Construct {
   readonly imageGenerationModelIds: string[];
   readonly endpointNames: string[];
   readonly agentNames: string[];
+  readonly crossAccountBedrockRoleArn: string;
 
   constructor(scope: Construct, id: string, props: BackendApiProps) {
     super(scope, id);
@@ -93,6 +94,11 @@ export class Api extends Construct {
       };
     }
 
+    // cross account access IAM role
+    const crossAccountBedrockRoleArn = this.node.tryGetContext(
+      'crossAccountBedrockRoleArn'
+    );
+
     // Lambda
     const predictFunction = new NodejsFunction(this, 'Predict', {
       runtime: Runtime.NODEJS_18_X,
@@ -102,6 +108,7 @@ export class Api extends Construct {
         MODEL_REGION: modelRegion,
         MODEL_IDS: JSON.stringify(modelIds),
         IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),
+        CROSS_ACCOUNT_BEDROCK_ROLE_ARN: crossAccountBedrockRoleArn,
       },
       bundling: {
         nodeModules: ['@aws-sdk/client-bedrock-runtime'],
@@ -118,6 +125,7 @@ export class Api extends Construct {
         IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),
         AGENT_REGION: agentRegion,
         AGENT_MAP: JSON.stringify(agentMap),
+        CROSS_ACCOUNT_BEDROCK_ROLE_ARN: crossAccountBedrockRoleArn,
       },
       bundling: {
         nodeModules: [
@@ -144,6 +152,7 @@ export class Api extends Construct {
         MODEL_REGION: modelRegion,
         MODEL_IDS: JSON.stringify(modelIds),
         IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),
+        CROSS_ACCOUNT_BEDROCK_ROLE_ARN: crossAccountBedrockRoleArn,
       },
     });
     table.grantWriteData(predictTitleFunction);
@@ -156,6 +165,7 @@ export class Api extends Construct {
         MODEL_REGION: modelRegion,
         MODEL_IDS: JSON.stringify(modelIds),
         IMAGE_GENERATION_MODEL_IDS: JSON.stringify(imageGenerationModelIds),
+        CROSS_ACCOUNT_BEDROCK_ROLE_ARN: crossAccountBedrockRoleArn,
       },
       bundling: {
         nodeModules: ['@aws-sdk/client-bedrock-runtime'],
@@ -183,15 +193,40 @@ export class Api extends Construct {
 
     // Bedrock は常に権限付与
     // Bedrock Policy
-    const bedrockPolicy = new PolicyStatement({
-      effect: Effect.ALLOW,
-      resources: ['*'],
-      actions: ['bedrock:*', 'logs:*'],
-    });
-    predictStreamFunction.role?.addToPrincipalPolicy(bedrockPolicy);
-    predictFunction.role?.addToPrincipalPolicy(bedrockPolicy);
-    predictTitleFunction.role?.addToPrincipalPolicy(bedrockPolicy);
-    generateImageFunction.role?.addToPrincipalPolicy(bedrockPolicy);
+    if (
+      typeof crossAccountBedrockRoleArn !== 'string' ||
+      crossAccountBedrockRoleArn === ''
+    ) {
+      const bedrockPolicy = new PolicyStatement({
+        effect: Effect.ALLOW,
+        resources: ['*'],
+        actions: ['bedrock:*', 'logs:*'],
+      });
+      predictStreamFunction.role?.addToPrincipalPolicy(bedrockPolicy);
+      predictFunction.role?.addToPrincipalPolicy(bedrockPolicy);
+      predictTitleFunction.role?.addToPrincipalPolicy(bedrockPolicy);
+      generateImageFunction.role?.addToPrincipalPolicy(bedrockPolicy);
+    } else {
+      // crossAccountBedrockRoleArn が指定されている場合のポリシー
+      const logsPolicy = new PolicyStatement({
+        effect: Effect.ALLOW,
+        actions: ['logs:*'],
+        resources: ['*'],
+      });
+      const assumeRolePolicy = new PolicyStatement({
+        effect: Effect.ALLOW,
+        actions: ['sts:AssumeRole'],
+        resources: [crossAccountBedrockRoleArn],
+      });
+      predictStreamFunction.role?.addToPrincipalPolicy(logsPolicy);
+      predictFunction.role?.addToPrincipalPolicy(logsPolicy);
+      predictTitleFunction.role?.addToPrincipalPolicy(logsPolicy);
+      generateImageFunction.role?.addToPrincipalPolicy(logsPolicy);
+      predictStreamFunction.role?.addToPrincipalPolicy(assumeRolePolicy);
+      predictFunction.role?.addToPrincipalPolicy(assumeRolePolicy);
+      predictTitleFunction.role?.addToPrincipalPolicy(assumeRolePolicy);
+      generateImageFunction.role?.addToPrincipalPolicy(assumeRolePolicy);
+    }
 
     const createChatFunction = new NodejsFunction(this, 'CreateChat', {
       runtime: Runtime.NODEJS_18_X,