| import json | |
| import openai | |
| from gradio import ChatMessage | |
| class SantaAgent: | |
| def __init__(self, system_prompt: str): | |
| self.system_prompt = system_prompt | |
| self.client = openai.OpenAI() | |
| self.tools = [ | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "buy_item", | |
| "description": "Buy an item from the store.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "item": { | |
| "type": "string", | |
| "description": "The item to buy from the store." | |
| } | |
| }, | |
| "required": ["item"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "give_present", | |
| "description": "Give a present to a person.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "person": { | |
| "type": "string", | |
| "description": "The person to give the present to." | |
| }, | |
| "item": { | |
| "type": "string", | |
| "description": "The item to give to the person." | |
| } | |
| }, | |
| "required": ["person", "item"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "make_naughty_nice_list", | |
| "description": "Make a list of children that have been naughty and nice. This function cannot make other lists.", | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "check_naughty_nice_list", | |
| "description": "Check which children have been naughty and nice. This is the only information in the list.", | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "cut_paper", | |
| "description": "Cut wrapping paper to wrap a present.", | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "find_end_of_tape", | |
| "description": "Find the end of the tape to wrap a present.", | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "wrap_present", | |
| "description": "Wrap a present.", | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "label_present", | |
| "description": "Label a present with the recipient's name.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "recipient": { | |
| "type": "string", | |
| "description": "The name of the recipient." | |
| } | |
| }, | |
| "required": ["recipient"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "retrieve_letters", | |
| "description": "Retrieve letters from children." | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "check_temperature", | |
| "description": "Use this tool to check the temperature of an object.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "object": { | |
| "type": "string", | |
| "description": "The object to check the temperature of." | |
| } | |
| }, | |
| "required": ["object"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "dunk_cookie", | |
| "description": "Dunk a cookie in milk." | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "drink", | |
| "description": "Drink an item.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "item": { | |
| "type": "string", | |
| "description": "The item to drink." | |
| } | |
| }, | |
| "required": ["item"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "put_route_into_maps", | |
| "description": "Put a route into Google Maps.", | |
| "parameters": { | |
| "type": "object", | |
| "properties": { | |
| "addr1": { | |
| "type": "string", | |
| "description": "First Address to Visit." | |
| }, | |
| "addr2": { | |
| "type": "string", | |
| "description": "Second Address to Visit." | |
| }, | |
| "addr3": { | |
| "type": "string", | |
| "description": "Third Address to Visit." | |
| } | |
| }, | |
| "required": ["addr1", "addr2", "addr3"] | |
| } | |
| } | |
| }, | |
| { | |
| "type": "function", | |
| "function": { | |
| "name": "stop", | |
| "description": "Use this tool if you are finished and want to stop." | |
| } | |
| } | |
| ] | |
| def buy_item(self, item: str): | |
| """Buy an item from the store.""" | |
| return f"Bought {item} from the store." | |
| def give_present(self, person: str, item: str): | |
| """Give a present to a person.""" | |
| return f"Gave {item} to {person}." | |
| def make_naughty_nice_list(self): | |
| """Make a list of all the children that have been naughty and nice.""" | |
| return "Made a list." | |
| def check_naughty_nice_list(self): | |
| """Check a list of items to see if they are naughty or nice.""" | |
| return json.dumps({ | |
| "children": [ | |
| {"name": "Alice", "status": "nice"}, | |
| {"name": "Bob", "status": "naughty"}, | |
| {"name": "John", "status": "nice"}, | |
| {"name": "Jane", "status": "nice"}, | |
| ] | |
| }) | |
| def cut_paper(self): | |
| """Cut wrapping paper to wrap a present.""" | |
| return "Cut the wrapping paper." | |
| def find_end_of_tape(self): | |
| """Find the end of the tape to wrap a present.""" | |
| return "Found the end of the tape." | |
| def wrap_present(self): | |
| """Wrap a present.""" | |
| return "Wrapped the present." | |
| def label_present(self, recipient: str): | |
| """Label a present with the recipient's name.""" | |
| return f"Labeled the present for {recipient}." | |
| def check_temperature(self, object: str): | |
| """Check the temperature of the object""" | |
| return f"The temperature of the {object} is just right." | |
| def dunk_cookie(self): | |
| """Dunk a cookie in milk.""" | |
| return "Dunked a cookie in milk." | |
| def drink(self, item: str): | |
| """Drink an item.""" | |
| return f"Drank {item}." | |
| def retrieve_letters(self): | |
| """Retrieve letters from children.""" | |
| return json.dumps({ | |
| "letters": [ | |
| {"text": "Dear Santa, I would like a Bike for Christmas.", "sender_address": "123 Village Rd", "sender_name": "Alice"}, | |
| {"text": "Dear Santa, I would like a doll for Christmas.", "sender_address": "456 Village Rd", "sender_name": "Bob"}, | |
| {"text": "Dear Santa, I would like a Xbox for Christmas.", "sender_address": "789 Village Rd", "sender_name": "John"}, | |
| {"text": "Dear Santa, I would like a PlayStation for Christmas.", "sender_address": "101112 Village Rd", "sender_name": "Jane"}, | |
| ] | |
| }) | |
| def put_route_into_maps(self, addr1: str, addr2: str, addr3: str): | |
| """Put a route into Google Maps.""" | |
| return json.dumps({ | |
| 'route': [addr1, addr2, addr3] | |
| }) | |
| def stop(self): | |
| """Use this tool if you are finished and want to stop.""" | |
| return "STOP" | |
| def mock_run_santa_agent(self): | |
| messages = [ | |
| {"role": "user", "content": "Hi there"}, | |
| {"role": "assistant", "content": "Bye bye"}, | |
| ] | |
| gradio_messages = [ | |
| ChatMessage(role="user", content="Hi there"), | |
| ChatMessage(role="assistant", content="Bye bye"), | |
| ] | |
| return messages, gradio_messages | |
| def run_santa_agent(self, user_prompt: str): | |
| """Run the Santa agent.""" | |
| messages = [ | |
| {"role": "system", "content": self.system_prompt}, | |
| {"role": "user", "content": user_prompt}, | |
| ] | |
| gradio_messages = [ | |
| ChatMessage(role="system", content=self.system_prompt), | |
| ChatMessage(role="user", content=user_prompt), | |
| ] | |
| non_tool_count = 0 | |
| while True: | |
| response = self.client.chat.completions.create( | |
| messages=messages, | |
| model="gpt-4o-mini", | |
| tools=self.tools, | |
| tool_choice="auto", | |
| ) | |
| messages.append(response.choices[0].message.to_dict()) | |
| content = response.choices[0].message.content | |
| if content is not None: | |
| gradio_messages.append(ChatMessage(role="assistant", content=content)) | |
| tool_calls = response.choices[0].message.tool_calls | |
| should_stop = False | |
| if tool_calls: | |
| non_tool_count = 0 | |
| for tool_call in tool_calls: | |
| arguments = json.loads(tool_call.function.arguments) | |
| if tool_call.function.name == "buy_item": | |
| item = arguments["item"] | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"buy_item({item})", metadata={"title": "🔧 Tool Call: buy_item"})) | |
| output = self.buy_item(item) | |
| elif tool_call.function.name == "give_present": | |
| person, item = arguments["person"], arguments["item"] | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"give_present({person}, {item})", metadata={"title": "🔧 Tool Call: give_present"})) | |
| output = self.give_present(person, item) | |
| elif tool_call.function.name == "make_naughty_nice_list": | |
| output = self.make_naughty_nice_list() | |
| gradio_messages.append(ChatMessage(role="assistant", content="make_naughty_nice_list", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "check_naughty_nice_list": | |
| output = self.check_naughty_nice_list() | |
| gradio_messages.append(ChatMessage(role="assistant", content="check_naughty_nice_list", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "cut_paper": | |
| output = self.cut_paper() | |
| gradio_messages.append(ChatMessage(role="assistant", content="cut_paper", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "find_end_of_tape": | |
| output = self.find_end_of_tape() | |
| gradio_messages.append(ChatMessage(role="assistant", content="find_end_of_tape", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "wrap_present": | |
| output = self.wrap_present() | |
| gradio_messages.append(ChatMessage(role="assistant", content="wrap_present", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "label_present": | |
| recipient = arguments["recipient"] | |
| output = self.label_present(recipient) | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"label_present({recipient})", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "retrieve_letters": | |
| output = self.retrieve_letters() | |
| gradio_messages.append(ChatMessage(role="assistant", content="retrieve_letters", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "check_temperature": | |
| object = arguments["object"] | |
| output = self.check_temperature(object) | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"check_temperature({object})", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "dunk_cookie": | |
| output = self.dunk_cookie() | |
| gradio_messages.append(ChatMessage(role="assistant", content="dunk_cookie", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "drink": | |
| item = arguments["item"] | |
| output = self.drink(item) | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"drink({item})", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "put_route_into_maps": | |
| addr1, addr2, addr3 = arguments["addr1"], arguments["addr2"], arguments["addr3"] | |
| output = self.put_route_into_maps(addr1, addr2, addr3) | |
| gradio_messages.append(ChatMessage(role="assistant", content=f"put_route_into_maps({addr1}, {addr2}, {addr3})", metadata={"title": f"🔧 Tool Call: {tool_call.function.name}"})) | |
| elif tool_call.function.name == "stop": | |
| output = self.stop() | |
| should_stop = True | |
| messages.append({"role": "tool", "content": output, "tool_call_id": tool_call.id}) | |
| if not should_stop: | |
| gradio_messages.append(ChatMessage(role="assistant", content=output, metadata={"title": f"🔧 Tool Output: {tool_call.function.name}"})) | |
| else: | |
| non_tool_count += 1 | |
| if non_tool_count >= 2: | |
| break | |
| if should_stop or len(messages) > 10: | |
| break | |
| return messages, gradio_messages | |