Skip to content

Commit f5dd73b

Browse files
authored
feat(stepfunctions-tasks): allow BedrockInvokeModel to use JsonPath (#30298)
### Issue # (if applicable) Closes #29229. ### Reason for this change When trying to use JsonPath to specify the S3 URIs that BedrockInvokeModel will read from and write from, you get an error. Example of the Error message: `jsii.errors.JavaScriptError: Error: Field references must be the entire string, cannot concatenate them (found 's3://${Token[prompt_bucket.348]}/${Token[prompt_key.349]}')` ### Description of changes Extended the inputPath property to be allowed as an input value for the task state. Instead of adding a new S3Uri props in current `BedrockInvokeModelProps` as proposed in the original issue, leveraged the `inputPath` property that is already defined in `sfn.TaskStateBaseProps` and being extended by `BedrockInvokeModelInputProps` and `BedrockInvokeModelOutputProps` **Limitation:** We cannot limit the resource policy to specific input token for which the value might be coming from the prompt, so had to keep it as [*] here. ### Description of how you validated changes Added unit tests. Successful deployment of integration tests in the account. ### Checklist - [x] Unit Tests - [x] Integration Tests - [x] Updated ReadMe - [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md) ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
1 parent 82b163d commit f5dd73b

File tree

11 files changed

+255
-19
lines changed

11 files changed

+255
-19
lines changed

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json

+1-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json

+3-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json

+33-2
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,25 @@
4141
]
4242
]
4343
}
44+
},
45+
{
46+
"Action": [
47+
"s3:GetObject",
48+
"s3:PutObject"
49+
],
50+
"Effect": "Allow",
51+
"Resource": {
52+
"Fn::Join": [
53+
"",
54+
[
55+
"arn:",
56+
{
57+
"Ref": "AWS::Partition"
58+
},
59+
":s3:::*"
60+
]
61+
]
62+
}
4463
}
4564
],
4665
"Version": "2012-10-17"
@@ -72,7 +91,19 @@
7291
{
7392
"Ref": "AWS::Region"
7493
},
75-
"::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"End\":true,\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:",
94+
"::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText\":\"Generate a list of five first names.\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt2\":{\"Next\":\"Prompt3\",\"Type\":\"Task\",\"ResultPath\":\"$\",\"ResultSelector\":{\"names.$\":\"$.Body.results[0].outputText\"},\"Resource\":\"arn:",
95+
{
96+
"Ref": "AWS::Partition"
97+
},
98+
":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:",
99+
{
100+
"Ref": "AWS::Partition"
101+
},
102+
":bedrock:",
103+
{
104+
"Ref": "AWS::Region"
105+
},
106+
"::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"End\":true,\"Type\":\"Task\",\"InputPath\":\"$.names\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
76107
{
77108
"Ref": "AWS::Partition"
78109
},
@@ -84,7 +115,7 @@
84115
{
85116
"Ref": "AWS::Region"
86117
},
87-
"::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}}},\"TimeoutSeconds\":30}"
118+
"::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
88119
]
89120
]
90121
},

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out

+1-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json

+1-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json

+2-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json

+41-2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts

+7-1
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,13 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
5252
resultPath: '$',
5353
});
5454

55-
const chain = sfn.Chain.start(prompt1).next(prompt2);
55+
const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
56+
model,
57+
inputPath: sfn.JsonPath.stringAt('$.names'),
58+
outputPath: sfn.JsonPath.stringAt('$.names'),
59+
});
60+
61+
const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3);
5662

5763
new sfn.StateMachine(stack, 'StateMachine', {
5864
definitionBody: sfn.DefinitionBody.fromChainable(chain),

packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md

+21
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,27 @@ const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
398398
names: sfn.JsonPath.stringAt('$.Body.results[0].outputText'),
399399
},
400400
});
401+
```
402+
### Using Input Path
403+
404+
Provide S3 URI as an input or output path to invoke a model
405+
406+
```ts
407+
408+
import * as bedrock from 'aws-cdk-lib/aws-bedrock';
409+
410+
const model = bedrock.FoundationModel.fromFoundationModelId(
411+
this,
412+
'Model',
413+
bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1,
414+
);
415+
416+
const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
417+
model,
418+
inputPath: sfn.JsonPath.stringAt('$.prompt'),
419+
outputPath: sfn.JsonPath.stringAt('$.prompt'),
420+
});
421+
401422
```
402423

403424
You can apply a guardrail to the invocation by setting `guardrail`.

packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts

+37-6
Original file line numberDiff line numberDiff line change
@@ -140,12 +140,14 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
140140

141141
constructor(scope: Construct, id: string, private readonly props: BedrockInvokeModelProps) {
142142
super(scope, id, props);
143+
143144
this.integrationPattern = props.integrationPattern ?? sfn.IntegrationPattern.REQUEST_RESPONSE;
144145

145146
validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS);
146147

147148
const isBodySpecified = props.body !== undefined;
148-
const isInputSpecified = props.input !== undefined && props.input.s3Location !== undefined;
149+
//Either specific props.input with bucket name and object key or input s3 path
150+
const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
149151

150152
if (isBodySpecified && isInputSpecified) {
151153
throw new Error('Either `body` or `input` must be specified, but not both.');
@@ -171,7 +173,21 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
171173
}),
172174
];
173175

174-
if (this.props.input !== undefined && this.props.input.s3Location !== undefined) {
176+
if (this.props.inputPath !== undefined) {
177+
policyStatements.push(
178+
new iam.PolicyStatement({
179+
actions: ['s3:GetObject'],
180+
resources: [
181+
Stack.of(this).formatArn({
182+
region: '',
183+
account: '',
184+
service: 's3',
185+
resource: '*',
186+
}),
187+
],
188+
}),
189+
);
190+
} else if (this.props.input !== undefined && this.props.input.s3Location !== undefined) {
175191
policyStatements.push(
176192
new iam.PolicyStatement({
177193
actions: ['s3:GetObject'],
@@ -188,7 +204,21 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
188204
);
189205
}
190206

191-
if (this.props.output !== undefined && this.props.output.s3Location !== undefined) {
207+
if (this.props.outputPath !== undefined) {
208+
policyStatements.push(
209+
new iam.PolicyStatement({
210+
actions: ['s3:PutObject'],
211+
resources: [
212+
Stack.of(this).formatArn({
213+
region: '',
214+
account: '',
215+
service: 's3',
216+
resource: '*',
217+
}),
218+
],
219+
}),
220+
);
221+
} else if (this.props.output !== undefined && this.props.output.s3Location !== undefined) {
192222
policyStatements.push(
193223
new iam.PolicyStatement({
194224
actions: ['s3:PutObject'],
@@ -241,10 +271,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
241271
Body: this.props.body?.value,
242272
Input: this.props.input?.s3Location ? {
243273
S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
244-
} : undefined,
274+
} : this.props.inputPath ? { S3Uri: this.props.inputPath } : undefined,
245275
Output: this.props.output?.s3Location ? {
246276
S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
247-
} : undefined,
277+
} : this.props.outputPath ? { S3Uri: this.props.outputPath }: undefined,
248278
GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
249279
GuardrailVersion: this.props.guardrail?.guardrailVersion,
250280
Trace: this.props.traceEnabled === undefined
@@ -254,5 +284,6 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
254284
: 'DISABLED',
255285
}),
256286
};
257-
}
287+
};
258288
}
289+

0 commit comments

Comments
 (0)