From 935f1ee18a28d6a724748fe643a0c10956e615ef Mon Sep 17 00:00:00 2001 From: juk0de Date: Sat, 2 Sep 2023 10:00:08 +0200 Subject: [PATCH] message: fixed Answer header for TXT format --- chatmastermind/message.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/chatmastermind/message.py b/chatmastermind/message.py index 384fb96..87de8e2 100644 --- a/chatmastermind/message.py +++ b/chatmastermind/message.py @@ -96,7 +96,7 @@ class AILine(str): def __new__(cls: Type[AILineInst], string: str) -> AILineInst: if not string.startswith(cls.prefix): - raise TagError(f"AILine '{string}' is missing prefix '{cls.prefix}'") + raise MessageError(f"AILine '{string}' is missing prefix '{cls.prefix}'") instance = super().__new__(cls, string) return instance @@ -116,7 +116,7 @@ class ModelLine(str): def __new__(cls: Type[ModelLineInst], string: str) -> ModelLineInst: if not string.startswith(cls.prefix): - raise TagError(f"ModelLine '{string}' is missing prefix '{cls.prefix}'") + raise MessageError(f"ModelLine '{string}' is missing prefix '{cls.prefix}'") instance = super().__new__(cls, string) return instance @@ -133,7 +133,7 @@ class Answer(str): A single answer with a defined header. """ tokens: int = 0 # tokens used by this answer - txt_header: ClassVar[str] = '=== ANSWER ===' + txt_header: ClassVar[str] = '==== ANSWER ====' yaml_key: ClassVar[str] = 'answer' def __new__(cls: Type[AnswerInst], string: str) -> AnswerInst: @@ -355,17 +355,20 @@ class Message(): try: pos = fd.tell() ai = AILine(fd.readline()).ai() - except TagError: + except MessageError: fd.seek(pos) # ModelLine (Optional) try: pos = fd.tell() model = ModelLine(fd.readline()).model() - except TagError: + except MessageError: fd.seek(pos) # Question and Answer text = fd.read().strip().split('\n') - question_idx = text.index(Question.txt_header) + 1 + try: + question_idx = text.index(Question.txt_header) + 1 + except ValueError: + raise MessageError(f"Question header '{Question.txt_header}' not found in '{file_path}'") try: answer_idx = text.index(Answer.txt_header) question = Question.from_list(text[question_idx:answer_idx])