(bedrock): implement prompt flows

[aws-rafams]

Describe the feature

Prompt flows for Amazon Bedrock offers the ability for you to use supported foundation models (FMs) to build workflows by linking prompts, foundational models, and other AWS services to create end-to-end solutions.
https://aws.amazon.com/bedrock/prompt-flows/

Use Case

Amazon Bedrock Prompt Flows accelerates the creation, testing, and deployment of workflows.

Proposed Solution

Implement a new module prompt.ts and prompt-flow.ts
Cloudformation Resource is already available: https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-bedrock-flow.html
https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-bedrock-prompt.html
CDK L1 should be available soon, we can already start talking about implementation details

Other Information

Can start working on the implementation right away as soon as L1 is released

Acknowledgements

• I may be able to implement this feature request
• This feature might incur a breaking change

[krokoko]

Thanks @aws-rafams ! As you mentioned, the feature is not exposed yet through the cdk, so it cannot be implemented now.
In terms of implementation, I would see:

• flow.ts, flow-version.ts and flow-alias.ts following the resources available in cfn. Maybe additional files for flow-nodes.ts and connections, conditions
• For default values, we can follow the console experience
• Helper functions to facilitate the creation of the different types of nodes as described here: https://docs.aws.amazon.com/bedrock/latest/userguide/flows-nodes.html
• Validation methods to detect configuration issues at synth time

[aws-rafams]

CDK L1s released today in 2.154.0, starting implementation

[krokoko]

Just FYI, this new feature requires the library to be updated to cdk v2.154.0, however I just found out that there is a bug in that version (see here: aws/aws-cdk#31183) which currently prevents the upgrade. The library won't build as a property was removed in cfnDataSource.
For your local development, you can still use v2.154.0 and comment the data source part that is causing issues

[aws-rafams]

Thanks for the heads up! Currently doing local development with v2.154.1.

BTW, there's a bug in the CloudFormation resource for AWS::Bedrock::Flow->FlowNode.Type, which is currently delaying the implementation. Some node types that are available via the API and console, such as Retrieval, Storage, and Agent, are not listed in the CloudFormation resource.

I've created an issue to track this problem. I would really appreciate a +1 or a comment to the issue to help raise its visibility and priority.

[krokoko]

Created a separate ticket (#664) for prompts/prompt management since the issue above is still open

[aws-rafams]

CloudFormation bug has been fixed, resuming implementation. https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-bedrock-flow-flownode.html#cfn-bedrock-flow-flownode-type

[aristideubertas]

Hi, any updates on the development of this implementation?

[Marcin-Duszynski]

@aws-rafams any updates?

[aws-rafams]

Hey everyone, sorry for not updating this thread. I am working on it but I have pause development a bit. Since Prompt Management and Prompt Flows are in Preview, their CloudFormation Specification is changing quite frequently and is hard to keep up with the changes and there are some bugs as well. For instance in Prompt Management, TopK was removed from the CloudFormation specification (seems like a bug to me), but this is impeding me to launch test apps I have already created.

Hopefully we see GA soon and the specification is more stable.

On another side, I'm also waiting on the refactor of Agents (#731) to come through as the structure will certainly impact how Agents are used within Prompt Flows.

Anyhow, I'll try to create a Draft PR as soon as possible so you can participate on the discussion and comment on the interface to instantiate nodes and a prompt flow.

[aws-rafams]

What are your thoughts on the current structure being developed?

import * as cdk from "aws-cdk-lib";
import { Construct } from "constructs";
import { PythonFunction } from "@aws-cdk/aws-lambda-python-alpha";
import { NodejsFunction } from "aws-cdk-lib/aws-lambda-nodejs";
import { FlowNode, FlowNodeDataType } from "../lib/flow-nodes";
import { Architecture, Runtime } from "aws-cdk-lib/aws-lambda";
import * as path from "path";
import { Prompt, PromptVariant } from "../lib/prompt";
import * as fs from "fs";
import { AttributeType, Table } from "aws-cdk-lib/aws-dynamodb";
import { bedrock, amazonaurora } from "@cdklabs/generative-ai-cdk-constructs";
import { Bucket } from "aws-cdk-lib/aws-s3";
import * as s3deploy from "aws-cdk-lib/aws-s3-deployment";
import { Flow, FlowDefinition } from "../lib/flow";

export class MainStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    /**************************************************************************
     *                              Resources
     *************************************************************************/
    // ----------------------------------------------------
    // DDB Tables
    // ----------------------------------------------------
    const propertyListingTable = new Table(this, "PropertyListingTable", {
      partitionKey: { name: "mls_id", type: AttributeType.STRING },
    });

    const loanTable = new Table(this, "LoanTable", {
      partitionKey: { name: "mls_id", type: AttributeType.STRING },
    });

    // ----------------------------------------------------
    // Data Loader
    // ----------------------------------------------------
    // const tablePopulatorFunction = new NodejsFunction(this, "TablePopulatorFunction", {
    //   description: "Lambda function to initially load the claims catalog.",
    //   runtime: Runtime.NODEJS_18_X,
    //   entry: path.join(__dirname, "lambda", "data-loader"),
    //   handler: "index.lambda_handler",
    //   timeout: cdk.Duration.seconds(60),
    //   architecture: Architecture.X86_64,
    //   environment: {
    //     PROPERTY_TABLE_NAME: propertyListingTable.tableName,
    //     LOAN_TABLE_NAME: loanTable.tableName,
    //   },
    //   depsLockFilePath: path.join(__dirname, "lambda", "data-loader", "package-lock.json"),
    // });

    // ----------------------------------------------------
    // Permissions
    // ----------------------------------------------------
    // loanTable.grantFullAccess(tablePopulatorFunction);
    // propertyListingTable.grantFullAccess(tablePopulatorFunction);

    // ----------------------------------------------------
    // Create Vector Store
    // ----------------------------------------------------
    const auroraDb = new amazonaurora.AmazonAuroraVectorStore(this, "AuroraPromptFlows", {
      embeddingsModelVectorDimension: bedrock.BedrockFoundationModel.COHERE_EMBED_MULTILINGUAL_V3.vectorDimensions!,
    });

    // ----------------------------------------------------
    // Create Managed RAG Architecture
    // ----------------------------------------------------
    const kb = new bedrock.KnowledgeBase(this, "KnowledgeBase", {
      vectorStore: auroraDb,
      embeddingsModel: bedrock.BedrockFoundationModel.TITAN_EMBED_TEXT_V2_1024,
      instruction:
        "Contains a curated list of FAQs, the Selling Guide. It establishes and communicates the rules of the road for eligible borrowers, loans, and processes to uphold loan quality.",
    });

    // ----------------------------------------------------
    // Add S3 DataSource
    // ----------------------------------------------------
    // Bucket to hold documents
    const infoBucket = new Bucket(this, "FAQBucket", {
      removalPolicy: cdk.RemovalPolicy.DESTROY, // Remove the bucket when the stack is destroyed
      autoDeleteObjects: true,
    });

    // Configure S3 as KB Data Source
    new bedrock.S3DataSource(this, "S3DataSource", {
      bucket: infoBucket,
      knowledgeBase: kb,
      dataSourceName: "info",
      chunkingStrategy: bedrock.ChunkingStrategy.DEFAULT,
    });

    // Prepopulate S3 bucket
    new s3deploy.BucketDeployment(this, "FAQFiles", {
      sources: [s3deploy.Source.asset(path.join(__dirname, "documents"))],
      destinationBucket: infoBucket,
    });
    /**************************************************************************
     *                        ACTION GROUP - LOAN
     *************************************************************************/
    // ----------------------------------------------------
    // Lambda Function
    // ----------------------------------------------------
    const agentLoanCalculatorFunction = new PythonFunction(this, "AgentLoanCalculatorFunction", {
      description: "Bedrock Insurance Agent Loan Affordability Calculator",
      runtime: Runtime.PYTHON_3_12,
      entry: path.join(__dirname, "lambda", "loan", "agent"),
      handler: "lambda_handler",
      timeout: cdk.Duration.seconds(30),
      architecture: Architecture.X86_64,
    });

    // ----------------------------------------------------
    // Action Group
    // ----------------------------------------------------
    const loanAffordabilityCalculatorAG = new bedrock.AgentActionGroup(this, "LoanAffordability", {
      actionGroupName: "loan-calculation",
      description:
        "Implements a calculator that helps users estimate how much they can afford to borrow based on their income and expenses.",
      actionGroupExecutor: {
        lambda: agentLoanCalculatorFunction,
      },
      actionGroupState: "ENABLED",
      skipResourceInUseCheckOnDelete: true,
      apiSchema: bedrock.ApiSchema.fromAsset(path.join(__dirname, "api-schemas", "loan_calculator.json")),
    });

    /**************************************************************************
     *                        ACTION GROUP - MLS
     *************************************************************************/
    // ----------------------------------------------------
    // Lambda Function
    // ----------------------------------------------------
    const mlsFunction = new PythonFunction(this, "MLSFunction", {
      description: "Bedrock Agent MLS Lookup function",
      runtime: Runtime.PYTHON_3_12,
      entry: path.join(__dirname, "lambda", "loan", "flow"),
      handler: "lambda_handler",
      timeout: cdk.Duration.seconds(30),
      environment: {
        PROPERTY_TABLE_NAME: propertyListingTable.tableName,
      },
      architecture: Architecture.X86_64,
    });

    // ----------------------------------------------------
    // Permissions
    // ----------------------------------------------------
    propertyListingTable.grantFullAccess(mlsFunction.grantPrincipal);

    // ----------------------------------------------------
    // Action Group
    // ----------------------------------------------------
    const mlsPropertyLookupAG = new bedrock.AgentActionGroup(this, "MLSLookup", {
      actionGroupName: "mls-lookup",
      description:
        "Implements a calculator that helps users estimate how much they can afford to borrow based on their income and expenses.",
      actionGroupExecutor: {
        lambda: mlsFunction,
      },
      actionGroupState: "ENABLED",
      skipResourceInUseCheckOnDelete: true,
      apiSchema: bedrock.ApiSchema.fromAsset(path.join(__dirname, "api-schemas", "mls_lookup.json")),
    });

    /**************************************************************************
     *                              AGENT
     *************************************************************************/
    const agent = new bedrock.Agent(this, "Agent", {
      name: "mortgage-processing-agent",
      description: "An agent to process mortgage applications",
      foundationModel: bedrock.BedrockFoundationModel.ANTHROPIC_CLAUDE_SONNET_V1_0,
      instruction:
        "You are a customer service agent who has access to knowledge about mortgage products and services. You can help customers apply for a mortgage and answer questions about loan terms, interest rates, and mortgage eligibility. You can guide customers through the steps to submit documents or get appraisals completed. You can explain refinance and modification options to customers and provide resources on mortgage assistance programs. You can also answer internal questions about loan underwriting process, credit requirements, and guidelines for mortgage servicers and lenders. Your goal is to provide excellent service to customers and help them through the homebuying and mortgage financing process. Always ask follow-up question to get general information required before giving the user an answer.",
    });

    agent.addActionGroups([loanAffordabilityCalculatorAG, mlsPropertyLookupAG]);

    /**************************************************************************
     *                              PROMPTS
     *************************************************************************/

    const rejectionPrompt = new Prompt(this, "RejectionLetterPrompt", {
      promptName: "rejection-prompt",
      description:
        "Use this prompt to genererate a rejection letter triggered by an unsatisfactory income to debt ratio",
      variants: [
        PromptVariant.text({
          variantName: "default",
          model: bedrock.BedrockFoundationModel.ANTHROPIC_CLAUDE_HAIKU_V1_0.asIModel(this),
          promptText: fs.readFileSync(path.join(__dirname, "prompts", "rejection-letter.txt")).toString(),
          promptVariables: ["income", "totalDebt", "loanAmount"],
          inferenceConfiguration: {
            temperature: 0,
            topP: 0.999,
          },
        }),
      ],
    });

    const processApplicationPrompt = new Prompt(this, "ProcessApplicationPrompt", {
      promptName: "process-application-prompt",
      description: "Use this prompt to genererate a question for an agent to process the mortgage application",
      variants: [
        PromptVariant.text({
          variantName: "default",
          model: bedrock.BedrockFoundationModel.ANTHROPIC_CLAUDE_SONNET_V1_0.asIModel(this),
          promptText: fs.readFileSync(path.join(__dirname, "prompts", "process-application.txt")).toString(),
          promptVariables: ["income", "creditScore", "totalDebt", "loanAmount", "mlsId"],
          inferenceConfiguration: {
            temperature: 0,
            topP: 0.999,
          },
        }),
      ],
    });

    // /**************************************************************************
    //  *                              NODES
    //  *************************************************************************/
    const inputNode = FlowNode.input({
      name: "input",
      inputDataType: FlowNodeDataType.OBJECT,
    });

    const loanCalculatorFunction = new PythonFunction(this, "LoanCalculatorFunction", {
      description: "Loan Affordability Calculator",
      runtime: Runtime.PYTHON_3_12,
      entry: path.join(__dirname, "lambda", "loan", "flow"),
      handler: "lambda_handler",
      timeout: cdk.Duration.seconds(30),
      architecture: Architecture.X86_64,
    });

    const loanCalculatorNode = FlowNode.lambdaFunction({
      name: "loan_calculator",
      inputs: [
        {
          name: "income",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.income",
          },
        },
        {
          name: "totalDebt",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.totalDebt",
          },
        },
        {
          name: "loanTerm",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.loanTerm",
          },
        },
      ],
      outputType: FlowNodeDataType.NUMBER,
      lambdaFunction: loanCalculatorFunction,
    });

    const conditionNode = FlowNode.conditionNode({
      name: "approvalCondition",
      inputs: [
        {
          name: "loanAmount",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.loanAmount",
          },
        },
        {
          name: "maximumAffordableLoan",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: loanCalculatorNode,
            expression: "$.data",
          },
        },
      ],
    });

    const rejectionPromptNode = FlowNode.prompt({
      name: "rejection_prompt_node",
      prompt: rejectionPrompt,
      inputs: [
        {
          name: "income",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.income",
          },
        },
        {
          name: "totalDebt",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.totalDebt",
          },
        },
        {
          name: "loanAmount",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.loanAmount",
          },
        },
      ],
    });

    const processApplicationPromptNode = FlowNode.prompt({
      name: "processApplication",
      inputs: [
        {
          name: "income",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.income",
          },
        },
        {
          name: "totalDebt",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.totalDebt",
          },
        },
        {
          name: "loanAmount",
          type: FlowNodeDataType.NUMBER,
          valueFrom: {
            sourceNode: inputNode,
            expression: "$.data.loanAmount",
          },
        },
      ],
      prompt: processApplicationPrompt,
    });

    conditionNode.addCondition({
      name: "can_afford_loan",
      conditionExpression: "loanAmount <= maximumAffordableLoan",
      transitionTo: processApplicationPromptNode,
    });

    const agentNode = FlowNode.agent({
      name: "mortgageProcessingAgent",
      agentAlias: {
        agentId: agent.agentId,
        aliasId: "TSTALIASID",
        aliasArn: cdk.Arn.format({
          service: "bedrock",
          region: cdk.Aws.REGION,
          account: cdk.Aws.ACCOUNT_ID,
          partition: cdk.Aws.PARTITION,
          resource: "agent-alias",
          resourceName: `${agent.agentId}/TSTALIASID`,
          arnFormat: cdk.ArnFormat.SLASH_RESOURCE_NAME,
        }),
      },
      agentInput: {
        type: FlowNodeDataType.STRING,
        valueFrom: {
          sourceNode: processApplicationPromptNode,
          expression: "$.data",
        },
      },
    });

    const processedOutputNode = FlowNode.output({
      name: "processed_output",
      outputData: {
        type: FlowNodeDataType.STRING,
        valueFrom: {
          sourceNode: agentNode,
        },
      },
    });

    conditionNode.addDefaultTransition({
      transitionTo: rejectionPromptNode,
    });

    const rejectedOutputNode = FlowNode.output({
      name: "reject_output",
      outputData: {
        type: FlowNodeDataType.STRING,
        valueFrom: {
          sourceNode: rejectionPromptNode,
        },
      },
    });

    const myFlow = new Flow(this, "MyFlow", {
      name: "mortgage-processing-flow",
      definition: FlowDefinition.fromNodes([
        inputNode,
        loanCalculatorNode,
        conditionNode,
        rejectionPromptNode,
        rejectedOutputNode,
        processApplicationPromptNode,
        agentNode,
        processedOutputNode,
      ]),
    });
  }
}

This code would generate this Prompt Flow:

agentAlias in an Agent Node will surely be simplified in the final release.
Cool thing is that it will manage role permissions as nodes are added/removed, not even the console is able to do this! (docs)