Amazon Bedrock: Enhance HR Support with Function Calling & Knowledge Bases

You can use the Amazon Bedrock API to give a model access to tools that can help it generate responses for messages that you send to the model. For example, you might have an HR chat application that lets employees find out information about their PTO balance, submit PTO requests, or query the knowledge base for HR policies. To answer a request for the PTO balance or to submit a PTO request, the model needs tools that can query the HR system and return the relevant information or perform the necessary actions. Similarly, to provide information about HR policies, the model needs a tool that can search the knowledge base and retrieve the relevant policy details. Tool use with models is also known as Function calling.

Please visit these great introduction by Jason Stehle to Getting started with the Amazon Bedrock Converse API and Intro to Tool Use with the Amazon Bedrock Converse API in previous posts.

This tutorial demonstrates how to use Amazon Bedrock with the Converse API , function calling and knowledge bases for a sample HR use case. The sample code shows how to retrieve information from a knowledge base, submit PTO requests, and query HR policies using a conversational AI model.

Even though this is an HR use case, you can easily extend this to other use cases very easily like customer support etc.

Key Concepts

In Amazon Bedrock, models don't directly call tools. Instead, when sending a message to a model, you also provide tool definitions that could help generate a response. If the model needs a tool, it responds with a request to call the tool, including input parameters. You then call the tool (e.g., an API, database, or Lambda function) and continue the conversation by providing the tool's result. Finally, the model generates a response using the tool results you sent.

To use tools with a model, you can use the Converse API (Converse or ConverseStream). While it's possible to use tools with base inference operations (InvokeModel or InvokeModelWithResponseStream), the Converse API is recommended for consistency across all Amazon Bedrock models that support tool use. Refer here for more detailed information

In this HR use case example, the code defines three tools:

pto_balance: Retrieves the PTO balance for an employee.
submit_pto: Submits a PTO request for an employee.
hr_policy: Retrieves information about a specific HR policy using Knowledge Bases for Amazon Bedrock

The model can request these tools when necessary to generate responses to HR-related queries.

Note: This is demo code for illustrative purposes only. Not intended for production use.

Prerequisites

Python 3.x
boto3 library
botocore library
langchain library
Select an AWS region that supports the Anthropic Claude 3 Sonnet model. I'm using us-west-2 (Oregon). You can check the documentation for model support by region.
Configure Amazon Bedrock model access for your account and region. Example here: https://catalog.workshops.aws/building-with-amazon-bedrock/en-US/prerequisites/bedrock-setup
If you don’t have your own local integrated development environment, you can try AWS Cloud9. Setup instructions here: https://catalog.workshops.aws/building-with-amazon-bedrock/en-US/prerequisites/cloud9-setup

Note : You can install required libraries using pip install

Code Overview

Import and initialize

We will also set up a logger object and configures the logging level to INFO for the current module

Tool Configuration

The tool_config dictionary defines the available tools for the model. Each tool has a name, description, and input schema. In this example, the following tools are defined:

pto_balance: Retrieves the PTO balance for an employee.
submit_pto: Submits a PTO request for an employee.
hr_policy: Retrieves information about a specific HR policy.

The main components of the code are:

retrieve(query, kbId, numberOfResults): Retrieves relevant documents from a knowledge base using the provided query.

For more information on how to create your knowledge base using Amazon Bedrock, please refer to the documentation. Here I have used some sample HR policy documents to sync in Amazon S3 data source.

You can also visit our aws samples for knowledge bases to get started with boto3 and its use with knowledge bases for bedrock

get_contexts(retrievalResults): Extracts the context from the retrieved documents for further processing. This can be changed based on your requirements.

get_pto_balance(employee_id): Retrieves the PTO balance for a given employee ID.

submit_pto_request(employee_id, start_date, end_date, reason): Submits a PTO request for an employee.

The InsufficientPTOBalanceError class is a custom exception that is raised when an employee tries to submit a PTO request but doesn't have enough PTO balance to cover the requested time off.

get_hr_policy(question): Retrieves information about a specific HR policy based on the provided question from Knowledge Bases for Amazon Bedrock and then pass to Amazon Bedrock to generate a response which can be returned as is without further processing.

generate_text(bedrock_client, model_id, tool_config, input_text): Generates text using the Amazon Bedrock model and handles tool use requests.

Handling Tool Requests

The generate_text function sends the user's input to the model using the converse API. If the model responds with a tool use request (stop_reason == 'tool_use'), the function handles the request based on the tool name:

For pto_balance, it calls the get_pto_balance function with the provided employee ID and returns the PTO balance.
For submit_pto, it calls the submit_pto_request function with the provided employee ID, start date, end date, and reason, and returns the result.
For hr_policy, it calls the get_hr_policy function with the user's input and returns the retrieved policy information from the knowledge base.

After handling the tool request, the function sends the tool result back to the model using the converse API, allowing the model to generate a final response.

Usage

Update the example_questions list with your own HR-related questions. Below is a sample set

Run the script :

Output : You can see based on user's input appropriate tools use is called and executed

Customization

Modify the get_pto_balance() and submit_pto_request() functions to integrate with your own HR systems and databases.
Update the get_hr_policy() function to retrieve information from your own HR knowledge base.
Adjust the tool_config dictionary to add, remove, or modify the available tools based on your specific HR use case.

Conclusion

This tutorial showcased how to use Amazon Bedrock with the Converse API and Knowledge bases for Amazon Bedrock for an HR use case. By leveraging conversational AI, tool use (function calling), and integrating with HR tools and knowledge bases, you can create a powerful system for handling HR-related queries and tasks. The sample code example demonstrates how to configure tools, handle tool requests, and generate responses using the Amazon Bedrock model.

Happy Building !

Any opinions in this post are those of the individual author and may not reflect the opinions of AWS.

Site Terms, Privacy, and more.