diff --git a/chatmastermind/storage.py b/chatmastermind/storage.py index be2bd88..0ce6f9e 100644 --- a/chatmastermind/storage.py +++ b/chatmastermind/storage.py @@ -5,9 +5,11 @@ from .utils import terminal_width, append_message, message_to_chat from typing import List, Dict, Any, Optional -def read_file(fname: str) -> Dict[str, Any]: +def read_file(fname: str, tags_only: bool = False) -> Dict[str, Any]: with open(fname, "r") as fd: - text = fd.read().split('\n') + if tags_only: + return {"tags": [x.strip() for x in fd.readline().strip().split(':')[1].strip().split(',')]} + text = fd.read().strip().split('\n') tags = [x.strip() for x in text.pop(0).split(':')[1].strip().split(',')] question_idx = text.index("=== QUESTION ===") + 1 answer_idx = text.index("==== ANSWER ====") @@ -91,7 +93,7 @@ def get_tags(config: Dict[str, Any], prefix: Optional[str]) -> List[str]: with open(file, 'r') as f: data = yaml.load(f, Loader=yaml.FullLoader) elif file.suffix == '.txt': - data = read_file(file) + data = read_file(file, tags_only=True) else: continue for tag in data.get('tags', []):