tests: added testcases for Message.from_file() and Message.to_file()

chatmastermind/configuration.py

@@ -63,7 +63,7 @@ class Config():
 
     def to_file(self, path: str) -> None:
         with open(path, 'w') as f:
-            yaml.dump(asdict(self), f)
+            yaml.dump(asdict(self), f, sort_keys=False)
 
     def as_dict(self) -> dict[str, Any]:
         return asdict(self)

chatmastermind/message.py

@@ -3,13 +3,15 @@ Module implementing message related functions and classes.
 """
 import pathlib
 import yaml
-from typing import Type, TypeVar, ClassVar, Optional, Any, Union
+from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final
 from dataclasses import dataclass, asdict
 from .tags import Tag, TagLine, TagError, match_tags
 
 QuestionInst = TypeVar('QuestionInst', bound='Question')
 AnswerInst = TypeVar('AnswerInst', bound='Answer')
 MessageInst = TypeVar('MessageInst', bound='Message')
+AILineInst = TypeVar('AILineInst', bound='AILine')
+ModelLineInst = TypeVar('ModelLineInst', bound='ModelLine')
 YamlDict = dict[str, Union[QuestionInst, AnswerInst, set[Tag]]]
 
 
@@ -55,6 +57,46 @@ def source_code(text: str, include_delims: bool = False) -> list[str]:
     return code_sections
 
 
+class AILine(str):
+    """
+    A line that represents the AI name in a '.txt' file..
+    """
+    prefix: Final[str] = 'AI:'
+
+    def __new__(cls: Type[AILineInst], string: str) -> AILineInst:
+        if not string.startswith(cls.prefix):
+            raise TagError(f"AILine '{string}' is missing prefix '{cls.prefix}'")
+        instance = super().__new__(cls, string)
+        return instance
+
+    def ai(self) -> str:
+        return self[len(self.prefix):].strip()
+
+    @classmethod
+    def from_ai(cls: Type[AILineInst], ai: str) -> AILineInst:
+        return cls(' '.join([cls.prefix, ai]))
+
+
+class ModelLine(str):
+    """
+    A line that represents the model name in a '.txt' file..
+    """
+    prefix: Final[str] = 'MODEL:'
+
+    def __new__(cls: Type[ModelLineInst], string: str) -> ModelLineInst:
+        if not string.startswith(cls.prefix):
+            raise TagError(f"ModelLine '{string}' is missing prefix '{cls.prefix}'")
+        instance = super().__new__(cls, string)
+        return instance
+
+    def model(self) -> str:
+        return self[len(self.prefix):].strip()
+
+    @classmethod
+    def from_model(cls: Type[ModelLineInst], model: str) -> ModelLineInst:
+        return cls(' '.join([cls.prefix, model]))
+
+
 class Question(str):
     """
     A single question with a defined header.
@@ -130,10 +172,15 @@ class Message():
     question: Question
     answer: Optional[Answer]
     tags: Optional[set[Tag]]
+    ai: Optional[str]
+    model: Optional[str]
     file_path: Optional[pathlib.Path]
+    # class variables
     file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
     tags_yaml_key: ClassVar[str] = 'tags'
     file_yaml_key: ClassVar[str] = 'file_path'
+    ai_yaml_key: ClassVar[str] = 'ai'
+    model_yaml_key: ClassVar[str] = 'model'
 
     @classmethod
     def from_dict(cls: Type[MessageInst], data: dict[str, Any]) -> MessageInst:
@@ -143,6 +190,8 @@ class Message():
         return cls(question=data[Question.yaml_key],
                    answer=data.get(Answer.yaml_key, None),
                    tags=set(data.get(cls.tags_yaml_key, [])),
+                   ai=data.get(cls.ai_yaml_key, None),
+                   model=data.get(cls.model_yaml_key, None),
                    file_path=data.get(cls.file_yaml_key, None))
 
     @classmethod
@@ -172,6 +221,8 @@ class Message():
         Create a Message from the given file. Expects the following file structures:
         For '.txt':
             * TagLine [Optional]
+            * AI [Optional]
+            * Model [Optional]
             * Question.txt_header
             * Question
             * Answer.txt_header [Optional]
@@ -180,6 +231,8 @@ class Message():
             * Question.yaml_key: single or multiline string
             * Answer.yaml_key: single or multiline string [Optional]
             * Message.tags_yaml_key: list of strings [Optional]
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
         Returns 'None' if the message does not fulfill the tag requirements.
         """
         if not file_path.exists():
@@ -190,16 +243,34 @@ class Message():
         tags: set[Tag] = set()
         question: Question
         answer: Optional[Answer] = None
+        ai: Optional[str] = None
+        model: Optional[str] = None
+        # TXT
         if file_path.suffix == '.txt':
             with open(file_path, "r") as fd:
+                # TagLine (Optional)
                 try:
+                    pos = fd.tell()
                     tags = TagLine(fd.readline()).tags()
                 except TagError:
-                    fd.seek(0, 0)  # allow message files without tags
+                    fd.seek(pos)
                 if tags_or or tags_and or tags_not:
                     # match with an empty set if the file has no tags
                     if not match_tags(tags, tags_or, tags_and, tags_not):
                         return None
+                # AILine (Optional)
+                try:
+                    pos = fd.tell()
+                    ai = AILine(fd.readline()).ai()
+                except TagError:
+                    fd.seek(pos)
+                # ModelLine (Optional)
+                try:
+                    pos = fd.tell()
+                    model = ModelLine(fd.readline()).model()
+                except TagError:
+                    fd.seek(pos)
+                # Question and Answer
                 text = fd.read().strip().split('\n')
                 question_idx = text.index(Question.txt_header) + 1
                 try:
@@ -208,8 +279,9 @@ class Message():
                     answer = Answer.from_list(text[answer_idx + 1:])
                 except ValueError:
                     question = Question.from_list(text[question_idx:])
-                return cls(question, answer, tags, file_path)
-        else:  # '.yaml'
+                return cls(question, answer, tags, ai, model, file_path)
+        # YAML
+        else:
             with open(file_path, "r") as fd:
                 data = yaml.load(fd, Loader=yaml.FullLoader)
                 if tags_or or tags_and or tags_not:
@@ -221,11 +293,13 @@ class Message():
                 data[cls.file_yaml_key] = file_path
                 return cls.from_dict(data)
 
-    def to_file(self, file_path: Optional[pathlib.Path]) -> None:
+    def to_file(self, file_path: Optional[pathlib.Path]) -> None:  # noqa: 11
         """
         Write a Message to the given file. Creates the following file structures:
         For '.txt':
             * TagLine
+            * AI [Optional]
+            * Model [Optional]
             * Question.txt_header
             * Question
             * Answer.txt_header
@@ -234,6 +308,8 @@ class Message():
             * Question.yaml_key: single or multiline string
             * Answer.yaml_key: single or multiline string
             * Message.tags_yaml_key: list of strings
+            * Message.ai_yaml_key: str [Optional]
+            * Message.model_yaml_key: str [Optional]
         """
         if file_path:
             self.file_path = file_path
@@ -241,20 +317,30 @@ class Message():
             raise MessageError("Got no valid path to write message")
         if self.file_path.suffix not in self.file_suffixes:
             raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
+        # TXT
         if self.file_path.suffix == '.txt':
             with open(self.file_path, "w") as fd:
                 msg_tags = self.tags or set()
                 fd.write(f'{TagLine.from_set(msg_tags)}\n')
+                if self.ai:
+                    fd.write(f'{AILine.from_ai(self.ai)}\n')
+                if self.model:
+                    fd.write(f'{ModelLine.from_model(self.model)}\n')
                 fd.write(f'{Question.txt_header}\n{self.question}\n')
                 fd.write(f'{Answer.txt_header}\n{self.answer}\n')
+        # YAML
         elif self.file_path.suffix == '.yaml':
             with open(self.file_path, "w") as fd:
                 data: YamlDict = {Question.yaml_key: str(self.question)}
                 if self.answer:
                     data[Answer.yaml_key] = str(self.answer)
+                if self.ai:
+                    data[self.ai_yaml_key] = self.ai
+                if self.model:
+                    data[self.model_yaml_key] = self.model
                 if self.tags:
                     data[self.tags_yaml_key] = sorted([str(tag) for tag in self.tags])
-                yaml.dump(data, fd)
+                yaml.dump(data, fd, sort_keys=False)
 
     def as_dict(self) -> dict[str, Any]:
         return asdict(self)

chatmastermind/tags.py

@@ -93,9 +93,9 @@ def match_tags(tags: set[Tag], tags_or: Optional[set[Tag]], tags_and: Optional[s
 
 class TagLine(str):
     """
-    A line of tags. It starts with a prefix ('TAGS:'), followed by a list of tags,
-    separated by the defaut separator (' '). Any operations on a TagLine will sort
-    the tags.
+    A line of tags in a '.txt' file. It starts with a prefix ('TAGS:'), followed by
+    a list of tags, separated by the defaut separator (' '). Any operations on a
+    TagLine will sort the tags.
     """
     # the prefix
     prefix: Final[str] = 'TAGS:'
@@ -116,7 +116,7 @@ class TagLine(str):
         """
         Create a new TagLine from a set of tags.
         """
-        return cls(' '.join([TagLine.prefix] + sorted([t for t in tags])))
+        return cls(' '.join([cls.prefix] + sorted([t for t in tags])))
 
     def tags(self) -> set[Tag]:
         """

tests/test_message.py

@@ -2,7 +2,7 @@ import pathlib
 import tempfile
 from typing import cast
 from .test_main import CmmTestCase
-from chatmastermind.message import source_code, Message, MessageError, Question, Answer
+from chatmastermind.message import source_code, Message, MessageError, Question, Answer, AILine, ModelLine
 from chatmastermind.tags import Tag, TagLine
 
 
@@ -89,7 +89,9 @@ class MessageToFileTxtTestCase(CmmTestCase):
         self.message = Message(Question('This is a question.'),
                                Answer('This is an answer.'),
                                {Tag('tag1'), Tag('tag2')},
-                               self.file_path)
+                               ai='ChatGPT',
+                               model='gpt-3.5-turbo',
+                               file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
@@ -101,6 +103,8 @@ class MessageToFileTxtTestCase(CmmTestCase):
         with open(self.file_path, "r") as fd:
             content = fd.read()
         expected_content = f"""{TagLine.prefix} tag1 tag2
+{AILine.prefix} ChatGPT
+{ModelLine.prefix} gpt-3.5-turbo
 {Question.txt_header}
 This is a question.
 {Answer.txt_header}
@@ -134,11 +138,15 @@ class MessageToFileYamlTestCase(CmmTestCase):
         self.message = Message(Question('This is a question.'),
                                Answer('This is an answer.'),
                                {Tag('tag1'), Tag('tag2')},
-                               self.file_path)
+                               ai='ChatGPT',
+                               model='gpt-3.5-turbo',
+                               file_path=self.file_path)
         self.message_multiline = Message(Question('This is a\nmultiline question.'),
                                          Answer('This is a\nmultiline answer.'),
                                          {Tag('tag1'), Tag('tag2')},
-                                         self.file_path)
+                                         ai='ChatGPT',
+                                         model='gpt-3.5-turbo',
+                                         file_path=self.file_path)
 
     def tearDown(self) -> None:
         self.file.close()
@@ -149,7 +157,14 @@ class MessageToFileYamlTestCase(CmmTestCase):
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
-        expected_content = "answer: This is an answer.\nquestion: This is a question.\ntags:\n- tag1\n- tag2\n"
+        expected_content = f"""{Question.yaml_key}: This is a question.
+{Answer.yaml_key}: This is an answer.
+{Message.ai_yaml_key}: ChatGPT
+{Message.model_yaml_key}: gpt-3.5-turbo
+{Message.tags_yaml_key}:
+- tag1
+- tag2
+"""
         self.assertEqual(content, expected_content)
 
     def test_to_file_yaml_multiline(self) -> None:
@@ -157,12 +172,14 @@ class MessageToFileYamlTestCase(CmmTestCase):
 
         with open(self.file_path, "r") as fd:
             content = fd.read()
-        expected_content = f"""{Answer.yaml_key}: |-
-  This is a
-  multiline answer.
-{Question.yaml_key}: |-
+        expected_content = f"""{Question.yaml_key}: |-
   This is a
   multiline question.
+{Answer.yaml_key}: |-
+  This is a
+  multiline answer.
+{Message.ai_yaml_key}: ChatGPT
+{Message.model_yaml_key}: gpt-3.5-turbo
 {Message.tags_yaml_key}:
 - tag1
 - tag2
@@ -181,29 +198,23 @@ This is a question.
 {Answer.txt_header}
 This is an answer.
 """)
-        self.file_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
-        self.file_path_no_tags = pathlib.Path(self.file_no_tags.name)
-        with open(self.file_path_no_tags, "w") as fd:
+        self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
+        self.file_path_min = pathlib.Path(self.file_min.name)
+        with open(self.file_path_min, "w") as fd:
             fd.write(f"""{Question.txt_header}
 This is a question.
-{Answer.txt_header}
-This is an answer.
-""")
-        self.file_no_answer = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
-        self.file_path_no_answer = pathlib.Path(self.file_no_answer.name)
-        with open(self.file_path_no_answer, "w") as fd:
-            fd.write(f"""{TagLine.prefix} tag1 tag2
-{Question.txt_header}
-This is a question.
 """)
 
     def tearDown(self) -> None:
         self.file.close()
-        self.file_no_tags.close()
+        self.file_min.close()
         self.file_path.unlink()
-        self.file_path_no_tags.unlink()
+        self.file_path_min.unlink()
 
-    def test_from_file_txt(self) -> None:
+    def test_from_file_txt_complete(self) -> None:
+        """
+        Read a complete message (with all optional values).
+        """
         message = Message.from_file(self.file_path)
         self.assertIsNotNone(message)
         self.assertIsInstance(message, Message)
@@ -213,24 +224,16 @@ This is a question.
             self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
             self.assertEqual(message.file_path, self.file_path)
 
-    def test_from_file_txt_no_tags(self) -> None:
-        message = Message.from_file(self.file_path_no_tags)
+    def test_from_file_txt_min(self) -> None:
+        """
+        Read a message with only required values.
+        """
+        message = Message.from_file(self.file_path_min)
         self.assertIsNotNone(message)
         self.assertIsInstance(message, Message)
         if message:  # mypy bug
             self.assertEqual(message.question, 'This is a question.')
-            self.assertEqual(message.answer, 'This is an answer.')
-            self.assertSetEqual(cast(set[Tag], message.tags), set())
-            self.assertEqual(message.file_path, self.file_path_no_tags)
-
-    def test_from_file_txt_no_answer(self) -> None:
-        message = Message.from_file(self.file_path_no_answer)
-        self.assertIsInstance(message, Message)
-        self.assertIsNotNone(message)
-        if message:  # mypy bug
-            self.assertEqual(message.question, 'This is a question.')
-            self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
-            self.assertEqual(message.file_path, self.file_path_no_answer)
+            self.assertEqual(message.file_path, self.file_path_min)
             self.assertIsNone(message.answer)
 
     def test_from_file_txt_tags_match(self) -> None:
@@ -248,18 +251,17 @@ This is a question.
         self.assertIsNone(message)
 
     def test_from_file_txt_no_tags_dont_match(self) -> None:
-        message = Message.from_file(self.file_path_no_tags, tags_or={Tag('tag1')})
+        message = Message.from_file(self.file_path_min, tags_or={Tag('tag1')})
         self.assertIsNone(message)
 
     def test_from_file_txt_no_tags_match_tags_not(self) -> None:
-        message = Message.from_file(self.file_path_no_tags, tags_not={Tag('tag1')})
+        message = Message.from_file(self.file_path_min, tags_not={Tag('tag1')})
         self.assertIsNotNone(message)
         self.assertIsInstance(message, Message)
         if message:  # mypy bug
             self.assertEqual(message.question, 'This is a question.')
-            self.assertEqual(message.answer, 'This is an answer.')
             self.assertSetEqual(cast(set[Tag], message.tags), set())
-            self.assertEqual(message.file_path, self.file_path_no_tags)
+            self.assertEqual(message.file_path, self.file_path_min)
 
     def test_from_file_not_exists(self) -> None:
         file_not_exists = pathlib.Path("example.txt")
@@ -282,31 +284,24 @@ class MessageFromFileYamlTestCase(CmmTestCase):
   - tag1
   - tag2
 """)
-        self.file_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
-        self.file_path_no_tags = pathlib.Path(self.file_no_tags.name)
-        with open(self.file_path_no_tags, "w") as fd:
+        self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
+        self.file_path_min = pathlib.Path(self.file_min.name)
+        with open(self.file_path_min, "w") as fd:
             fd.write(f"""
 {Question.yaml_key}: |-
   This is a question.
-{Answer.yaml_key}: |-
-  This is an answer.
-""")
-        self.file_no_answer = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
-        self.file_path_no_answer = pathlib.Path(self.file_no_answer.name)
-        with open(self.file_path_no_answer, "w") as fd:
-            fd.write(f"""
-{Question.yaml_key}: |-
-  This is a question.
-{Message.tags_yaml_key}:
-  - tag1
-  - tag2
 """)
 
     def tearDown(self) -> None:
         self.file.close()
         self.file_path.unlink()
+        self.file_min.close()
+        self.file_path_min.unlink()
 
-    def test_from_file_yaml(self) -> None:
+    def test_from_file_yaml_complete(self) -> None:
+        """
+        Read a complete message (with all optional values).
+        """
         message = Message.from_file(self.file_path)
         self.assertIsInstance(message, Message)
         self.assertIsNotNone(message)
@@ -316,24 +311,17 @@ class MessageFromFileYamlTestCase(CmmTestCase):
             self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
             self.assertEqual(message.file_path, self.file_path)
 
-    def test_from_file_yaml_no_tags(self) -> None:
-        message = Message.from_file(self.file_path_no_tags)
+    def test_from_file_yaml_min(self) -> None:
+        """
+        Read a message with only the required values.
+        """
+        message = Message.from_file(self.file_path_min)
         self.assertIsInstance(message, Message)
         self.assertIsNotNone(message)
         if message:  # mypy bug
             self.assertEqual(message.question, 'This is a question.')
-            self.assertEqual(message.answer, 'This is an answer.')
             self.assertSetEqual(cast(set[Tag], message.tags), set())
-            self.assertEqual(message.file_path, self.file_path_no_tags)
-
-    def test_from_file_yaml_no_answer(self) -> None:
-        message = Message.from_file(self.file_path_no_answer)
-        self.assertIsInstance(message, Message)
-        self.assertIsNotNone(message)
-        if message:  # mypy bug
-            self.assertEqual(message.question, 'This is a question.')
-            self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
-            self.assertEqual(message.file_path, self.file_path_no_answer)
+            self.assertEqual(message.file_path, self.file_path_min)
             self.assertIsNone(message.answer)
 
     def test_from_file_not_exists(self) -> None:
@@ -357,15 +345,47 @@ class MessageFromFileYamlTestCase(CmmTestCase):
         self.assertIsNone(message)
 
     def test_from_file_yaml_no_tags_dont_match(self) -> None:
-        message = Message.from_file(self.file_path_no_tags, tags_or={Tag('tag1')})
+        message = Message.from_file(self.file_path_min, tags_or={Tag('tag1')})
         self.assertIsNone(message)
 
     def test_from_file_yaml_no_tags_match_tags_not(self) -> None:
-        message = Message.from_file(self.file_path_no_tags, tags_not={Tag('tag1')})
+        message = Message.from_file(self.file_path_min, tags_not={Tag('tag1')})
         self.assertIsNotNone(message)
         self.assertIsInstance(message, Message)
         if message:  # mypy bug
             self.assertEqual(message.question, 'This is a question.')
-            self.assertEqual(message.answer, 'This is an answer.')
             self.assertSetEqual(cast(set[Tag], message.tags), set())
-            self.assertEqual(message.file_path, self.file_path_no_tags)
+            self.assertEqual(message.file_path, self.file_path_min)
+
+
+class TagsFromFileTestCase(CmmTestCase):
+    def setUp(self) -> None:
+        self.file_txt = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
+        self.file_path_txt = pathlib.Path(self.file_txt.name)
+        with open(self.file_path_txt, "w") as fd:
+            fd.write(f"""{TagLine.prefix} tag1 tag2
+{Question.txt_header}
+This is a question.
+{Answer.txt_header}
+This is an answer.
+""")
+        self.file_yaml = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
+        self.file_path_yaml = pathlib.Path(self.file_yaml.name)
+        with open(self.file_path_yaml, "w") as fd:
+            fd.write(f"""
+{Question.yaml_key}: |-
+  This is a question.
+{Answer.yaml_key}: |-
+  This is an answer.
+{Message.tags_yaml_key}:
+  - tag1
+  - tag2
+""")
+
+    def test_tags_from_file_txt(self) -> None:
+        tags = Message.tags_from_file(self.file_path_txt)
+        self.assertSetEqual(tags, {Tag('tag1'), Tag('tag2')})
+
+    def test_tags_from_file_yaml(self) -> None:
+        tags = Message.tags_from_file(self.file_path_yaml)
+        self.assertSetEqual(tags, {Tag('tag1'), Tag('tag2')})