You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
83 lines
3.0 KiB
83 lines
3.0 KiB
from langchain.schema import AIMessage, HumanMessage, SystemMessage, ChatMessage
|
|
|
|
|
|
class Message(object):
|
|
def __init__(self, timestamp, user_name, message, event_id=None, user_id=None, room_name=None, room_id=None, role=None):
|
|
self.timestamp = timestamp
|
|
self.user_name = user_name
|
|
self.message = message
|
|
self.event_id = event_id
|
|
self.user_id = user_id
|
|
self.room_name = room_name
|
|
self.room_id = room_id
|
|
self.role = role
|
|
|
|
@classmethod
|
|
def from_matrix(cls, room, event):
|
|
return cls(event.server_timestamp/1000, room.user_name(event.sender), event.body, event.event_id, event.sender, room.display_name, room.room_id)
|
|
|
|
def to_langchain(self):
|
|
if self.role == "human":
|
|
return HumanMessage(
|
|
content=self.message,
|
|
role=self.user_name, # "chat"
|
|
additional_kwargs={
|
|
"timestamp": self.timestamp,
|
|
"user_name": self.user_name,
|
|
"event_id": self.event_id,
|
|
"user_id": self.user_id,
|
|
"room_name": self.room_name,
|
|
"room_id": self.room_id,
|
|
}
|
|
)
|
|
elif self.role == "ai":
|
|
return AIMessage(
|
|
content=self.message,
|
|
role=self.user_name, # "chat"
|
|
additional_kwargs={
|
|
"timestamp": self.timestamp,
|
|
"user_name": self.user_name,
|
|
"event_id": self.event_id,
|
|
"user_id": self.user_id,
|
|
"room_name": self.room_name,
|
|
"room_id": self.room_id,
|
|
}
|
|
)
|
|
elif self.role == "system":
|
|
return SystemMessage(
|
|
content=self.message,
|
|
role=self.user_name, # "chat"
|
|
additional_kwargs={
|
|
"timestamp": self.timestamp,
|
|
"user_name": self.user_name,
|
|
"event_id": self.event_id,
|
|
"user_id": self.user_id,
|
|
"room_name": self.room_name,
|
|
"room_id": self.room_id,
|
|
}
|
|
)
|
|
else:
|
|
return ChatMessage(
|
|
content=self.message,
|
|
role=self.user_name, # "chat"
|
|
additional_kwargs={
|
|
"timestamp": self.timestamp,
|
|
"user_name": self.user_name,
|
|
"event_id": self.event_id,
|
|
"user_id": self.user_id,
|
|
"room_name": self.room_name,
|
|
"room_id": self.room_id,
|
|
}
|
|
)
|
|
|
|
def is_from(self, user_id):
|
|
return self.user_id == user_id
|
|
|
|
def is_command(self):
|
|
return self.message.startswith('!')
|
|
|
|
def is_error(self):
|
|
return self.message.startswith('<ERROR>')
|
|
|
|
def __str__(self):
|
|
return str("{}: {}".format(self.user_name, self.message))
|
|
|