46 lines
973 B
Python
46 lines
973 B
Python
from core.types import BaseRole
|
|
|
|
from langchain.schema import HumanMessage, SystemMessage, AIMessage, ChatMessage
|
|
|
|
|
|
class User(BaseRole):
|
|
def __init__(self):
|
|
super().__init__("user")
|
|
|
|
def generate(self, message):
|
|
return HumanMessage(content=message)
|
|
|
|
|
|
class AI(BaseRole):
|
|
def __init__(self):
|
|
super().__init__("ai")
|
|
|
|
def generate(self, message):
|
|
return AIMessage(content=message)
|
|
|
|
|
|
class System(BaseRole):
|
|
def __init__(self):
|
|
super().__init__("system")
|
|
|
|
def generate(self, message):
|
|
return SystemMessage(content=message)
|
|
|
|
|
|
class Uploader(BaseRole):
|
|
def __init__(self):
|
|
super().__init__("uploader")
|
|
self._file = None
|
|
|
|
def generate(self, message):
|
|
return HumanMessage(content=message)
|
|
|
|
|
|
class Sender(BaseRole):
|
|
def __init__(self):
|
|
super().__init__("sender")
|
|
self._file = None
|
|
|
|
def generate(self, message):
|
|
return AIMessage(content=message)
|