Enable the Chat Protocol

1
from datetime import datetime
2
from uuid import uuid4
3
4
from openai import OpenAI
5
from uagents import Context, Protocol, Agent
6
from uagents.experimental.chat_agent.protocol import build_llm_message_history
7
from uagents_core.contrib.protocols.chat import (
8
    ChatAcknowledgement,
9
    ChatMessage,
10
    EndSessionContent,
11
    StartSessionContent,
12
    TextContent,
13
    chat_protocol_spec,
14
)
15
16
##
17
### Example Expert Assistant
18
##
19
## This chat example is a barebones demonstration of how to attach a chat protocol to an agent
20
## and customize its behavior. In this example, we prompt the ASI-1 model to answer questions
21
## on a specific subject only.
22
##
23
24
def create_text_chat(text: str, end_session: bool = False) -> ChatMessage:
25
    content = [TextContent(type="text", text=text)]
26
    if end_session:
27
        content.append(EndSessionContent(type="end-session"))
28
    return ChatMessage(timestamp=datetime.utcnow(), msg_id=uuid4(), content=content)
29
30
# the subject that this assistant is an expert in
31
subject_matter = "The Sun"
32
33
SYSTEM_PROMPT = (
34
    f"You are a helpful assistant who only answers questions about {subject_matter}. "
35
    "If the user asks about any other topics, you should politely say that you do not know about them."
36
)
37
38
client = OpenAI(
39
    # By default, we are using the ASI-1 LLM endpoint and model
40
    base_url='https://api.asi1.ai/v1',
41
42
    # You can get an ASI-1 api key by creating an account at https://asi1.ai/developer
43
    api_key='INSERT_YOUR_API_HERE',
44
)
45
46
agent = Agent()
47
48
# We create a new protocol which is compatible with the chat protocol spec. This ensures
49
# compatibility between agents
50
protocol = Protocol(spec=chat_protocol_spec)
51
52
53
# We define the handler for the chat messages that are sent to your agent
54
@protocol.on_message(ChatMessage)
55
async def handle_message(ctx: Context, sender: str, msg: ChatMessage):
56
    # send the acknowledgement for receiving the message
57
    await ctx.send(
58
        sender,
59
        ChatAcknowledgement(timestamp=datetime.now(), acknowledged_msg_id=msg.msg_id),
60
    )
61
62
    text = msg.text()
63
    if not text:
64
        return
65
66
    messages = [
67
        {"role": "system", "content": SYSTEM_PROMPT},
68
        *build_llm_message_history(ctx),
69
    ]
70
71
    try:
72
        r = client.chat.completions.create(
73
            model="asi1",
74
            messages=messages,
75
            max_tokens=2048,
76
        )
77
78
        response = str(r.choices[0].message.content)
79
    except Exception as e:
80
        ctx.logger.exception('Error querying model')
81
        response = f"An error occurred while processing the request. Please try again later. {e}"
82
83
    await ctx.send(sender, create_text_chat(response))
84
85
86
@protocol.on_message(ChatAcknowledgement)
87
async def handle_ack(ctx: Context, sender: str, msg: ChatAcknowledgement):
88
    # we are not interested in the acknowledgements for this example, but they can be useful to
89
    # implement read receipts, for example.
90
    pass
91
92
93
# attach the protocol to the agent
94
agent.include(protocol, publish_manifest=True)
95
96
if __name__ == "__main__":
97
    agent.run()