File size: 5,041 Bytes
afbbf55 d6bfc1a afbbf55 679197d afbbf55 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
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": "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 stop(self):
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),
]
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:
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 == "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}"}))
if should_stop or len(messages) > 10:
break
return messages, gradio_messages
|