diff --git a/memor/__init__.py b/memor/__init__.py index 4bc9405..44a32e3 100644 --- a/memor/__init__.py +++ b/memor/__init__.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- """Memor modules.""" from .params import MEMOR_VERSION - +from .prompt import Prompt __version__ = MEMOR_VERSION diff --git a/memor/prompt.py b/memor/prompt.py new file mode 100644 index 0000000..95e5998 --- /dev/null +++ b/memor/prompt.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +"""Prompt class.""" +import enum +import datetime +import json + + +class Role(enum.Enum): + """Role enum.""" + + SYSTEM = 0 + USER = 1 + ASSISTANT = 2 + + DEFAULT = USER + + +class Prompt: + """Prompt class.""" + + def __init__( + self, + message, + responses=[], + role=Role.DEFAULT, + temperature=None, + model=None, + date=datetime.datetime.now()): + """Prompt object initiator.""" + self.message = message + self.responses = responses + self.role = role + self.temperature = temperature + self.model = model + self.date = date + + def add_response(self, response, index=None): + """Add a response to the prompt object.""" + if index is None: + self.responses.append(response) + else: + self.responses.insert(index, response) + + def remove_response(self, index): + """Remove a response from the prompt object.""" + self.responses.pop(index) + + def get_message(self): + """Get the prompt message.""" + return self.message + + def to_json(self): + """Convert the prompt to a JSON object.""" + data = { + "message": self.message, + "responses": self.responses, + "role": str(self.role), + "temperature": self.temperature, + "model": self.model, + "date": str(self.date) + } + return json.dumps(data, indent=4)