From e55d7d929710fae0152e3bb40e81f1a2a8068b76 Mon Sep 17 00:00:00 2001 From: Rafael Mosca Date: Wed, 6 Nov 2024 19:33:17 +0100 Subject: [PATCH] feat(bedrock): prompt flows first def --- .../bedrock/prompt-flows/flow-connections.ts | 133 +++ .../bedrock/prompt-flows/flow-node-props.ts | 415 +++++++++ .../bedrock/prompt-flows/flow-nodes.ts | 857 ++++++++++++++++++ .../bedrock/prompt-flows/flow-version.ts | 56 ++ src/cdk-lib/bedrock/prompt-flows/flow.ts | 283 ++++++ 5 files changed, 1744 insertions(+) create mode 100644 src/cdk-lib/bedrock/prompt-flows/flow-connections.ts create mode 100644 src/cdk-lib/bedrock/prompt-flows/flow-node-props.ts create mode 100644 src/cdk-lib/bedrock/prompt-flows/flow-nodes.ts create mode 100644 src/cdk-lib/bedrock/prompt-flows/flow-version.ts create mode 100644 src/cdk-lib/bedrock/prompt-flows/flow.ts diff --git a/src/cdk-lib/bedrock/prompt-flows/flow-connections.ts b/src/cdk-lib/bedrock/prompt-flows/flow-connections.ts new file mode 100644 index 00000000..841f97b9 --- /dev/null +++ b/src/cdk-lib/bedrock/prompt-flows/flow-connections.ts @@ -0,0 +1,133 @@ +/** + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ + +import { CfnFlow } from "aws-cdk-lib/aws-bedrock"; +import { FlowNode } from "./flow-nodes"; + +export enum ConnectionType { + DATA = "Data", + CONDITIONAL = "Conditional", +} + +/** + * Properties require for all flow connections. + */ +export interface ConnectionCommonProps { + /** + * The name of the connection. + */ + readonly name: string; + /** + * The node that the connection starts at. + */ + readonly source: FlowNode; + /** + * The node that the connection ends at. + */ + readonly target: FlowNode; +} + +export interface DataConnectionProps extends ConnectionCommonProps { + /** + * The configuration of a connection originating from a simple node. + */ + readonly config: DataConnectionConfiguration; +} + +export interface ConditionalConnectionProps extends ConnectionCommonProps { + /** + * The name of the condition + */ + readonly condition?: string; +} + +export interface DataConnectionConfiguration { + /** + * The name of the output in the source node that the connection begins from. + */ + readonly sourceOutput: string; + /** + * The name of the input in the target node that the connection ends at. + */ + readonly targetInput: string; +} + +/** + * + */ +export class FlowConnection { + // ------------------------------------------------------ + // Static Methods + // ------------------------------------------------------ + /** + * Static method to create a conditional connection. + */ + public static conditional(props: ConditionalConnectionProps): FlowConnection { + return new FlowConnection({ + name: props.name, + source: props.source.name, + target: props.target.name, + type: ConnectionType.CONDITIONAL, + configuration: { + conditional: { + condition: props.condition, + }, + }, + }); + } + /** + * Static method to create a data connection. + */ + public static data(props: DataConnectionProps): FlowConnection { + return new FlowConnection({ + name: props.name, + source: props.source.name, + target: props.target.name, + type: ConnectionType.DATA, + configuration: { + data: props.config, + }, + }); + } + + // ------------------------------------------------------ + // Constructor + // ------------------------------------------------------ + readonly name: string; + readonly source: string; + readonly target: string; + readonly type: ConnectionType; + readonly configuration: CfnFlow.FlowConnectionConfigurationProperty; + + protected constructor(props: any) { + this.name = props.name; + this.source = props.source; + this.target = props.target; + this.type = props.type; + this.configuration = props.configuration; + } + + asCfnProperty(): CfnFlow.FlowConnectionProperty { + return { + name: this.name, + source: this.source, + target: this.target, + type: this.type, + configuration: this.configuration, + }; + } + + // ------------------------------------------------------ + // Properties + // ------------------------------------------------------ +} diff --git a/src/cdk-lib/bedrock/prompt-flows/flow-node-props.ts b/src/cdk-lib/bedrock/prompt-flows/flow-node-props.ts new file mode 100644 index 00000000..3a5ebbd0 --- /dev/null +++ b/src/cdk-lib/bedrock/prompt-flows/flow-node-props.ts @@ -0,0 +1,415 @@ +/** + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ + +import { IAlias, IFunction } from "aws-cdk-lib/aws-lambda"; +import { IBucket } from "aws-cdk-lib/aws-s3"; +import { FlowNode } from "./flow-nodes"; +import { Prompt } from "../prompts/prompt"; +import { IKnowledgeBase } from "../knowledge-base"; +import { BedrockFoundationModel } from "../models"; +import { IAgentAlias } from "../agent-alias"; + +/****************************************************************************** + * COMMON + *****************************************************************************/ +/** + * The type of flow node. + */ +export enum FlowNodeType { + INPUT = "Input", + OUTPUT = "Output", + KNOWLEDGE_BASE = "KnowledgeBase", + CONDITION = "Condition", + LEX = "Lex", + PROMPT = "Prompt", + LAMBDA = "LambdaFunction", + AGENT = "Agent", + STORAGE = "Storage", + RETRIEVAL = "Retrieval", + ITERATOR = "Iterator", + COLLECTOR = "Collector", +} + +/** + * The data type of the input. If the input doesn't match this type at runtime, + * a validation error will be thrown. + */ +export enum FlowNodeDataType { + STRING = "String", + NUMBER = "Number", + BOOLEAN = "Boolean", + OBJECT = "Object", + ARRAY = "Array", +} + +export interface FlowNodeInputSource { + /** + * The flow node to use as source. + */ + readonly sourceNode: FlowNode; + /** + * The name of the output to use as source. + * @default "Inferred if single output from source node" + */ + readonly sourceOutputName?: string; + /** + * The expression to use for the output. + * @default "$.data" + */ + readonly expression?: string; +} + +export interface UntypedFlowNodeInput { + /** + * Where the data from this input will be taken from. + */ + readonly valueFrom: FlowNodeInputSource; +} + +export interface TypedFlowNodeInput extends UntypedFlowNodeInput { + /** + * The flow node data type to use for the input. + * @default FlowNodeDataType.STRING + */ + readonly type?: FlowNodeDataType; +} + +export interface NamedTypedFlowNodeInput extends TypedFlowNodeInput { + /** + * The name for the input. + */ + readonly name: string; +} + +export interface NamedFlowNodeOutput { + /** + * The name for the output. + */ + readonly name: string; + /** + * The flow node data type to use for the output. + * @default FlowNodeDataType.STRING + */ + readonly type?: FlowNodeDataType; +} + +/****************************************************************************** + * INPUT / OUTPUT + *****************************************************************************/ +/** + * Properties common to all node types + */ +export interface BaseNodeProps { + /** + * The name of the node. + */ + readonly name: string; +} + +// ------------------------------------------------------------------ +// INPUT Node +// ------------------------------------------------------------------ +/** + * The properties for an Input Data node. + */ +export interface InputDataNodeProps extends BaseNodeProps { + /** + * The input data type to use. + * @default FlowNodeDataType.STRING + */ + readonly inputDataType?: FlowNodeDataType; +} + +// ------------------------------------------------------------------ +// OUTPUT Node +// ------------------------------------------------------------------ +/** + * The properties for an Output Data node. + */ +export interface OutputDataNodeProps extends BaseNodeProps { + /** + * The output data to use. + */ + readonly outputData: TypedFlowNodeInput; +} + +/****************************************************************************** + * PROCESSING + *****************************************************************************/ +// ------------------------------------------------------------------ +// LAMBDA Node +// ------------------------------------------------------------------ +/** + * The properties for a Lambda Function Node. + */ +export interface LambdaFunctionNodeProps extends BaseNodeProps { + /** + * The Lambda function to use. + */ + readonly lambdaFunction: IFunction; + /** + * The Lambda function alias to use in the node + * + * @default "$LATEST" + */ + readonly functionAlias?: IAlias; + /** + * The inputs to use. + */ + readonly inputs: NamedTypedFlowNodeInput[]; + /** + * The output type to use for the `functionResponse` output. + * @default FlowNodeDataType.STRING + */ + readonly functionResponseType?: FlowNodeDataType; +} + +// ------------------------------------------------------------------ +// AGENT Node +// ------------------------------------------------------------------ +/** + * The properties for an Agent Node. + */ +export interface AgentNodeProps extends BaseNodeProps { + /** + * The Agent Alias this node refers to. + */ + readonly agentAlias: IAgentAlias; + /** + * The prompt to send to the agent. + */ + readonly agentInput: TypedFlowNodeInput; + /** + * Any prompt attributes to send alongside the prompt. + * + * @default - No prompt attributes. + */ + readonly promptAttributes?: TypedFlowNodeInput; + /** + * Any session attributes to send alongside the prompt. + * + * @default - No session attributes. + */ + readonly sessionAttributes?: TypedFlowNodeInput; +} + +// ------------------------------------------------------------------ +// PROMPT Node +// ------------------------------------------------------------------ +/** + * The properties for a Prompt Node. + */ +export interface PromptNodeProps extends BaseNodeProps { + /** + * The prompt to use. + */ + readonly prompt: Prompt; + /** + * The inputs for the prompt. + * @default - No input + */ + readonly inputs?: NamedTypedFlowNodeInput[]; +} + +/****************************************************************************** + * STORAGE + *****************************************************************************/ +// ------------------------------------------------------------------ +// Knowledge Base Node - Retrieve +// ------------------------------------------------------------------ +/** + * The properties for a Knowledge Base with retrieve functionality node. + * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html + */ +export interface KnowledgeBaseRetrieveNodeProps extends BaseNodeProps { + /** + * The Knowledge base to use. + */ + readonly knowledgeBase: IKnowledgeBase; + /** + * The input expression to use for the retrieval query. + */ + readonly retrievalQuery: UntypedFlowNodeInput; +} + +// ------------------------------------------------------------------ +// Knowledge Base Node - RetrieveAndGenerate +// ------------------------------------------------------------------ +/** + * The properties for a Knowledge Base with retrieve and generate functionality node. + * @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html + */ +export interface KnowledgeBaseRetrieveAndGenerateNodeProps extends KnowledgeBaseRetrieveNodeProps { + /** + * The FM model to use to generate responses from the information it retrieves. + */ + readonly model: BedrockFoundationModel; +} + +// ------------------------------------------------------------------ +// S3 Retieval Node +// ------------------------------------------------------------------ +/** + * The properties for an S3 Retrieval Data node. + */ +export interface S3RetrievalNodeProps extends BaseNodeProps { + /** + * The S3 bucket to use. + */ + readonly bucket: IBucket; + /** + * The object key input to use. + */ + readonly objectKeyInput: UntypedFlowNodeInput; +} + +// ------------------------------------------------------------------ +// S3 Storage Node +// ------------------------------------------------------------------ +/** + * The properties for an S3 Storage Data node. + */ +export interface S3StorageNodeProps extends BaseNodeProps { + /** + * The S3 bucket to use. + */ + readonly bucket: IBucket; + /** + * The content input to use. + */ + readonly contentInput: TypedFlowNodeInput; + /** + * The object key input to use. + */ + readonly objectKeyInput: UntypedFlowNodeInput; +} + +/****************************************************************************** + * AI + *****************************************************************************/ +// ------------------------------------------------------------------ +// LEX Node +// ------------------------------------------------------------------ +/** + * The properties for a Lex Node. + */ +export interface LexBotNodeProps extends BaseNodeProps { + /** + * The Lex bot alias to use. + */ + readonly botAliasArn: string; + /** + * The locale to use. + * @default "en_US" + */ + readonly localeId?: string; + /** + * The utterance to send to the bot. + */ + readonly inputText: TypedFlowNodeInput; + /** + * Any request attributes to send alongside the prompt. + * @see https://docs.aws.amazon.com/lexv2/latest/dg/context-mgmt-request-attribs.html + */ + readonly requestAttributes?: TypedFlowNodeInput; + /** + * Any session attributes to send alongside the prompt. + * @see https://docs.aws.amazon.com/lexv2/latest/dg/context-mgmt-session-attribs.html + */ + readonly sessionAttributes?: TypedFlowNodeInput; +} + +/****************************************************************************** + * LOGIC + *****************************************************************************/ +// ------------------------------------------------------------------ +// CONDITION Node +// ------------------------------------------------------------------ +/** + * The properties for a Condition Node. + */ +export interface ConditionNodeProps extends BaseNodeProps { + /** + * The inputs to use for the node. + */ + readonly inputs: NamedTypedFlowNodeInput[]; +} + +// ------------------------------------------------------------------ +// ITERATOR Node +// ------------------------------------------------------------------ +/** + * The properties for an Iterator Node. + */ +export interface IteratorNodeProps extends BaseNodeProps { + /** + * The configuration to use for the `array` input. + */ + readonly arrayInput: UntypedFlowNodeInput; + /** + * The array item data type to use. + * @default FlowNodeDataType.STRING + */ + readonly arrayItemDataType?: FlowNodeDataType; +} + +// ------------------------------------------------------------------ +// COLLECTOR Node +// ------------------------------------------------------------------ +/** + * The properties for a Collector Node. + */ +export interface CollectorNodeProps extends BaseNodeProps { + /** + * The configuration for the `arrayItem` input. + */ + readonly arrayItemInput: TypedFlowNodeInput; + + /** + * The configuration for the `arraySize` input. + */ + readonly arraySizeInput: UntypedFlowNodeInput; +} + +// ------------------------------------------------------------------ +// Condition Logic +// ------------------------------------------------------------------ +/** + * The configuration for a predefined condition (condition and name already set). + */ +export interface PredefinedCondition { + /** + * Where to go next in case the condition is satisfied + */ + readonly transitionTo: FlowNode; +} + +/** + * The configuration for a condition whose name has already been set. + */ +export interface FlowCondition extends PredefinedCondition { + /** + * The condition to use. + */ + readonly conditionExpression?: string; +} + +/** + * The configuration for a condition whose name, condition, and transition needs to be specified.. + */ +export interface NamedFlowCondition extends FlowCondition { + /** + * A name for the condition that you can reference. + */ + readonly name: string; +} diff --git a/src/cdk-lib/bedrock/prompt-flows/flow-nodes.ts b/src/cdk-lib/bedrock/prompt-flows/flow-nodes.ts new file mode 100644 index 00000000..627dcf9e --- /dev/null +++ b/src/cdk-lib/bedrock/prompt-flows/flow-nodes.ts @@ -0,0 +1,857 @@ +/** + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ + +import { Aws, aws_bedrock as bedrock, IResolvable, Lazy } from "aws-cdk-lib"; +import { CfnFlow } from "aws-cdk-lib/aws-bedrock"; +import { PolicyStatement } from "aws-cdk-lib/aws-iam"; +import { FlowConnection } from "./flow-connections"; +import { + AgentNodeProps, + CollectorNodeProps, + ConditionNodeProps, + FlowNodeDataType, + FlowNodeType, + InputDataNodeProps, + IteratorNodeProps, + KnowledgeBaseRetrieveAndGenerateNodeProps, + KnowledgeBaseRetrieveNodeProps, + LambdaFunctionNodeProps, + LexBotNodeProps, + NamedFlowCondition, + NamedFlowNodeOutput, + NamedTypedFlowNodeInput, + OutputDataNodeProps, + PredefinedCondition, + PromptNodeProps, + S3RetrievalNodeProps, + S3StorageNodeProps, +} from "./flow-node-props"; +import { PromptVariant } from "../prompts/prompt"; + +/****************************************************************************** + * CONSTRUCT PROPS + *****************************************************************************/ +/** + * Defines the properties to create a Flow Node + * Not exported as its use is internal-only. + */ +interface FlowNodeProps { + /** + * The name of the flow node. + */ + readonly name: string; + /** + * The configurations for the node. + */ + readonly configuration: bedrock.CfnFlow.FlowNodeConfigurationProperty; + /** + * The type of node. + */ + readonly type: FlowNodeType; + /** + * The inputs to the node. + */ + readonly inputs: NamedTypedFlowNodeInput[]; + /** + * The outputs of the node. + */ + readonly outputs: NamedFlowNodeOutput[]; + /** + * The needed IAM policy statements needed to add to a Prompt Flow Execution Role in order + * to use this node correctly within a Prompt Flow. + * @default No Statements + */ + readonly neededPolicyStatements?: PolicyStatement[]; + /** + * The data connections established for this node. + */ + readonly dataConnections?: FlowConnection[]; +} + +/****************************************************************************** + * CONSTRUCT + *****************************************************************************/ +/** + * Class to create a new managed Flow Node. + */ +export class FlowNode { + // ------------------------------------------------------ + // Basic Nodes + // ------------------------------------------------------ + /** + * Creates an input node. This can only be the first node in the flow. + * + * Outputs: + * - `document` (User-defined type) + */ + public static input(props: InputDataNodeProps): FlowNode { + return new FlowNode({ + name: props.name, + type: FlowNodeType.INPUT, + configuration: { input: {} }, + inputs: [], + outputs: [ + { + name: "document", + type: props.inputDataType, + }, + ], + }); + } + + /** + * Creates an output node. This will be the last node in the flow. + * You can specify multiple output nodes, one for every possible + * execution branch. + * + * Inputs: + * - `document` (User-defined type) + */ + public static output(props: OutputDataNodeProps): FlowNode { + return new FlowNode({ + name: props.name, + type: FlowNodeType.OUTPUT, + configuration: { output: {} }, + inputs: [ + { + name: "document", + type: props.outputData?.type, + valueFrom: { + sourceNode: props.outputData.valueFrom.sourceNode, + sourceOutputName: props.outputData.valueFrom.sourceOutputName, + expression: props.outputData.valueFrom.expression, + }, + }, + ], + outputs: [], + }); + } + + // ------------------------------------------------------ + // Code Nodes + // ------------------------------------------------------ + /** + * Creates a Lambda Function node which lets you call a Lambda function + * in which you can define code to carry out business logic. + * + * Inputs: + * - User defined inputs. + * + * Outputs: + * - `functionResponse` (User-defined type) + */ + public static lambdaFunction(props: LambdaFunctionNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.LAMBDA, + configuration: { + lambdaFunction: { + lambdaArn: props.lambdaFunction.functionArn, + }, + }, + inputs: props.inputs, + outputs: [ + { + name: "functionResponse", + type: props.functionResponseType, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["lambda:InvokeFunction"], + resources: [props.lambdaFunction.functionArn, `${props.lambdaFunction.functionArn}:*`], + conditions: { + StringEquals: { + "aws:ResourceAccount": Aws.ACCOUNT_ID, + }, + }, + }), + ], + }); + } + + // ------------------------------------------------------ + // Orchestration Nodes + // ------------------------------------------------------ + /** + * Creates an Agent node that allows you to run an agent. + * + * Inputs: + * - `agentInputText` (String) + * - `promptAttributes` (Optional - User defined - defaults to Object) + * - `sessionAttributes` (Optional - User defined - defaults to Object) + * + * Outputs: + * - `agentResponse` (String) + */ + public static agent(props: AgentNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.AGENT, + configuration: { + agent: { + agentAliasArn: props.agentAlias.aliasArn, + }, + }, + inputs: [ + { + name: "agentInputText", + type: props.agentInput?.type, + valueFrom: { + sourceNode: props.agentInput.valueFrom.sourceNode, + sourceOutputName: props.agentInput.valueFrom.sourceOutputName, + expression: props.agentInput.valueFrom.expression, + }, + }, + ...(props.promptAttributes + ? [ + { + name: "promptAttributes", + type: props.promptAttributes.type ?? FlowNodeDataType.OBJECT, + valueFrom: props.promptAttributes.valueFrom, + }, + ] + : []), + ...(props.sessionAttributes + ? [ + { + name: "sessionAttributes", + type: props.sessionAttributes.type ?? FlowNodeDataType.OBJECT, + valueFrom: props.sessionAttributes.valueFrom, + }, + ] + : []), + ], + outputs: [ + { + name: "agentResponse", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["bedrock:InvokeAgent"], + resources: [props.agentAlias.aliasArn], + }), + ], + }); + } + + /** + * Creates a prompt node that allows you to run a prompt. + * + * Inputs: + * - User defined inputs. The amount of inputs must be equal to the defined variables in the prompt. + * + * Outputs: + * - `modelCompletion` (String) + */ + public static prompt(props: PromptNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.PROMPT, + configuration: { + prompt: { + sourceConfiguration: { + resource: { + promptArn: props.prompt.promptArn, + }, + }, + }, + }, + inputs: props.inputs ?? [], + outputs: [ + { + name: "modelCompletion", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["bedrock:GetPrompt"], + resources: [props.prompt.promptArn], + }), + // For those variants that have model ID, grant ability to invoke model + ...props.prompt.variants + .filter((variant: PromptVariant) => variant.modelId) + .map( + (variant: PromptVariant) => + new PolicyStatement({ + actions: ["bedrock:InvokeModel"], + resources: [variant.modelId!], + }) + ), + ], + }); + } + + // ------------------------------------------------------ + // Data Nodes + // ------------------------------------------------------ + /** + * Creates a Knowledge base node that returns retrieved results. + * + * Inputs: + * - `retrievalQuery` (String) + * + * Outputs: + * - `retrievalResults` (Array) + */ + public static knowledgeBaseRetrieve(props: KnowledgeBaseRetrieveNodeProps): FlowNode { + return new FlowNode({ + name: props.name, + type: FlowNodeType.KNOWLEDGE_BASE, + configuration: { + knowledgeBase: { + knowledgeBaseId: props.knowledgeBase.knowledgeBaseId, + }, + }, + inputs: [ + { + name: "retrievalQuery", + type: FlowNodeDataType.STRING, + valueFrom: props.retrievalQuery.valueFrom, + }, + ], + outputs: [ + { + name: "retrievalResults", + type: FlowNodeDataType.ARRAY, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["bedrock:Retrieve"], + resources: [props.knowledgeBase.knowledgeBaseArn], + }), + ], + }); + } + + /** + * Creates a Knowledge base node that returns generates a response based on + * retrieved results. + * + * Inputs: + * - `retrievalQuery` (String) + * + * Outputs: + * - `outputText` (String) + */ + public static knowledgeBaseRetrieveAndGenerate(props: KnowledgeBaseRetrieveAndGenerateNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.KNOWLEDGE_BASE, + configuration: { + knowledgeBase: { + knowledgeBaseId: props.knowledgeBase.knowledgeBaseId, + modelId: props.model.modelId, + }, + }, + inputs: [ + { + name: "retrievalQuery", + type: FlowNodeDataType.STRING, + valueFrom: props.retrievalQuery.valueFrom, + }, + ], + outputs: [ + { + name: "outputText", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["bedrock:RetrieveAndGenerate"], + resources: [props.knowledgeBase.knowledgeBaseArn], + }), + ], + }); + } + + /** + * Creates an S3 retrieval node. + * + * Inputs: + * - `objectKey` (String) + * + * Outputs: + * - `s3Content` (String) + */ + public static s3Retrieval(props: S3RetrievalNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.RETRIEVAL, + configuration: { + retrieval: { + serviceConfiguration: { + s3: { + bucketName: props.bucket.bucketName, + }, + }, + }, + }, + inputs: [ + { + name: "objectKey", + type: FlowNodeDataType.STRING, + valueFrom: props.objectKeyInput.valueFrom, + }, + ], + outputs: [ + { + name: "s3Content", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["s3:GetObject"], + resources: [props.bucket.arnForObjects("*")], + conditions: { + StringEquals: { + "aws:ResourceAccount": Aws.ACCOUNT_ID, + }, + }, + }), + ], + }); + } + + /** + * Creates an S3 storage node. + * + * Inputs: + * - `content` (User-defined type) + * - `objectKey` (String) + * + * Outputs: + * - `s3Uri` (String) + */ + public static s3Storage(props: S3StorageNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.STORAGE, + configuration: { + storage: { + serviceConfiguration: { + s3: { + bucketName: props.bucket.bucketName, + }, + }, + }, + }, + inputs: [ + { + name: "content", + type: props.contentInput.type, + valueFrom: props.contentInput.valueFrom, + }, + { + name: "objectKey", + type: FlowNodeDataType.STRING, + valueFrom: props.objectKeyInput.valueFrom, + }, + ], + outputs: [ + { + name: "s3Uri", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["s3:PutObject"], + resources: [props.bucket.bucketArn, props.bucket.arnForObjects("*")], + conditions: { + StringEquals: { + "aws:ResourceAccount": Aws.ACCOUNT_ID, + }, + }, + }), + ], + }); + } + // ------------------------------------------------------ + // AI Nodes + // ------------------------------------------------------ + /** + * Creates a Lex Node + * + * Inputs: + * - `content` (User-defined type) + * - `objectKey` (String) + * + * Outputs: + * - `s3Uri` (String) + */ + public static lexBot(props: LexBotNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.LEX, + configuration: { + lex: { + botAliasArn: props.botAliasArn, + localeId: props.localeId ?? "en_US", + }, + }, + inputs: [ + { + name: "inputText", + type: props.inputText?.type ?? FlowNodeDataType.STRING, + valueFrom: props.inputText.valueFrom, + }, + ...(props.requestAttributes + ? [ + { + name: "requestAttributes", + type: props.requestAttributes?.type ?? FlowNodeDataType.OBJECT, + valueFrom: props.requestAttributes.valueFrom, + }, + ] + : []), + ...(props.sessionAttributes + ? [ + { + name: "sessionAttributes", + type: props.sessionAttributes?.type ?? FlowNodeDataType.OBJECT, + valueFrom: props.sessionAttributes.valueFrom, + }, + ] + : []), + ], + outputs: [ + { + name: "predictedIntent", + type: FlowNodeDataType.STRING, + }, + ], + neededPolicyStatements: [ + new PolicyStatement({ + actions: ["lex:RecognizeUtterance"], + resources: [props.botAliasArn], + conditions: { + StringEquals: { + "aws:ResourceAccount": Aws.ACCOUNT_ID, + }, + }, + }), + ], + }); + } + + // ------------------------------------------------------ + // Logic Nodes + // ------------------------------------------------------ + /** + * Creates a condition node that allows you to conditionally execute a flow. + */ + public static conditionNode(props: ConditionNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.CONDITION, + configuration: {}, + inputs: props.inputs, + outputs: [], + }); + } + + /** + * Creates an iterator node that allows you to iterate over an array. + * + * This node takes an input array and iteratively sends each item of the array + * as an output to the following node in the flow. It also returns the size of + * the array in the output. + * + * The output flow node at the end of the flow iteration will return a response + * for each member of the array. If you want to return only one consolidated + * response, you can include a collector node downstream from this iterator node. + * + * Inputs: + * - `array` (Array) + * + * Outputs: + * - `arrayItem` (User-defined type - defaults to String) + * - `arraySize` (Number) + */ + public static iteratorNode(props: IteratorNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.ITERATOR, + configuration: { iterator: undefined }, + inputs: [ + { + name: "array", + type: FlowNodeDataType.ARRAY, + valueFrom: props.arrayInput.valueFrom, + }, + ], + outputs: [ + { + name: "arrayItem", + type: props.arrayItemDataType ?? FlowNodeDataType.STRING, + }, + { + name: "arraySize", + type: FlowNodeDataType.NUMBER, + }, + ], + }); + } + + /** + * Creates a collector node that allows you to collect results from an iterator. + * Inputs: + * - `arraySize` (Number) + * - `arrayItem` (User-defined type - defaults to String) + * + * Outputs: + * - `collectedArray` (Array) + */ + public static collectorNode(props: CollectorNodeProps) { + return new FlowNode({ + name: props.name, + type: FlowNodeType.COLLECTOR, + configuration: { + collector: undefined, + }, + inputs: [ + { + name: "arraySize", + type: FlowNodeDataType.NUMBER, + valueFrom: props.arraySizeInput.valueFrom, + }, + { + name: "arrayItem", + type: props.arrayItemInput.type ?? FlowNodeDataType.STRING, + valueFrom: props.arrayItemInput.valueFrom, + }, + ], + outputs: [ + { + name: "collectedArray", + type: FlowNodeDataType.ARRAY, + }, + ], + }); + } + // ------------------------------------------------------ + // Properties + // ------------------------------------------------------ + /** + * The name of the flow node. + */ + public readonly name: string; + /** + * The type of node. + */ + public readonly type: FlowNodeType; + /** + * The inputs to the node. + */ + public readonly inputs: NamedTypedFlowNodeInput[]; + /** + * The outputs of the node. + */ + public readonly outputs: NamedFlowNodeOutput[]; + /** + * The needed policy statements in order to use this node correctly within + * a Prompt Flow. + */ + public readonly neededPolicyStatements?: PolicyStatement[]; + /** + * The configurations for the node. + */ + public readonly connections: FlowConnection[]; + /** + * The configurations for the node. + */ + protected configuration: bedrock.CfnFlow.FlowNodeConfigurationProperty; + /** + * The configurations for the node. + */ + protected readonly conditions: NamedFlowCondition[]; + + // ------------------------------------------------------ + // Constructor + // ------------------------------------------------------ + protected constructor(props: FlowNodeProps) { + this.name = props.name; + this.configuration = props.configuration; + this.type = props.type; + this.inputs = props.inputs; + this.outputs = props.outputs; + this.neededPolicyStatements = props.neededPolicyStatements; + this.connections = this._computeConnections(); + this.conditions = []; + } + + // ------------------------------------------------------ + // Methods + // ------------------------------------------------------ + // /** + // * Returns the CloudFormation representation of this Flow node. + // */ + asNodeCfnProperty(): bedrock.CfnFlow.FlowNodeProperty | IResolvable { + if (this.type == FlowNodeType.CONDITION) { + return { + name: this.name, + configuration: { + condition: { + conditions: this._computeConditions(), + }, + }, + type: this.type, + inputs: this.inputs?.flatMap((item) => { + return { + name: item.name, + type: item.type ?? FlowNodeDataType.STRING, + expression: item.valueFrom.expression ?? "$.data", + } as CfnFlow.FlowNodeInputProperty; + }), + outputs: this.outputs?.flatMap((item) => { + return { + name: item.name, + type: item.type ?? FlowNodeDataType.STRING, + } as CfnFlow.FlowNodeOutputProperty; + }), + }; + } else { + return { + name: this.name, + configuration: this.configuration, + type: this.type, + inputs: this.inputs?.flatMap((item) => { + return { + name: item.name, + type: item.type ?? FlowNodeDataType.STRING, + expression: item.valueFrom.expression ?? "$.data", + } as CfnFlow.FlowNodeInputProperty; + }), + outputs: this.outputs?.flatMap((item) => { + return { + name: item.name, + type: item.type ?? FlowNodeDataType.STRING, + } as CfnFlow.FlowNodeOutputProperty; + }), + }; + } + } + + /** + * + * @internal + */ + _computeConnections() { + return this.inputs?.flatMap((input) => { + return FlowConnection.data({ + name: `${this.name}_${input.name}`, + source: input.valueFrom.sourceNode, + target: this, + config: { + sourceOutput: input.valueFrom.sourceOutputName ?? input.valueFrom.sourceNode.outputs[0].name, + targetInput: input.name, + }, + }); + }); + } + /** + * + * @internal + */ + _computeDataConnections(): FlowConnection[] { + return this.inputs?.flatMap((input) => { + return FlowConnection.data({ + name: `${this.name}_${input.name}`, + source: input.valueFrom.sourceNode, + target: this, + config: { + sourceOutput: input.valueFrom.sourceOutputName ?? input.valueFrom.sourceNode.outputs[0].name, + targetInput: input.name, + }, + }); + }); + } + + /** + * + * @internal + */ + _computeConditionalConnections(): FlowConnection[] { + if (this.type != FlowNodeType.CONDITION) { + return []; + } else { + if (this.conditions.length < 2) { + throw new Error( + "Condition nodes must have configured at least a condition, and a default transition. Use the appropriate methods" + ); + } else { + return this.conditions.flatMap((condition) => { + return FlowConnection.conditional({ + name: `${this.name}_${condition.name}`, + source: this, + target: condition.transitionTo, + condition: condition.conditionExpression, + }); + }); + } + } + } + + /** + * + * @internal + */ + _computeConditions(): IResolvable { + if (this.conditions.length < 2) { + throw new Error( + "Condition nodes must have configured at least a condition, and a default transition. Use the appropriate methods" + ); + } else { + return Lazy.any( + { + produce: () => { + return this.conditions.flatMap((item: NamedFlowCondition) => { + return { + name: item.name, + expression: item.conditionExpression, + } as bedrock.CfnFlow.FlowConditionProperty; + }); + }, + }, + { omitEmptyArray: true } + ); + } + } + + addCondition(config: NamedFlowCondition) { + if (this.type != FlowNodeType.CONDITION) { + throw new Error("Only condition nodes can have conditions"); + } else { + this.conditions.push(config); + this.connections.push( + FlowConnection.conditional({ + name: `${this.name}_${config.name}`, + source: this, + target: config.transitionTo, + condition: config.name, + }) + ); + } + } + + addDefaultTransition(config: PredefinedCondition) { + this.addCondition({ name: "default", ...config }); + } +} diff --git a/src/cdk-lib/bedrock/prompt-flows/flow-version.ts b/src/cdk-lib/bedrock/prompt-flows/flow-version.ts new file mode 100644 index 00000000..768f9e67 --- /dev/null +++ b/src/cdk-lib/bedrock/prompt-flows/flow-version.ts @@ -0,0 +1,56 @@ +/** + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ + +import { Construct } from "constructs"; +import * as bedrock from "aws-cdk-lib/aws-bedrock"; +import { Flow } from "./flow"; + +export interface FlowVersionProps { + /** + * The flow that this version belongs to. + */ + readonly flow: Flow; + /** + * A description of the flow version + */ + readonly description?: string; +} + +/** + * Creates a version of the flow that you can deploy. + */ +export class FlowVersion extends Construct { + /** + * The version of the flow. + */ + public readonly version: string; + + /** + * The flow that this version belongs to. + */ + public readonly flow: Flow; + + private readonly _resource: bedrock.CfnFlowVersion; + constructor(scope: Construct, id: string, props: FlowVersionProps) { + super(scope, id); + + this.flow = props.flow; + + this._resource = new bedrock.CfnFlowVersion(this, `FlowVersion-${this.flow._hash.slice(0, 16)}`, { + description: props.description, + flowArn: props.flow.flowArn, + }); + + this.version = this._resource.attrVersion; + } +} diff --git a/src/cdk-lib/bedrock/prompt-flows/flow.ts b/src/cdk-lib/bedrock/prompt-flows/flow.ts new file mode 100644 index 00000000..7f25051f --- /dev/null +++ b/src/cdk-lib/bedrock/prompt-flows/flow.ts @@ -0,0 +1,283 @@ +/** + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance + * with the License. A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES + * OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions + * and limitations under the License. + */ + +import { Construct } from "constructs"; +import { ArnFormat, aws_iam as iam, IResource, Resource, Stack } from "aws-cdk-lib"; +import { IKey } from "aws-cdk-lib/aws-kms"; +import { aws_bedrock as bedrock } from "aws-cdk-lib"; +import * as fs from "fs"; +import { FlowNode } from "./flow-nodes"; +import { md5hash } from "aws-cdk-lib/core/lib/helpers-internal"; +import { FlowVersion } from "./flow-version"; + +export enum ModelType { + CUSTOM = "custom-model", + FOUNDATIONAL = "foundation-model", + PROVISIONED = "provisioned-model", +} + +/****************************************************************************** + * COMMON + *****************************************************************************/ +/** + * Represents a Prompt Flow, either created with CDK or imported. + */ +export interface IFlow extends IResource { + /** + * The service role used by the flow. + */ + readonly executionRole: iam.IRole; + /** + * The ARN of the prompt flow. + * @example "arn:aws:bedrock:us-east-1:123456789012:flow/EUTIGM37LX" + */ + readonly flowArn: string; + /** + * The ID of the prompt flow. + * @example "EUTIGM37LX" + */ + readonly flowId: string; + /** + * The version of the flow. + */ + flowVersion: string; + /** + * Method to create a Flow Version. + */ + createVersion(description: string): string; +} + +/** + * Represents a Prompt Flow Definition. + */ +export class FlowDefinition { + /** + * The Amazon S3 location of the JSON flow definition. Appropriate Role permissions + * need to be manually added from the console or with a custom role in CDK. + */ + public static fromS3(location: bedrock.CfnFlow.S3LocationProperty) { + return new FlowDefinition({ definitionS3Location: location }); + } + + /** + * The definition of the flow must be provided as a JSON-formatted string. + * Note that the string must match the format in + * [FlowDefinition](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-bedrock-flow-flowdefinition.html) . + */ + public static fromNodes(nodes: FlowNode[]) { + return new FlowDefinition({ + definition: { + nodes: nodes.map((item) => item.asNodeCfnProperty()), + connections: nodes.flatMap((item) => item.connections.map((conn) => conn.asCfnProperty())), + }, + nodes: nodes, + }); + } + + /** + * The definition of the flow provided as a local JSON file. + * Note that the file must match the format in + * [FlowDefinition](https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-bedrock-flow-flowdefinition.html) . + */ + public static fromFile(path: string) { + const definition = fs.readFileSync(path, "utf8"); + return new FlowDefinition({ definitionString: definition }); + } + + public readonly definition?: bedrock.CfnFlow.FlowDefinitionProperty; + public readonly definitionS3Location?: bedrock.CfnFlow.S3LocationProperty; + public readonly definitionString?: string; + public readonly nodes?: FlowNode[]; + + protected constructor(config: { + definition?: bedrock.CfnFlow.FlowDefinitionProperty; + definitionS3Location?: bedrock.CfnFlow.S3LocationProperty; + definitionString?: string; + nodes?: FlowNode[]; + }) { + this.definition = config.definition; + this.definitionS3Location = config.definitionS3Location; + this.definitionString = config.definitionString; + this.nodes = config.nodes; + } +} + +/****************************************************************************** + * PROPS FOR NEW CONSTRUCT + *****************************************************************************/ +export interface FlowProps { + /** + * The name of the prompt flow. + */ + readonly name: string; + /** + * The prompt flow content definition. + */ + readonly definition: FlowDefinition; + /** + * A map that specifies the mappings for placeholder variables in the prompt flow definition. + */ + readonly substitutions?: { [key: string]: any }; + /** + * The description of what the prompt flow does. + * @default - No description provided. + */ + readonly description?: string; + /** + * The KMS key that the prompt is encrypted with. + * @default - AWS owned and managed key + */ + readonly encryptionKey?: IKey; +} + +/****************************************************************************** + * NEW CONSTRUCT DEFINITION + *****************************************************************************/ +/** + * Creates a prompt flow that you can use to send an input through various + * steps to yield an output. + * @see https://docs.aws.amazon.com/bedrock/latest/userguide/flows.html + * @resource AWS::Bedrock::Flow + */ +export class Flow extends Resource implements IFlow { + /** + * The service role used by the flow. + */ + public readonly executionRole: iam.Role; + /** + * The ARN of the prompt flow. + * @example "arn:aws:bedrock:us-east-1:123456789012:flow/EUTIGM37LX" + */ + public readonly flowArn: string; + /** + * The ID of the prompt flow. + * @example "EUTIGM37LX" + */ + public readonly flowId: string; + /** + * The version of the prompt flow. + * Defaults to "DRAFT" if no explicit version is created. + */ + public flowVersion: string; + /** + * The computed hash of the flow properties. + * @internal + */ + public readonly _hash: string; + /** + * The L1 Flow resource + */ + private readonly _resource: bedrock.CfnFlow; + + constructor(scope: Construct, id: string, props: FlowProps) { + super(scope, id); + + // ------------------------------------------------------ + // Execution Role + // ------------------------------------------------------ + // Create base execution role + this.executionRole = new iam.Role(this, "AmazonBedrockExecutionRoleForFlows_", { + assumedBy: new iam.ServicePrincipal("bedrock.amazonaws.com").withConditions({ + StringEquals: { + "aws:SourceAccount": Stack.of(this).account, + }, + ArnLike: { + "aws:SourceArn": [ + Stack.of(this).formatArn({ + service: "bedrock", + resource: "flow", + resourceName: "*", + arnFormat: ArnFormat.SLASH_RESOURCE_NAME, + }), + ], + }, + }), + }); + + // ------------------------------------------------------ + // CFN Props + // ------------------------------------------------------ + const { definition, definitionS3Location, definitionString } = props.definition; + const cfnProps: bedrock.CfnFlowProps = { + name: props.name, + executionRoleArn: this.executionRole.roleArn, + description: props.description, + definition, + definitionS3Location, + definitionString, + }; + + this._hash = md5hash(JSON.stringify(cfnProps)); + + // ------------------------------------------------------ + // L1 Instance + // ------------------------------------------------------ + this._resource = new bedrock.CfnFlow(this, "Resource", cfnProps); + + this.flowArn = this._resource.attrArn; + this.flowId = this._resource.attrId; + this.flowVersion = this._resource.attrVersion; + + // ------------------------------------------------------ + // Add appropriate permissions to the role + // ------------------------------------------------------ + // Grant Get Flow Permissions + this.executionRole.addToPrincipalPolicy( + new iam.PolicyStatement({ + actions: ["bedrock:GetFlow"], + resources: [this.flowArn], + }) + ); + + // If data must be encrypted with custom KMS key, add appropriate permissions + if (props.encryptionKey) { + props.encryptionKey.grantEncryptDecrypt(this.executionRole); + } + + // If definition is provided as objects, inspect it, and add appropriate + // permissions to the execution role + if (props?.definition) { + for (const node of props.definition.nodes!) { + if (node.neededPolicyStatements) { + for (const statement of node.neededPolicyStatements) { + this.executionRole.addToPrincipalPolicy(statement); + } + } + } + } + } + + // // ------------------------------------------------------ + // // Define via Method + // // ------------------------------------------------------ + // public fromDefinition(definition: FlowDefinition){ + // this. + // } + // ------------------------------------------------------ + // Create Version + // ------------------------------------------------------ + /** + * Create a version for the guardrail. + * @param description The description of the version. + * @returns The guardrail version. + */ + public createVersion(description?: string): string { + const flowVersion = new FlowVersion(this, `FlowVersion-${this._hash.slice(0, 16)}`, { + flow: this, + description: description, + }); + + this.flowVersion = flowVersion.version; + return this.flowVersion; + } +}