Splain main.py to several files.

chatmastermind/storage.py

@@ -0,0 +1,57 @@
+import yaml
+import io
+import pathlib
+from .utils import terminal_width, append_message, message_to_chat
+from typing import List, Dict, Any, Optional
+
+
+def save_answers(question: str,
+                 answers: list[str],
+                 tags: list[str],
+                 otags: Optional[list[str]]
+                 ) -> None:
+    wtags = otags or tags
+    for num, answer in enumerate(answers, start=1):
+        title = f'-- ANSWER {num} '
+        title_end = '-' * (terminal_width() - len(title))
+        print(f'{title}{title_end}')
+        print(answer)
+        with open(f"{num:02d}.yaml", "w") as fd:
+            with io.StringIO() as f:
+                yaml.dump({'question': question},
+                          f,
+                          default_style="|",
+                          default_flow_style=False)
+                fd.write(f.getvalue().replace('"question":', "question:", 1))
+            with io.StringIO() as f:
+                yaml.dump({'answer': answer},
+                          f,
+                          default_style="|",
+                          default_flow_style=False)
+                fd.write(f.getvalue().replace('"answer":', "answer:", 1))
+            yaml.dump({'tags': wtags},
+                      fd,
+                      default_flow_style=False)
+
+
+def create_chat(question: Optional[str],
+                tags: Optional[List[str]],
+                extags: Optional[List[str]],
+                config: Dict[str, Any]
+                ) -> List[Dict[str, str]]:
+    chat = []
+    append_message(chat, 'system', config['system'].strip())
+    for file in sorted(pathlib.Path(config['db']).iterdir()):
+        if file.suffix == '.yaml':
+            with open(file, 'r') as f:
+                data = yaml.load(f, Loader=yaml.FullLoader)
+            data_tags = set(data.get('tags', []))
+            tags_match = \
+                not tags or data_tags.intersection(tags)
+            extags_do_not_match = \
+                not extags or not data_tags.intersection(extags)
+            if tags_match and extags_do_not_match:
+                message_to_chat(data, chat)
+    if question:
+        append_message(chat, 'user', question)
+    return chat