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.
Figure 2. GenAIWorkflow state machine
- 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.
- Calls Amazon Athena
StartQueryExecution
to execute the Athena query. - 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, aNextToken
is returned, indicating subsequent queries need to be made in order to retrieve all records. - 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.
- For reach record Id, we call Amazon Athena
StartQueryExecution
API to start running the query. - This step gets the query results, which is the report content associated with the record Id.
- Using a Choice State, this step determines if any results are returned before continuing with the processing logic.
- 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.
- 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.
- 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.
- 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.
- Saves the structured output data in a DynamoDB table using the
PutItem
API. - If no results are returned, skip the iteration.
- A
Choice state
is used to determine the presence of aNextToken
value. - If
NextToken
is present, we call the Amazon AthenaGetQueryResults
API again with theNextToken
to fetch the next batch of results. - 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.
Figure 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.