Skip to content

Instantly share code, notes, and snippets.

@mrmaheshrajput
Last active August 15, 2024 11:32
Show Gist options
  • Save mrmaheshrajput/794d530c58d9ecadf832b383f8bd1241 to your computer and use it in GitHub Desktop.
Save mrmaheshrajput/794d530c58d9ecadf832b383f8bd1241 to your computer and use it in GitHub Desktop.
from datetime import datetime
import json
from typing import Any, Dict, List
import boto3
from botocore.exceptions import ClientError
# Initialize a Boto3 session and create a Bedrock runtime client
session = boto3.Session()
region = "us-east-1" # us-west-2 has better runtime quota
bedrock_client = session.client(service_name = 'bedrock-runtime', region_name = region)
# Define available models with their respective request limits
available_models = {
"sonnet3-5": "anthropic.claude-3-5-sonnet-20240620-v1:0", # 50 requests per min
"sonnet": "anthropic.claude-3-sonnet-20240229-v1:0", # 500 requests per min
"llama31-70b": "meta.llama3-1-70b-instruct-v1:0", # 400 requests per min
"llama31-405b": "meta.llama3-1-405b-instruct-v1:0", # 50 requests per min
}
modelId = available_models["sonnet3-5"] # Select model for conversation
class FakeDatabase:
"""Sample fake database implementation."""
def __init__(self):
self.customers = [
{"id": "1213210", "name": "John Doe", "email": "john@gmail.com", "phone": "123-456-7890", "username": "johndoe"},
{"id": "2837622", "name": "Priya Patel", "email": "priya@candy.com", "phone": "987-654-3210", "username": "priya123"},
{"id": "3924156", "name": "Liam Nguyen", "email": "lnguyen@yahoo.com", "phone": "555-123-4567", "username": "liamn"},
{"id": "4782901", "name": "Aaliyah Davis", "email": "aaliyahd@hotmail.com", "phone": "111-222-3333", "username": "adavis"},
{"id": "5190753", "name": "Hiroshi Nakamura", "email": "hiroshi@gmail.com", "phone": "444-555-6666", "username": "hiroshin"},
{"id": "6824095", "name": "Fatima Ahmed", "email": "fatimaa@outlook.com", "phone": "777-888-9999", "username": "fatimaahmed"},
{"id": "7135680", "name": "Alejandro Rodriguez", "email": "arodriguez@protonmail.com", "phone": "222-333-4444", "username": "alexr"},
{"id": "8259147", "name": "Megan Anderson", "email": "megana@gmail.com", "phone": "666-777-8888", "username": "manderson"},
{"id": "9603481", "name": "Kwame Osei", "email": "kwameo@yahoo.com", "phone": "999-000-1111", "username": "kwameo"},
{"id": "1057426", "name": "Mei Lin", "email": "meilin@gmail.com", "phone": "333-444-5555", "username": "mlin"}
]
self.orders = [
{"id": "24601", "customer_id": "1213210", "product": "Wireless Headphones", "quantity": 1, "price": 79.99, "status": "Shipped"},
{"id": "13579", "customer_id": "1213210", "product": "Smartphone Case", "quantity": 2, "price": 19.99, "status": "Processing"},
{"id": "97531", "customer_id": "2837622", "product": "Bluetooth Speaker", "quantity": 1, "price": "49.99", "status": "Shipped"},
{"id": "86420", "customer_id": "3924156", "product": "Fitness Tracker", "quantity": 1, "price": 129.99, "status": "Delivered"},
{"id": "54321", "customer_id": "4782901", "product": "Laptop Sleeve", "quantity": 3, "price": 24.99, "status": "Shipped"},
{"id": "19283", "customer_id": "5190753", "product": "Wireless Mouse", "quantity": 1, "price": 34.99, "status": "Processing"},
{"id": "74651", "customer_id": "6824095", "product": "Gaming Keyboard", "quantity": 1, "price": 89.99, "status": "Delivered"},
{"id": "30298", "customer_id": "7135680", "product": "Portable Charger", "quantity": 2, "price": 29.99, "status": "Shipped"},
{"id": "47652", "customer_id": "8259147", "product": "Smartwatch", "quantity": 1, "price": 199.99, "status": "Processing"},
{"id": "61984", "customer_id": "9603481", "product": "Noise-Cancelling Headphones", "quantity": 1, "price": 149.99, "status": "Shipped"},
{"id": "58243", "customer_id": "1057426", "product": "Wireless Earbuds", "quantity": 2, "price": 99.99, "status": "Delivered"},
{"id": "90357", "customer_id": "1213210", "product": "Smartphone Case", "quantity": 1, "price": 19.99, "status": "Shipped"},
{"id": "28164", "customer_id": "2837622", "product": "Wireless Headphones", "quantity": 2, "price": 79.99, "status": "Processing"}
]
def get_user(self, key:str, value:str) -> Dict[str, str]:
"""Return metadata of user."""
if key in {"email", "phone", "username"}:
for customer in self.customers:
if customer[key] == value:
return customer
return f"Couldn't find a user with {key} of {value}"
else:
raise ValueError(f"Invalid key: {key}")
return None
def get_order_by_id(self, order_id: str) -> Dict[str, str]:
"""Return metadata of the order using order id."""
for order in self.orders:
if order["id"] == order_id:
return order
return None
def get_customer_orders(self, customer_id: str) -> List[Dict[str, str]]:
"""Return a list of orders for a specific customer."""
return [order for order in self.orders if order["customer_id"] == customer_id]
def cancel_order(self, order_id: str) -> str:
"""Cancel an order if it's in 'Processing' status."""
order = self.get_order_by_id(order_id)
if order:
if order["status"] == "Processing":
order["status"] = "Cancelled"
return "Cancelled the order"
else:
return "Order has already shipped. Can't cancel it."
return "Can't find that order!"
# Define all the tools avilable to the model
tool_config = {
"tools": [
{
"toolSpec": {
"name": "get_user",
"description": "Looks up a user by email, phone, or username.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"key": {
"type": "string",
"enum": ["email", "phone", "username"],
"description": "The attribute to search for a user by (email, phone, or username).",
},
"value": {
"type": "string",
"description": "The value to match for the specified attribute.",
},
},
"required": ["key", "value"],
}
},
}
},
{
"toolSpec": {
"name": "get_order_by_id",
"description": "Retrieves the details of a specific order based on the order ID. Returns the order ID, product name, quantity, price, and order status.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"order_id": {
"type": "string",
"description": "The unique identifier for the order.",
}
},
"required": ["order_id"],
}
},
}
},
{
"toolSpec": {
"name": "get_customer_orders",
"description": "Retrieves the list of orders belonging to a user based on a user's customer id.",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"customer_id": {
"type": "string",
"description": "The customer_id belonging to the user",
}
},
"required": ["customer_id"],
}
},
}
},
{
"toolSpec": {
"name": "cancel_order",
"description": "Cancels an order based on a provided order_id. Only orders that are 'processing' can be cancelled",
"inputSchema": {
"json": {
"type": "object",
"properties": {
"order_id": {
"type": "string",
"description": "The order_id pertaining to a particular order",
}
},
"required": ["order_id"],
}
},
}
},
],
"toolChoice": {"auto": {}},
}
db = FakeDatabase()
def process_tool_call(tool_name: str, tool_input: Any) -> Any:
"""Process the tool call based on the tool name and input."""
if tool_name == "get_user":
return db.get_user(tool_input["key"], tool_input["value"])
elif tool_name == "get_order_by_id":
return db.get_order_by_id(tool_input["order_id"])
elif tool_name == "get_customer_orders":
return db.get_customer_orders(tool_input["customer_id"])
elif tool_name == "cancel_order":
return db.cancel_order(tool_input["order_id"])
def simple_chat():
"""Main chat function that interacts with the user and the LLM."""
system_prompt = """
You are a customer support chat bot for an online retailer called TechNova.
Your job is to help users look up their account, orders, and cancel orders.
Be helpful and brief in your responses.
You have access to a set of tools, but only use them when needed.
If you do not have enough information to use a tool correctly, ask a user follow up questions to get the required inputs.
Do not call any of the tools unless you have the required data from a user.
"""
# Initial user message
user_message = input("\nUser: ")
messages = [{"role": "user", "content": [{"text": user_message}]}]
while True:
# If the last message is from the assistant, get another input from the user
if messages[-1].get("role") == "assistant":
user_message = input("\nUser: ")
messages.append({"role": "user", "content": [{"text": user_message}]})
# Parameters for API request to the Bedrock model
converse_api_params = {
"modelId": modelId,
"system": [{"text": system_prompt}],
"messages": messages,
"inferenceConfig": {"maxTokens": 4096},
"toolConfig": tool_config, # Pass the tool config
}
# Get response from Bedrock model
response = bedrock_client.converse(**converse_api_params)
# Append assistant's message to the conversation
messages.append(
{"role": "assistant", "content": response["output"]["message"]["content"]}
)
# If the model wants to use a tool, process the tool call
if response["stopReason"] == "tool_use":
tool_use = response["output"]["message"]["content"][
-1
] # Naive approach assumes only 1 tool is called at a time
tool_id = tool_use["toolUse"]["toolUseId"]
tool_name = tool_use["toolUse"]["name"]
tool_input = tool_use["toolUse"]["input"]
print(f"Claude wants to use the {tool_name} tool")
print(f"Tool Input:")
print(json.dumps(tool_input, indent=2))
# Run the underlying tool functionality on the fake database
tool_result = process_tool_call(tool_name, tool_input)
print(f"\nTool Result:")
print(json.dumps(tool_result, indent=2))
# Append tool result message
messages.append(
{
"role": "user",
"content": [
{
"toolResult": {
"toolUseId": tool_id,
"content": [{"text": str(tool_result)}],
}
}
],
}
)
else:
# If the model does not want to use a tool, just print the text response
print(
"\nTechNova Support:"
+ f"{response['output']['message']['content'][0]['text']}"
)
if __name__ == "__main__":
simple_chat()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment