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 1 commit
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,
s3inputPath: sfn.JsonPath.stringAt('$.names'),
s3outputPath: sfn.JsonPath.stringAt('$.names'),
});

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 s3inputPath?: 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 s3outputPath?: string;
shikha372 marked this conversation as resolved.
Show resolved Hide resolved

/**
* 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.s3inputPath !== 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.s3inputPath !== 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.s3outputPath !== 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.s3inputPath ? { S3Uri: this.props.s3inputPath } : undefined,
shikha372 marked this conversation as resolved.
Show resolved Hide resolved
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.s3outputPath ? { S3Uri: this.props.s3outputPath }: undefined,
GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
GuardrailVersion: this.props.guardrail?.guardrailVersion,
Trace: this.props.traceEnabled === undefined
Expand Down
Loading