Prompt Class Added

memor/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 """Memor modules."""
 from .params import MEMOR_VERSION
-
+from .prompt import Prompt
 
 __version__ = MEMOR_VERSION

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)