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

pip install boto3 langchain

import logging
import json
import boto3
from botocore.exceptions import ClientError
from datetime import datetime
from botocore.client import Config
from langchain.prompts import PromptTemplate
import json

bedrock_client = boto3.client(service_name='bedrock-runtime')

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

tool_config = {
"tools": [
{
"toolSpec": {
"name": "pto_balance",
"description": "Get the PTO balance for an employee.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"employee_id": {
"type": "string",
"description": "The employee ID for which you want to retrieve the PTO balance."
}
},
"required": [
"employee_id"
]
}
}
}
},
{
"toolSpec": {
"name": "submit_pto",
"description": "Submit a PTO request for an employee.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"employee_id": {
"type": "string",
"description": "The employee ID for which you want to submit a PTO request."
},
"start_date": {
"type": "string",
"description": "The start date of the PTO request in the format YYYY-MM-DD."
},
"end_date": {
"type": "string",
"description": "The end date of the PTO request in the format YYYY-MM-DD."
},
"reason": {
"type": "string",
"description": "The reason for the PTO request."
}
},
"required": [
"employee_id",
"start_date",
"end_date",
"reason"
]
}
}
}
},
{
"toolSpec": {
"name": "hr_policy",
"description": "Get information about a specific HR policy.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"question": {
"type": "string",
"description": "The HR policy question to query the knowledge base."
}
},
"required": [
"question"
]
}
}
}
}
]
}

def retrieve(query, kbId, numberOfResults=3):
bedrock_config = Config(connect_timeout=120, read_timeout=120, retries={'max_attempts': 0})

bedrock_agent_client = boto3.client("bedrock-agent-runtime",config=bedrock_config)
return bedrock_agent_client.retrieve(
retrievalQuery= {
'text': query
},
knowledgeBaseId=kbId,
retrievalConfiguration= {
'vectorSearchConfiguration': {
'numberOfResults': numberOfResults,
'overrideSearchType': "HYBRID"
}
}
)

def get_contexts(retrievalResults):
contexts = []
#print(retrievalResults)
for retrievedResult in retrievalResults:
text = retrievedResult['content']['text']
# Remove the "Document 1: " prefix if it exists
if text.startswith("Document 1: "):
text = text[len("Document 1: "):]
contexts.append(text)
contexts_string = ', '.join(contexts)
return contexts_string

def get_pto_balance(employee_id):
"""Returns the PTO balance for the given employee ID."""
pto_balance = 0
if employee_id == "E1234":
pto_balance = 10
elif employee_id == "E5678":
pto_balance = 5
else:
raise ValueError(f"Invalid employee ID: {employee_id}")
return pto_balance

def submit_pto_request(employee_id, start_date, end_date, reason):
"""Submits a PTO request for the given employee ID, start date, end date, and reason."""
pto_balance = get_pto_balance(employee_id)
start_date = datetime.strptime(start_date, "%Y-%m-%d")
end_date = datetime.strptime(end_date, "%Y-%m-%d")
days = (end_date - start_date).days + 1
if days > pto_balance:
raise InsufficientPTOBalanceError(f"Insufficient PTO balance for employee {employee_id}.")
# Simulating PTO request submission
return f"PTO request submitted for employee {employee_id} from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')} for {reason}."

class InsufficientPTOBalanceError(Exception):
"""Raised when an employee doesn't have enough PTO balance."""
pass

def get_hr_policy(question):
response_retrieve = retrieve(question, "<Your_KnowledgeBase_ID>")["retrievalResults"]
contexts = get_contexts(response_retrieve)
PROMPT_TEMPLATE = """DOCUMENT:
{context}
QUESTION:
{message}
INSTRUCTIONS:
Answer the user's QUESTION using only the DOCUMENT text above.
Keep your answer strictly grounded in the facts provided. Do not refer to the "DOCUMENT," "documents," "provided text," ,"based on.." or any similar phrases in your answer.
If the provided text contains the facts to answer the QUESTION, include all relevant details in your answer.
If the provided text doesn’t contain the facts to answer the QUESTION, respond only with "I don't know" and do not add any further information.
"""
prompt = PromptTemplate(template=PROMPT_TEMPLATE,
input_variables=["context","message"])
prompt_final = prompt.format(context=contexts,
message=question)
body = json.dumps({
"anthropic_version": "bedrock-2023-05-31",
"max_tokens": 1000,
"messages": [
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt_final
}
]
}
]
})
model_id = 'anthropic.claude-3-sonnet-20240229-v1:0'
accept = 'application/json'
content_type = 'application/json'
response = bedrock_client.invoke_model(body=body, modelId=model_id, accept=accept, contentType=content_type)
response_body = json.loads(response.get('body').read())
response_text = response_body.get('content')[0]['text']
return response_text

def generate_text(bedrock_client, model_id, tool_config, input_text):
"""Generates text using the supplied Amazon Bedrock model. If necessary,
the function handles tool use requests and sends the result to the model."""
logger.info("Generating text with model %s", model_id)

# Create the initial message from the user input.
messages = [{
"role": "user",
"content": [{"text": input_text}]
}]

response = bedrock_client.converse(
modelId=model_id,
messages=messages,
toolConfig=tool_config
)

output_message = response['output']['message']
messages.append(output_message)
stop_reason = response['stopReason']

if stop_reason == 'tool_use':
# Tool use requested. Call the tool and send the result to the model.
tool_requests = response['output']['message']['content']
print(tool_requests)
for tool_request in tool_requests:
if 'toolUse' in tool_request:
tool = tool_request['toolUse']
logger.info("Requesting tool %s. Request: %s",
tool['name'], tool['toolUseId'])

if tool['name'] == 'pto_balance':
tool_result = {}
try:
pto_balance = get_pto_balance(tool['input']['employee_id'])
tool_result = {
"toolUseId": tool['toolUseId'],
"content": [{"json": {"pto_balance": pto_balance}}]
}
except ValueError as err:
tool_result = {
"toolUseId": tool['toolUseId'],
"content": [{"text": err.args[0]}],
"status": 'error'
}

tool_result_message = {
"role": "user",
"content": [
{
"toolResult": tool_result
}
]
}
messages.append(tool_result_message)

elif tool['name'] == 'submit_pto':
tool_result = {}
try:
pto_request_result = submit_pto_request(
tool['input']['employee_id'],
tool['input']['start_date'],
tool['input']['end_date'],
tool['input']['reason']
)
tool_result = {
"toolUseId": tool['toolUseId'],
"content": [{"text": pto_request_result}]
}
except (ValueError, InsufficientPTOBalanceError) as err:
tool_result = {
"toolUseId": tool['toolUseId'],
"content": [{"text": err.args[0]}],
"status": 'error'
}

tool_result_message = {
"role": "user",
"content": [
{
"toolResult": tool_result
}
]
}
messages.append(tool_result_message)

elif tool['name'] == 'hr_policy':
tool_result = {}
policy = get_hr_policy(input_text)
policy = {
"text": policy
}

print(json.dumps(policy, indent=4))
return

# Send the tool result to the model.
response = bedrock_client.converse(
modelId=model_id,
messages=messages,
toolConfig=tool_config
)
print(response)
output_message = response['output']['message']

# Print the final response from the model.
for content in output_message['content']:
print(json.dumps(content, indent=4))

example_questions = [
"What is the PTO balance for employee E5678?",
"Submit a PTO request for employee E1234 from 2023-07-10 to 2023-07-12 for a family vacation.",
"What is the dress code policy?",
"Can you tell me how many PTO days employee E1234 has left?",
"I want to request PTO for employee E5678 starting from 2023-08-01 to 2023-08-05 for a personal trip.",
"What are the standard working hours?",
"I need to know the remaining PTO balance for employee E9012.",
"Employee E9012 needs to take PTO from 2023-09-15 to 2023-09-20 for a medical procedure.",
"What is the overtime policy?",
"What is the policy for vacation",
"What is the policy for work from home?"
]

for question in example_questions:
try:
print(f"Question: {question}")
generate_text(bedrock_client, model_id, tool_config, question)
print("\n")
except ClientError as err:
message = err.response['Error']['Message']
logger.error("A client error occurred: %s", message)
print(f"A client error occurred: {message}")

Question: What is the PTO balance for employee E5678?
INFO:__main__:Generating text with model anthropic.claude-3-sonnet-20240229-v1:0
[{'text': 'To get the PTO balance for an employee, we can use the "pto_balance" tool:'}, {'toolUse': {'toolUseId': 'tooluse_00kL9hLpQOelFoaQjbDSlw', 'name': 'pto_balance', 'input': {'employee_id': 'E5678'}}}]
INFO:__main__:Requesting tool pto_balance. Request: tooluse_00kL9hLpQOelFoaQjbDSlw
{
"text": "The PTO balance for employee E5678 is 5 days."
}

Question: Submit a PTO request for employee E1234 from 2023-07-10 to 2023-07-12 for a family vacation.
INFO:__main__:Generating text with model anthropic.claude-3-sonnet-20240229-v1:0
[{'text': 'Here is how to submit the PTO request using the tool:'}, {'toolUse': {'toolUseId': 'tooluse_ywjLWDjfSISVP91AfBAmag', 'name': 'submit_pto', 'input': {'employee_id': 'E1234', 'start_date': '2023-07-10', 'end_date': '2023-07-12', 'reason': 'family vacation'}}}]
INFO:__main__:Requesting tool submit_pto. Request: tooluse_ywjLWDjfSISVP91AfBAmag
{
"text": "The PTO request has been successfully submitted for the specified dates and reason."
}

Question: What is the dress code policy?
INFO:__main__:Generating text with model anthropic.claude-3-sonnet-20240229-v1:0
[{'text': 'To get information about the dress code policy, we can use the `hr_policy` tool like this:'}, {'toolUse': {'toolUseId': 'tooluse_ZGMnL4HSRfKlZ9S8wbSZoA', 'name': 'hr_policy', 'input': {'question': 'What is the dress code policy?'}}}]
INFO:__main__:Requesting tool hr_policy. Request: tooluse_ZGMnL4HSRfKlZ9S8wbSZoA
Dress Code Policy Purpose: This policy outlines the dress code requirements for employees to maintain a professional appearance and create a positive work environment. It applies to all employees, regardless of their position or department. Policy: 1.1 All employees are expected to dress in a manner appropriate for their job roles and responsibilities. Refer to the Job Descriptions for specific requirements. 1.2 Business casual attire is required for office-based employees, including collared shirts, dress slacks, skirts, and dresses. See the Business Casual Guidelines for more information. 1.3 Employees in customer-facing roles should wear company-provided uniforms or adhere to a more formal dress code. Consult the Uniform Policy for details. 1.4 Clothing with offensive or inappropriate language, graphics, or logos is strictly prohibited. Review the Inappropriate Attire Examples for guidance. 1.5 Supervisors have the authority to determine the appropriateness of an employee's attire and may send employees home to change if necessary. Refer to the Supervisor Responsibilities for more information. Compliance: Failure to comply with this dress code policy may result in disciplinary action, up to and including termination of employment. Progressive disciplinary steps are outlined in the Disciplinary Action Policy., Document 3: Social Media Policy Purpose: This policy provides guidelines for employees' use of social media platforms, both personally and professionally, to protect the company's reputation and confidential information. It applies to all employees, contractors, and temporary workers. Policy: 3.1 Employees are expected to conduct themselves professionally on social media and avoid posting content that may harm the company's reputation. Review the Social Media Best Practices for guidance. 3.2 Confidential company information, including client data, financial information, and intellectual property, should never be shared on social media. Refer to the Confidentiality Agreement for more details. 3.3 Employees should not post disparaging comments about the company, its employees, clients, or competitors. See the Anti-Harassment Policy for more information. 3.4 When posting content related to the company, employees must clearly state that their views are personal and do not necessarily reflect those of the company. Use the Social Media Disclaimer when necessary. 3.5 The company reserves the right to monitor employees' public social media activity and take appropriate action if necessary. Consult the Employee Monitoring Policy for details. Compliance: Violations of this social media policy may lead to disciplinary action, up to and including termination of employment. Employees should exercise good judgment and seek guidance from their supervisors or the HR department if unsure about the appropriateness of their social media activity. Refer to the Social Media FAQs for common questions and answers., Document 2: Vacation Leave Policy Purpose: This policy outlines the guidelines for employees to request and take vacation leave. It ensures fair and consistent application of vacation benefits across the organization. Policy: 2.1 Full-time employees are eligible for paid vacation leave based on their years of service with the company. See the Vacation Accrual Schedule for details. 2.2 Vacation leave accrues at a rate of X days per month, up to a maximum of Y days per year. Refer to the Vacation Accrual Calculator to determine your accrual rate. 2.3 Employees must submit vacation leave requests to their supervisors at least Z weeks in advance for approval. Use the Vacation Request Form to submit your request. 2.4 Vacation leave requests will be approved based on business needs and the order in which they are received. Check the Vacation Scheduling Guidelines for more information. 2.5 Unused vacation leave may be carried over to the next year, up to a maximum of W days. Consult the Vacation Carryover Policy for details. Compliance: Employees are encouraged to use their earned vacation leave to maintain a healthy work-life balance. Managers should work with their teams to ensure adequate coverage during employee vacations. Refer to the Manager's Guide to Vacation Planning for best practices.
{
"text": "All employees are expected to dress in a manner appropriate for their job roles and responsibilities. Business casual attire is required for office-based employees, including collared shirts, dress slacks, skirts, and dresses. Employees in customer-facing roles should wear company-provided uniforms or adhere to a more formal dress code. Clothing with offensive or inappropriate language, graphics, or logos is strictly prohibited. Supervisors have the authority to determine the appropriateness of an employee's attire and may send employees home to change if necessary."
}