Last active
August 15, 2024 11:32
-
-
Save mrmaheshrajput/794d530c58d9ecadf832b383f8bd1241 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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