Skip to content

Commit 3a3bbd8

Browse files
toshikwatoshikwa
authored
Add cross-account support for knowledge base, video generation, and speech-to-speech (#1044)
1 parent 1c11b36 commit 3a3bbd8

21 files changed

+447
-151
lines changed

‎docs/en/DEPLOY_OPTION.md

Lines changed: 72 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1778,10 +1778,16 @@ To use Bedrock from a different AWS account, you need to create one IAM role in
17781778
- `GenerativeAiUseCasesStack-APIPredictService`
17791779
- `GenerativeAiUseCasesStack-APIPredictStreamService`
17801780
- `GenerativeAiUseCasesStack-APIGenerateImageService`
1781+
- `GenerativeAiUseCasesStack-APIGenerateVideoService`
1782+
- `GenerativeAiUseCasesStack-APIListVideoJobsService`
1783+
- `GenerativeAiUseCasesStack-SpeechToSpeechTaskService`
1784+
- `GenerativeAiUseCasesStack-RagKnowledgeBaseRetrieve` (Only when using Knowledge Base)
1785+
- `GenerativeAiUseCasesStack-APIGetFileDownloadSigned` (Only when using Knowledge Base)
17811786

17821787
For details on how to specify Principals, refer to: [AWS JSON Policy Elements: Principal](https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_policies_elements_principal.html)
17831788

1784-
Principal configuration example (set in the different account)
1789+
<details>
1790+
<summary>Principal configuration example (set in the different account)</summary>
17851791

17861792
```json
17871793
{
@@ -1794,20 +1800,77 @@ Principal configuration example (set in the different account)
17941800
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictTitleServiceXXX-XXXXXXXXXXXX",
17951801
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictServiceXXXXXXXX-XXXXXXXXXXXX",
17961802
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictStreamServiceXX-XXXXXXXXXXXX",
1797-
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateImageServiceXX-XXXXXXXXXXXX"
1803+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateImageServiceXX-XXXXXXXXXXXX",
1804+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateVideoServiceXX-XXXXXXXXXXXX",
1805+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIListVideoJobsServiceXX-XXXXXXXXXXXX",
1806+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-SpeechToSpeechTaskService-XXXXXXXXXXXX",
1807+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-RagKnowledgeBaseRetrieveX-XXXXXXXXXXXX",
1808+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGetFileDownloadSignedU-XXXXXXXXXXXX"
17981809
]
17991810
},
1800-
"Action": "sts:AssumeRole",
1801-
"Condition": {}
1811+
"Action": "sts:AssumeRole"
18021812
}
18031813
]
18041814
}
18051815
```
18061816

1817+
</details>
1818+
1819+
<details>
1820+
<summary>Policy configuration example (set in the different account)</summary>
1821+
1822+
```json
1823+
{
1824+
"Version": "2012-10-17",
1825+
"Statement": [
1826+
{
1827+
"Sid": "AllowBedrockInvokeModel",
1828+
"Effect": "Allow",
1829+
"Action": [
1830+
"bedrock:InvokeModel*",
1831+
"bedrock:Rerank",
1832+
"bedrock:GetInferenceProfile",
1833+
"bedrock:GetAsyncInvoke",
1834+
"bedrock:ListAsyncInvokes"
1835+
],
1836+
"Resource": ["*"]
1837+
},
1838+
{
1839+
"Sid": "AllowS3PutObjectToVideoTempBucket",
1840+
"Effect": "Allow",
1841+
"Action": ["s3:PutObject"],
1842+
"Resource": ["arn:aws:s3:::<video-temp-bucket-name>/*"]
1843+
},
1844+
{
1845+
"Sid": "AllowBedrockRetrieveFromKnowledgeBase",
1846+
"Effect": "Allow",
1847+
"Action": ["bedrock:RetrieveAndGenerate*", "bedrock:Retrieve*"],
1848+
"Resource": [
1849+
"arn:aws:bedrock:<region>:<account-id>:knowledge-base/<knowledge-base-id>"
1850+
]
1851+
},
1852+
{
1853+
"Sid": "AllowS3GetPresignedUrl",
1854+
"Effect": "Allow",
1855+
"Action": ["s3:GetObject*"],
1856+
"Resource": ["arn:aws:s3:::<knowledge-base-datasource-bucket-name>/*"]
1857+
}
1858+
]
1859+
}
1860+
```
1861+
1862+
</details>
1863+
18071864
Set the following parameter:
18081865

18091866
- `crossAccountBedrockRoleArn` ... The ARN of the IAM role created in advance in the different account
18101867

1868+
When using Knowledge Base, you'll need to include these additional parameters:
1869+
1870+
- `ragKnowledgeBaseEnabled` ... Set to `true` to enable Knowledge Base
1871+
- `ragKnowledgeBaseId` ... Knowledge Base ID created in advance in the different account
1872+
- Knowledge Base must exist in the `modelRegion`
1873+
18111874
**Edit [parameter.ts](/packages/cdk/parameter.ts)**
18121875

18131876
```typescript
@@ -1816,6 +1879,8 @@ const envs: Record<string, Partial<StackInput>> = {
18161879
dev: {
18171880
crossAccountBedrockRoleArn:
18181881
'arn:aws:iam::AccountID:role/PreCreatedRoleName',
1882+
ragKnowledgeBaseEnabled: true, // Only when using Knowledge Base
1883+
ragKnowledgeBaseId: 'XXXXXXXXXX', // Only when using Knowledge Base
18191884
},
18201885
};
18211886
```
@@ -1826,7 +1891,9 @@ const envs: Record<string, Partial<StackInput>> = {
18261891
// cdk.json
18271892
{
18281893
"context": {
1829-
"crossAccountBedrockRoleArn": "arn:aws:iam::AccountID:role/PreCreatedRoleName"
1894+
"crossAccountBedrockRoleArn": "arn:aws:iam::AccountID:role/PreCreatedRoleName",
1895+
"ragKnowledgeBaseEnabled": true, // Only when using Knowledge Base
1896+
"ragKnowledgeBaseId": "XXXXXXXXXX" // Only when using Knowledge Base
18301897
}
18311898
}
18321899
```

‎docs/ja/DEPLOY_OPTION.md

Lines changed: 72 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1785,10 +1785,16 @@ const envs: Record<string, Partial<StackInput>> = {
17851785
- `GenerativeAiUseCasesStack-APIPredictService`
17861786
- `GenerativeAiUseCasesStack-APIPredictStreamService`
17871787
- `GenerativeAiUseCasesStack-APIGenerateImageService`
1788+
- `GenerativeAiUseCasesStack-APIGenerateVideoService`
1789+
- `GenerativeAiUseCasesStack-APIListVideoJobsService`
1790+
- `GenerativeAiUseCasesStack-SpeechToSpeechTaskService`
1791+
- `GenerativeAiUseCasesStack-RagKnowledgeBaseRetrieve` (Knowledge Base 利用時のみ)
1792+
- `GenerativeAiUseCasesStack-APIGetFileDownloadSigned` (Knowledge Base 利用時のみ)
17881793

17891794
Principal の指定方法について詳細を確認したい場合はこちらを参照ください: [AWS JSON ポリシーの要素: Principal](https://docs.aws.amazon.com/ja_jp/IAM/latest/UserGuide/reference_policies_elements_principal.html)
17901795

1791-
Principal 設定例 (別アカウントにて設定)
1796+
<details>
1797+
<summary>Principal 設定例 (別アカウントにて設定)</summary>
17921798

17931799
```json
17941800
{
@@ -1801,20 +1807,77 @@ Principal 設定例 (別アカウントにて設定)
18011807
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictTitleServiceXXX-XXXXXXXXXXXX",
18021808
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictServiceXXXXXXXX-XXXXXXXXXXXX",
18031809
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIPredictStreamServiceXX-XXXXXXXXXXXX",
1804-
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateImageServiceXX-XXXXXXXXXXXX"
1810+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateImageServiceXX-XXXXXXXXXXXX",
1811+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGenerateVideoServiceXX-XXXXXXXXXXXX",
1812+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIListVideoJobsServiceXX-XXXXXXXXXXXX",
1813+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-SpeechToSpeechTaskService-XXXXXXXXXXXX",
1814+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-RagKnowledgeBaseRetrieveX-XXXXXXXXXXXX",
1815+
"arn:aws:iam::111111111111:role/GenerativeAiUseCasesStack-APIGetFileDownloadSignedU-XXXXXXXXXXXX"
18051816
]
18061817
},
1807-
"Action": "sts:AssumeRole",
1808-
"Condition": {}
1818+
"Action": "sts:AssumeRole"
18091819
}
18101820
]
18111821
}
18121822
```
18131823

1824+
</details>
1825+
1826+
<details>
1827+
<summary>ポリシー設定例 (別アカウントにて設定)</summary>
1828+
1829+
```json
1830+
{
1831+
"Version": "2012-10-17",
1832+
"Statement": [
1833+
{
1834+
"Sid": "AllowBedrockInvokeModel",
1835+
"Effect": "Allow",
1836+
"Action": [
1837+
"bedrock:InvokeModel*",
1838+
"bedrock:Rerank",
1839+
"bedrock:GetInferenceProfile",
1840+
"bedrock:GetAsyncInvoke",
1841+
"bedrock:ListAsyncInvokes"
1842+
],
1843+
"Resource": ["*"]
1844+
},
1845+
{
1846+
"Sid": "AllowS3PutObjectToVideoTempBucket",
1847+
"Effect": "Allow",
1848+
"Action": ["s3:PutObject"],
1849+
"Resource": ["arn:aws:s3:::<video-temp-bucket-name>/*"]
1850+
},
1851+
{
1852+
"Sid": "AllowBedrockRetrieveFromKnowledgeBase",
1853+
"Effect": "Allow",
1854+
"Action": ["bedrock:RetrieveAndGenerate*", "bedrock:Retrieve*"],
1855+
"Resource": [
1856+
"arn:aws:bedrock:<region>:<account-id>:knowledge-base/<knowledge-base-id>"
1857+
]
1858+
},
1859+
{
1860+
"Sid": "AllowS3GetPresignedUrl",
1861+
"Effect": "Allow",
1862+
"Action": ["s3:GetObject*"],
1863+
"Resource": ["arn:aws:s3:::<knowledge-base-datasource-bucket-name>/*"]
1864+
}
1865+
]
1866+
}
1867+
```
1868+
1869+
</details>
1870+
18141871
パラメータには以下の値を設定します。
18151872

18161873
- `crossAccountBedrockRoleArn` ... 別アカウントで事前に作成した IAM ロールの ARN です
18171874

1875+
Knowledge Base を利用する場合は、下記パラメーターも指定します。
1876+
1877+
- `ragKnowledgeBaseEnabled` ... Knowledge Base を有効化する場合は `true` とします
1878+
- `ragKnowledgeBaseId` ... 別アカウントに事前構築した Knowledge Base の ID です
1879+
- Knowledge Base は `modelRegion` に存在する必要があります
1880+
18181881
**[parameter.ts](/packages/cdk/parameter.ts) を編集**
18191882

18201883
```typescript
@@ -1823,6 +1886,8 @@ const envs: Record<string, Partial<StackInput>> = {
18231886
dev: {
18241887
crossAccountBedrockRoleArn:
18251888
'arn:aws:iam::アカウントID:role/事前に作成したロール名',
1889+
ragKnowledgeBaseEnabled: true, // Knowledge Base を利用する場合のみ
1890+
ragKnowledgeBaseId: 'XXXXXXXXXX', // Knowledge Base を利用する場合のみ
18261891
},
18271892
};
18281893
```
@@ -1833,7 +1898,9 @@ const envs: Record<string, Partial<StackInput>> = {
18331898
// cdk.json
18341899
{
18351900
"context": {
1836-
"crossAccountBedrockRoleArn": "arn:aws:iam::アカウントID:role/事前に作成したロール名"
1901+
"crossAccountBedrockRoleArn": "arn:aws:iam::アカウントID:role/事前に作成したロール名",
1902+
"ragKnowledgeBaseEnabled": true, // Knowledge Base を利用する場合のみ
1903+
"ragKnowledgeBaseId": "XXXXXXXXXX" // Knowledge Base を利用する場合のみ
18371904
}
18381905
}
18391906
```

‎packages/cdk/lambda/getFileDownloadSignedUrl.ts

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,23 @@ import { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
22
import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
33
import { GetObjectCommand, S3Client } from '@aws-sdk/client-s3';
44
import { GetFileDownloadSignedUrlRequest } from 'generative-ai-use-cases';
5+
import { initKnowledgeBaseS3Client } from './utils/bedrockClient';
6+
7+
const MODEL_REGION = process.env.MODEL_REGION as string;
58

69
export const handler = async (
710
event: APIGatewayProxyEvent
811
): Promise<APIGatewayProxyResult> => {
912
try {
1013
const req = event.queryStringParameters as GetFileDownloadSignedUrlRequest;
1114

12-
const client = new S3Client({
13-
region: req.region,
14-
});
15+
// We pass `s3Type` parameter since Knowledge Base may need to reference S3 in a different account
16+
const client =
17+
req.s3Type === 'knowledgeBase'
18+
? await initKnowledgeBaseS3Client({
19+
region: req.region ?? MODEL_REGION,
20+
})
21+
: new S3Client({ region: req.region });
1522
const command = new GetObjectCommand({
1623
Bucket: req.bucketName,
1724
Key: req.filePrefix,

‎packages/cdk/lambda/repositoryVideoJob.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,13 @@ import {
1515
GetAsyncInvokeCommand,
1616
ValidationException,
1717
} from '@aws-sdk/client-bedrock-runtime';
18-
import { initBedrockClient } from './utils/bedrockApi';
1918
import { CopyVideoJobParams } from './copyVideoJob';
2019
import {
2120
LambdaClient,
2221
InvokeCommand,
2322
InvocationType,
2423
} from '@aws-sdk/client-lambda';
24+
import { initBedrockRuntimeClient } from './utils/bedrockClient';
2525

2626
const BUCKET_NAME: string = process.env.BUCKET_NAME!;
2727
const TABLE_NAME: string = process.env.TABLE_NAME!;
@@ -90,7 +90,7 @@ const checkAndUpdateJob = async (
9090
job: VideoJob
9191
): Promise<'InProgress' | 'Completed' | 'Failed' | 'Finalizing'> => {
9292
try {
93-
const client = await initBedrockClient(job.region);
93+
const client = await initBedrockRuntimeClient({ region: job.region });
9494
const command = new GetAsyncInvokeCommand({
9595
invocationArn: job.invocationArn,
9696
});

‎packages/cdk/lambda/retrieveKnowledgeBase.ts

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,10 @@
11
import * as lambda from 'aws-lambda';
2-
import {
3-
BedrockAgentRuntimeClient,
4-
RetrieveCommand,
5-
} from '@aws-sdk/client-bedrock-agent-runtime';
2+
import { RetrieveCommand } from '@aws-sdk/client-bedrock-agent-runtime';
63
import { RetrieveKnowledgeBaseRequest } from 'generative-ai-use-cases';
4+
import { initBedrockAgentRuntimeClient } from './utils/bedrockClient';
75

86
const KNOWLEDGE_BASE_ID = process.env.KNOWLEDGE_BASE_ID;
9-
const MODEL_REGION = process.env.MODEL_REGION;
7+
const MODEL_REGION = process.env.MODEL_REGION as string;
108

119
exports.handler = async (
1210
event: lambda.APIGatewayProxyEvent
@@ -25,23 +23,18 @@ exports.handler = async (
2523
};
2624
}
2725

28-
const kb = new BedrockAgentRuntimeClient({
29-
region: MODEL_REGION,
30-
});
26+
const client = await initBedrockAgentRuntimeClient({ region: MODEL_REGION });
3127
const retrieveCommand = new RetrieveCommand({
3228
knowledgeBaseId: KNOWLEDGE_BASE_ID,
33-
retrievalQuery: {
34-
text: query,
35-
},
29+
retrievalQuery: { text: query },
3630
retrievalConfiguration: {
3731
vectorSearchConfiguration: {
3832
numberOfResults: 10,
3933
overrideSearchType: 'HYBRID',
4034
},
4135
},
4236
});
43-
44-
const retrieveRes = await kb.send(retrieveCommand);
37+
const retrieveRes = await client.send(retrieveCommand);
4538

4639
return {
4740
statusCode: 200,

‎packages/cdk/lambda/speechToSpeechTask.ts

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ import { events, EventsChannel } from 'aws-amplify/data';
33
import { fromNodeProviderChain } from '@aws-sdk/credential-providers';
44
import { randomUUID } from 'crypto';
55
import {
6-
BedrockRuntimeClient,
76
InvokeModelWithBidirectionalStreamCommand,
87
InvokeModelWithBidirectionalStreamInput,
98
InvokeModelWithBidirectionalStreamCommandOutput,
@@ -15,9 +14,12 @@ import {
1514
SpeechToSpeechEvent,
1615
Model,
1716
} from 'generative-ai-use-cases';
17+
import { initBedrockRuntimeClient } from './utils/bedrockClient';
1818

1919
Object.assign(global, { WebSocket: require('ws') });
2020

21+
const MODEL_REGION = process.env.MODEL_REGION as string;
22+
2123
const MAX_AUDIO_INPUT_QUEUE_SIZE = 200;
2224
const MIN_AUDIO_OUTPUT_QUEUE_SIZE = 10;
2325
const MAX_AUDIO_OUTPUT_PER_BATCH = 20;
@@ -422,8 +424,8 @@ export const handler = async (event: { channelId: string; model: Model }) => {
422424

423425
console.log('promptName', promptName);
424426

425-
const bedrock = new BedrockRuntimeClient({
426-
region: event.model.region,
427+
const bedrockRuntimeClient = await initBedrockRuntimeClient({
428+
region: event.model.region ?? MODEL_REGION,
427429
requestHandler: new NodeHttp2Handler({
428430
requestTimeout: 300000,
429431
sessionTimeout: 300000,
@@ -514,7 +516,7 @@ export const handler = async (event: { channelId: string; model: Model }) => {
514516

515517
console.log('Async iterator created');
516518

517-
const response = await bedrock.send(
519+
const response = await bedrockRuntimeClient.send(
518520
new InvokeModelWithBidirectionalStreamCommand({
519521
modelId: event.model.modelId,
520522
body: asyncIterator,

0 commit comments

Comments
 (0)