AWS for Industries

Orchestrating Clinical Generative AI Workflows Using AWS Step Functions

In the rapidly evolving field of healthcare, leveraging generative AI has the potential to revolutionize patient care and streamline clinical processes. However, orchestrating these complex workflows while ensuring data security and compliance with regulations, such as HIPAA, can be challenging. A common pattern involves starting with the source data, choosing a generative AI model, and using prompts to perform specific tasks on each of the data records. In this blog post, we’ll explore how AWS Step Functions, combined with other AWS services, can be used to build a secure, extensible, and scalable solution for orchestrating clinical generative AI workflows.

Use case example

Figure 1: Architectural diagram

To demonstrate this architecture, we’ve built a healthcare-focused solution that combines AWS HealthLake, generative AI services like Amazon Bedrock, and other AWS components to summarize clinical notes. Built using the AWS Cloud Development Kit (CDK) for easy deployment and management, the solution is orchestrated by an AWS Step Functions State Machine called GenAIWorkflow. Let’s examine each step of this workflow below.

GenAIWorkflow state machine

Figure 2. GenAIWorkflow state machine

  1. Start by fetching query parameters from AWS Systems Manager Parameter Store. Parameters include the Amazon Athena query that retrieves a list of record Ids to process, and AWS HealthLake datastore details.
  2. Calls Amazon Athena StartQueryExecution to execute the Athena query.
  3. Amazon Athena GetQueryResults retrieves the results of the query, which is a list of record Ids. Depending on the query and result set, a maximum number of results are returned. If there are additional results, a NextToken is returned, indicating subsequent queries need to be made in order to retrieve all records.
  4. A Map state in Distributed mode iterates through the list of returned record Ids, and process them concurrently. By default, Step Functions uses a concurrency value of 10.
  5. For reach record Id, we call Amazon Athena StartQueryExecution API to start running the query.
  6. This step gets the query results, which is the report content associated with the record Id.
  7. Using a Choice State, this step determines if any results are returned before continuing with the processing logic.
  8. This is an optional step to use an AWS Lambda function to perform any necessary formatting of the Athena query results, e.g. combining values from multiple columns into one value, or creating a JSON record from query results.
  9. Retrieves the prompt template for the specific task from the Parameter Store. This prompt should be customized to your use-case and should also be optimized for the model you wish to invoke.
  10. This step creates the prompt by combining the prompt template and report content, and then calls the Amazon Bedrock API. Retry mechanisms can be configured here to catch different types of Amazon Bedrock API invocation exceptions and retry the API.
  11. Another Lambda function is invoked to parse the Amazon Bedrock API response and convert the output data into a format that can be saved in an Amazon DynamoDB table.
  12. Saves the structured output data in a DynamoDB table using the PutItem API.
  13. If no results are returned, skip the iteration.
  14. A Choice state is used to determine the presence of a NextToken value.
  15. If NextToken is present, we call the Amazon Athena GetQueryResults API again with the NextToken to fetch the next batch of results.
  16. The workflow succeeds when there are no more records for processing.

Benefits

Using AWS Step Functions to orchestrate generative AI workflows offers several advantages. At its core, Step Functions provides a visual workflow orchestration tool that seamlessly integrates with various AWS services, making it ideal for managing complex clinical generative AI workflows. Security is paramount in this solution, particularly with Step Functions recently released support for Customer Managed Keys to encrypt State Machine and Activity resources, enabling you to encrypt your workflow definitions and execution data using your own encryption keys. This enhanced security extends to robust handling of Protected Health Information (PHI) through AWS KMS Customer Managed Keys (CMKs) and comprehensive encryption both in transit and at rest across components. The architecture leverages secure data processing and storage through services like Amazon Athena, AWS Lambda, and Amazon DynamoDB. The solution is also highly extensible, storing SQL queries and generative AI prompts as parameters in AWS Systems Manager Parameter Store, which enables modifications without code changes. Additionally, Lambda functions can be updated to accommodate use-case specific logic before or after the generative AI step.

This solution also excels in operational reliability and performance. Step Functions provides built-in error handling and retry mechanisms, ensuring workflow resilience through automatic retries or graceful failure handling, with the option to use redrive functionality to restart execution from failed steps. The screenshot below shows the visual editor within Step Functions and the robust error handling capabilities it provides. The architecture’s use of serverless technologies like Lambda and Athena enables automatic scaling based on workload, optimizing both performance and cost-effectiveness by consuming resources only when needed. Finally, comprehensive monitoring and logging capabilities are built into Step Functions, allowing you to track workflow progress and examine the input and output data at each step of the process.

Error handling for Step FunctionsFigure 3. Error handling for Step Functions

Next Steps

By leveraging AWS Step Functions and other AWS services, we can build secure, extensible, and scalable solutions for orchestrating clinical generative AI workflows. This empowers healthcare organizations to harness the potential of generative AI while ensuring the highest standards of data protection and compliance. We encourage you to explore the solution, adapt it to your specific needs, and unlock new possibilities in patient care and clinical processes. We value community collaboration and welcome your contributions – whether through submitting a pull request to enhance the solution’s functionality, reporting any issues you encounter via GitHub issues, or sharing your experiences and suggestions in the comments section below. Your feedback and involvement will help us continue to improve and strengthen this solution to better serve the healthcare community’s needs.

Qing Liu

Qing Liu

Qing Liu is a Senior Solution Architect at AWS. Qing has more than 10 years of experience working in healthcare IT industry. He is passionate about using healthcare data to drive better insights and improve patient outcomes. In his spare time, he likes to play tennis with his wife and friends.

Nick Ragusa

Nick Ragusa

Nick Ragusa is a Principal Solutions Architect on the AWS Education team, where he helps educational institutions and Academic Medical Centers harness cloud technology to transform both learning experiences and patient outcomes. When he's not architecting cloud solutions, you'll find him cheering from the sidelines at his three children's sporting events, logging thoughtful miles on crisp autumn morning runs, or satisfying his sweet tooth with a carefully chosen dessert.