Skip to content

Commit 5b059b9

Browse files
authored
fix(stepfunctions-tasks): add back BedrockInvokeModel to use JsonPath (#31325)
Reverts #31308 After discussion with the team, we are working on the "roll-forward fix" PR for this issue and not going ahead with this revert #31305 Reverting to resolve the existing conflicts.
1 parent 03ebca8 commit 5b059b9

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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 3 additions & 3 deletions
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

Lines changed: 33 additions & 2 deletions
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 1 addition & 1 deletion
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

Lines changed: 2 additions & 2 deletions
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

Lines changed: 41 additions & 2 deletions
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

Lines changed: 7 additions & 1 deletion
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

Lines changed: 21 additions & 0 deletions
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

Lines changed: 37 additions & 6 deletions
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)