Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(stepfunctions-tasks): fix bedrock input/output path in step-funct… #31305

Merged
merged 19 commits into from
Sep 5, 2024
Merged

fix(stepfunctions-tasks): fix bedrock input/output path in step-funct… #31305

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
d504493
fix(stepfunctions-tasks): fix bedrock input/output path in step-funct…
shikha372 Sep 3, 2024
8f0d77f
fix(stepfunctions-tasks): fix bedrock input/output path in step-funct…
shikha372 Sep 3, 2024
7682fa2
update readme
shikha372 Sep 3, 2024
6c6ac5f
move input and output s3 uri to bedrock invoke model props
shikha372 Sep 3, 2024
070af9e
adding feature flag
shikha372 Sep 4, 2024
33cd3bc
addressing comments and last feedback
shikha372 Sep 5, 2024
47d60bd
update feature flag
shikha372 Sep 5, 2024
7589bd4
fix small nits
shikha372 Sep 5, 2024
38b80bd
fix small nits
shikha372 Sep 5, 2024
211e23f
fix README
shikha372 Sep 5, 2024
af747c8
set feature flag default to true
shikha372 Sep 5, 2024
a52bdcf
fix integ
shikha372 Sep 5, 2024
451b442
fix nit
shikha372 Sep 5, 2024
0b7d4e8
fix last nit hopefully
shikha372 Sep 5, 2024
4415b89
Update features.ts
shikha372 Sep 5, 2024
2ce70ff
set flag type to api default
shikha372 Sep 5, 2024
dfa2739
fix nit
shikha372 Sep 5, 2024
0a7c25b
final fix for successful build
shikha372 Sep 5, 2024
523c05d
Merge branch 'main' into bedrock_input_path
mergify[bot] Sep 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...ock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 3 additions & 3 deletions ...g.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 13 additions & 1 deletion ...invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,19 @@
{
"Ref": "AWS::Region"
},
"::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:",
"::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\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
{
"Ref": "AWS::Partition"
},
":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:",
{
"Ref": "AWS::Partition"
},
":bedrock:",
{
"Ref": "AWS::Region"
},
"::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
{
"Ref": "AWS::Partition"
},
Expand Down
2 changes: 1 addition & 1 deletion ...rk-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion ...integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions ...eg/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

22 changes: 21 additions & 1 deletion ...-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

23 changes: 21 additions & 2 deletions ...k-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,13 +52,32 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
resultPath: '$',
});

/** Test for Bedrock Output Path */
const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
model,
inputPath: sfn.JsonPath.stringAt('$.names'),
body: sfn.TaskInput.fromObject(
{
inputText: sfn.JsonPath.format(
'Alphabetize this list of first names:\n{}',
sfn.JsonPath.stringAt('$.names'),
),
textGenerationConfig: {
maxTokenCount: 100,
temperature: 1,
},
},
),
outputPath: sfn.JsonPath.stringAt('$.names'),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can help me to understand why body is required instead of inputPath?

Copy link
Contributor Author

@shikha372 shikha372 Sep 3, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just added this integ test is to make sure, we are not breaking any existing use case for outputPath or inputPath from the base class, looking back at why it might have been marked as required could be a decision taken while designing this s3 input parameter, i don't see it something as required in API doc

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In order to test the inputPath field, you can do a Pass state to transform the previous output into the format expected for the InvokeModel body.

});

const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3);
/** Test for Bedrock s3 URI Path */
const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
model,
s3InputUri: sfn.JsonPath.stringAt('$.names'),
s3OutputUri: sfn.JsonPath.stringAt('$.names'),
});

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you run the verification steps listed at the top of this class?
Stack verification steps

const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);

new sfn.StateMachine(stack, 'StateMachine', {
definitionBody: sfn.DefinitionBody.fromChainable(chain),
Expand Down
28 changes: 23 additions & 5 deletions packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,24 @@ export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
*/
readonly output?: BedrockInvokeModelOutputProps;
shikha372 marked this conversation as resolved.
Show resolved Hide resolved

/**
* The destination location where the API response is written.
shikha372 marked this conversation as resolved.
Show resolved Hide resolved
*
* This field can be used to specify s3 URI in the form of token
*
* @default - The API response body is returned in the result.
*/
readonly s3InputUri?: string;
shikha372 marked this conversation as resolved.
Show resolved Hide resolved

/**
* The source location where the API response is written.
*
* This field can be used to specify s3 URI in the form of token
*
* @default - The API response body is returned in the result.
*/
readonly s3OutputUri?: string;

/**
* The guardrail is applied to the invocation
*
Expand Down Expand Up @@ -147,7 +165,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {

const isBodySpecified = props.body !== undefined;
//Either specific props.input with bucket name and object key or input s3 path
const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.s3InputUri !== undefined);

if (isBodySpecified && isInputSpecified) {
throw new Error('Either `body` or `input` must be specified, but not both.');
Expand All @@ -173,7 +191,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
}),
];

if (this.props.inputPath !== undefined) {
if (this.props.s3InputUri !== undefined) {
policyStatements.push(
new iam.PolicyStatement({
actions: ['s3:GetObject'],
Expand Down Expand Up @@ -204,7 +222,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
);
}

if (this.props.outputPath !== undefined) {
if (this.props.s3OutputUri !== undefined) {
policyStatements.push(
new iam.PolicyStatement({
actions: ['s3:PutObject'],
Expand Down Expand Up @@ -271,10 +289,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
Body: this.props.body?.value,
Input: this.props.input?.s3Location ? {
S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
} : this.props.inputPath ? { S3Uri: this.props.inputPath } : undefined,
} : this.props.s3InputUri ? { S3Uri: this.props.s3InputUri } : undefined,
Output: this.props.output?.s3Location ? {
S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
} : this.props.outputPath ? { S3Uri: this.props.outputPath }: undefined,
} : this.props.s3OutputUri ? { S3Uri: this.props.s3OutputUri }: undefined,
GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
GuardrailVersion: this.props.guardrail?.guardrailVersion,
Trace: this.props.traceEnabled === undefined
Expand Down
10 changes: 4 additions & 6 deletions packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,8 +210,8 @@ describe('Invoke Model', () => {

const task = new BedrockInvokeModel(stack, 'Invoke', {
model,
inputPath: sfn.JsonPath.stringAt('$.prompt'),
outputPath: sfn.JsonPath.stringAt('$.prompt'),
s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
});

new sfn.StateMachine(stack, 'StateMachine', {
Expand All @@ -234,8 +234,6 @@ describe('Invoke Model', () => {
],
},
End: true,
InputPath: '$.prompt',
OutputPath: '$.prompt',
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there other unit tests that test this functionality?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are two unit tests, one to check the expected template and one to check the permissions populated in policy

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The title of this unit test is "invoke model allows input and output json path". I think this test was originally meant to test the original behavior of inputPath and outputPath. So we need an additional new test for these new s3InputUri and s3OutputUri params, AND fix this unit test to go back to testing the original behavior of inputPath and outputPath.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update: I see in the history now that you had actually added this test in the last change. I do still think we should have both tests (inputPath and s3InputUri)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @clareliguori , this leads to one question can body be skipped in case of an inputPath ? since the earlier validation seems to added only for body and input and didn't point to a definition where we might be using these base props.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh I see, yes you still need body. With inputPath, you would do something like body=sfn.JsonPath.object_at("$")

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah seems like I totally missed adding unit tests for inputPath and outputPath

Parameters: {
ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
Input: {
Expand Down Expand Up @@ -330,8 +328,8 @@ describe('Invoke Model', () => {
// WHEN
const task = new BedrockInvokeModel(stack, 'Invoke', {
model,
inputPath: sfn.JsonPath.stringAt('$.prompt'),
outputPath: sfn.JsonPath.stringAt('$.prompt'),
s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
});

new sfn.StateMachine(stack, 'StateMachine', {
Expand Down
Loading