Enable the Chat Protocol

1
import os
2
from datetime import datetime
3
from uuid import uuid4
4
5
from openai import OpenAI
6
from uagents import Context, Protocol, Agent
7
from uagents.experimental.chat_agent.protocol import build_llm_message_history
8
from uagents_core.contrib.protocols.chat import (
9
    ChatAcknowledgement,
10
    ChatMessage,
11
    EndSessionContent,
12
    StartSessionContent,
13
    TextContent,
14
    chat_protocol_spec,
15
)
16
17
##
18
### Example Expert Assistant
19
##
20
## This chat example is a barebones demonstration of how to attach a chat protocol to an agent
21
## and customize its behavior. In this example, we prompt the ASI-1 model to answer questions
22
## on a specific subject only.
23
##
24
25
def create_text_chat(text: str, end_session: bool = False) -> ChatMessage:
26
    content = [TextContent(type="text", text=text)]
27
    if end_session:
28
        content.append(EndSessionContent(type="end-session"))
29
    return ChatMessage(timestamp=datetime.utcnow(), msg_id=uuid4(), content=content)
30
31
# the subject that this assistant is an expert in
32
subject_matter = "The Sun"
33
34
SYSTEM_PROMPT = (
35
    f"You are a helpful assistant who only answers questions about {subject_matter}. "
36
    "If the user asks about any other topics, you should politely say that you do not know about them."
37
)
38
39
# Agentverse injects both variables into Hosted Agents, so no API key is needed here.
40
# Running outside Agentverse? Set them yourself, using a key from https://asi1.ai/developer
41
client = OpenAI(
42
    base_url=os.environ["ASI1_BASE_URL"],
43
    api_key=os.environ["ASI1_API_KEY"],
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()