Compare commits

..

7 Commits

4 changed files with 87 additions and 193 deletions

View File

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

View File

@ -3,15 +3,13 @@ Module implementing message related functions and classes.
"""
import pathlib
import yaml
from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final
from typing import Type, TypeVar, ClassVar, Optional, Any, Union
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]]]
@ -57,46 +55,6 @@ 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.
@ -172,15 +130,10 @@ 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:
@ -190,8 +143,6 @@ 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
@ -221,8 +172,6 @@ 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]
@ -231,8 +180,6 @@ 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():
@ -243,34 +190,16 @@ 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(pos)
fd.seek(0, 0) # allow message files without tags
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:
@ -279,9 +208,8 @@ class Message():
answer = Answer.from_list(text[answer_idx + 1:])
except ValueError:
question = Question.from_list(text[question_idx:])
return cls(question, answer, tags, ai, model, file_path)
# YAML
else:
return cls(question, answer, tags, file_path)
else: # '.yaml'
with open(file_path, "r") as fd:
data = yaml.load(fd, Loader=yaml.FullLoader)
if tags_or or tags_and or tags_not:
@ -293,13 +221,11 @@ class Message():
data[cls.file_yaml_key] = file_path
return cls.from_dict(data)
def to_file(self, file_path: Optional[pathlib.Path]) -> None: # noqa: 11
def to_file(self, file_path: Optional[pathlib.Path]) -> None:
"""
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
@ -308,8 +234,6 @@ 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
@ -317,30 +241,20 @@ 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, sort_keys=False)
yaml.dump(data, fd)
def as_dict(self) -> dict[str, Any]:
return asdict(self)

View File

@ -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 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.
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.
"""
# 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([cls.prefix] + sorted([t for t in tags])))
return cls(' '.join([TagLine.prefix] + sorted([t for t in tags])))
def tags(self) -> set[Tag]:
"""

View File

@ -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, AILine, ModelLine
from chatmastermind.message import source_code, Message, MessageError, Question, Answer
from chatmastermind.tags import Tag, TagLine
@ -89,9 +89,7 @@ class MessageToFileTxtTestCase(CmmTestCase):
self.message = Message(Question('This is a question.'),
Answer('This is an answer.'),
{Tag('tag1'), Tag('tag2')},
ai='ChatGPT',
model='gpt-3.5-turbo',
file_path=self.file_path)
self.file_path)
def tearDown(self) -> None:
self.file.close()
@ -103,8 +101,6 @@ 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}
@ -138,15 +134,11 @@ class MessageToFileYamlTestCase(CmmTestCase):
self.message = Message(Question('This is a question.'),
Answer('This is an answer.'),
{Tag('tag1'), Tag('tag2')},
ai='ChatGPT',
model='gpt-3.5-turbo',
file_path=self.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')},
ai='ChatGPT',
model='gpt-3.5-turbo',
file_path=self.file_path)
self.file_path)
def tearDown(self) -> None:
self.file.close()
@ -157,14 +149,7 @@ class MessageToFileYamlTestCase(CmmTestCase):
with open(self.file_path, "r") as fd:
content = fd.read()
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
"""
expected_content = "answer: This is an answer.\nquestion: This is a question.\ntags:\n- tag1\n- tag2\n"
self.assertEqual(content, expected_content)
def test_to_file_yaml_multiline(self) -> None:
@ -172,14 +157,12 @@ class MessageToFileYamlTestCase(CmmTestCase):
with open(self.file_path, "r") as fd:
content = fd.read()
expected_content = f"""{Question.yaml_key}: |-
This is a
multiline question.
{Answer.yaml_key}: |-
expected_content = f"""{Answer.yaml_key}: |-
This is a
multiline answer.
{Message.ai_yaml_key}: ChatGPT
{Message.model_yaml_key}: gpt-3.5-turbo
{Question.yaml_key}: |-
This is a
multiline question.
{Message.tags_yaml_key}:
- tag1
- tag2
@ -198,23 +181,29 @@ This is a question.
{Answer.txt_header}
This is an answer.
""")
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:
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:
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_min.close()
self.file_no_tags.close()
self.file_path.unlink()
self.file_path_min.unlink()
self.file_path_no_tags.unlink()
def test_from_file_txt_complete(self) -> None:
"""
Read a complete message (with all optional values).
"""
def test_from_file_txt(self) -> None:
message = Message.from_file(self.file_path)
self.assertIsNotNone(message)
self.assertIsInstance(message, Message)
@ -224,16 +213,24 @@ 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_min(self) -> None:
"""
Read a message with only required values.
"""
message = Message.from_file(self.file_path_min)
def test_from_file_txt_no_tags(self) -> None:
message = Message.from_file(self.file_path_no_tags)
self.assertIsNotNone(message)
self.assertIsInstance(message, Message)
if message: # mypy bug
self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.file_path, self.file_path_min)
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.assertIsNone(message.answer)
def test_from_file_txt_tags_match(self) -> None:
@ -251,17 +248,18 @@ 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_min, tags_or={Tag('tag1')})
message = Message.from_file(self.file_path_no_tags, 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_min, tags_not={Tag('tag1')})
message = Message.from_file(self.file_path_no_tags, 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_min)
self.assertEqual(message.file_path, self.file_path_no_tags)
def test_from_file_not_exists(self) -> None:
file_not_exists = pathlib.Path("example.txt")
@ -284,24 +282,31 @@ class MessageFromFileYamlTestCase(CmmTestCase):
- tag1
- tag2
""")
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:
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:
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_complete(self) -> None:
"""
Read a complete message (with all optional values).
"""
def test_from_file_yaml(self) -> None:
message = Message.from_file(self.file_path)
self.assertIsInstance(message, Message)
self.assertIsNotNone(message)
@ -311,17 +316,24 @@ 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_min(self) -> None:
"""
Read a message with only the required values.
"""
message = Message.from_file(self.file_path_min)
def test_from_file_yaml_no_tags(self) -> None:
message = Message.from_file(self.file_path_no_tags)
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_min)
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.assertIsNone(message.answer)
def test_from_file_not_exists(self) -> None:
@ -345,47 +357,15 @@ class MessageFromFileYamlTestCase(CmmTestCase):
self.assertIsNone(message)
def test_from_file_yaml_no_tags_dont_match(self) -> None:
message = Message.from_file(self.file_path_min, tags_or={Tag('tag1')})
message = Message.from_file(self.file_path_no_tags, 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_min, tags_not={Tag('tag1')})
message = Message.from_file(self.file_path_no_tags, 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_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')})
self.assertEqual(message.file_path, self.file_path_no_tags)