From e4d055b90033773bba712f8ab49cf9339d7d83ed Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Sat, 12 Aug 2023 12:20:49 +0200
Subject: [PATCH] Fix the max_tokens, temperature, and model setup.

---
 chatmastermind/main.py  | 26 ++++++++++++++++++--------
 chatmastermind/utils.py | 11 +++++------
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/chatmastermind/main.py b/chatmastermind/main.py
index 843885c..3ed387c 100755
--- a/chatmastermind/main.py
+++ b/chatmastermind/main.py
@@ -21,7 +21,7 @@ def tags_completer(prefix, parsed_args, **kwargs):
     return get_tags_unique(config, prefix)
 
 
-def read_config(path: str):
+def read_config(path: str) -> ConfigType:
     with open(path, 'r') as f:
         config = yaml.load(f, Loader=yaml.FullLoader)
     return config
@@ -81,6 +81,15 @@ def ask_cmd(args: argparse.Namespace, config: ConfigType) -> None:
     """
     Handler for the 'ask' command.
     """
+    if type(config['openai']) is not dict:
+        raise RuntimeError('Configuration openai is not a dict.')
+    config_openai = config['openai']
+    if args.max_tokens:
+        config_openai['max_tokens'] = args.max_tokens
+    if args.temperature:
+        config_openai['temperature'] = args.temperature
+    if args.model:
+        config_openai['model'] = args.model
     chat, question, tags = create_question_with_hist(args, config)
     print_chat_hist(chat, False, args.only_source_code)
     otags = args.output_tags or []
@@ -211,13 +220,14 @@ def main() -> int:
     config = read_config(args.config)
 
     # modify config according to args
-    openai_api_key(config['openai']['api_key'])
-    if args.max_tokens:
-        config['openai']['max_tokens'] = args.max_tokens
-    if args.temperature:
-        config['openai']['temperature'] = args.temperature
-    if args.model:
-        config['openai']['model'] = args.model
+    if type(config['openai']) is dict:
+        config_openai = config['openai']
+    else:
+        RuntimeError("Configuration openai is not a dict.")
+    if type(config_openai['api_key']) is str:
+        openai_api_key(config_openai['api_key'])
+    else:
+        raise RuntimeError("Configuration openai.api_key is not a string.")
 
     command.func(command, config)
 
diff --git a/chatmastermind/utils.py b/chatmastermind/utils.py
index 2a58dae..fba8296 100644
--- a/chatmastermind/utils.py
+++ b/chatmastermind/utils.py
@@ -1,8 +1,7 @@
 import shutil
 from pprint import PrettyPrinter
-from typing import List, Dict, Union
 
-ConfigType = Dict[str, Union[str, Dict[str, Union[str, int]]]]
+ConfigType = dict[str, str | dict[str, str | int | float]]
 
 
 def terminal_width() -> int:
@@ -31,15 +30,15 @@ def print_tag_args(tags: list[str], extags: list[str], otags: list[str]) -> None
         print()
 
 
-def append_message(chat: List[Dict[str, str]],
+def append_message(chat: list[dict[str, str]],
                    role: str,
                    content: str
                    ) -> None:
     chat.append({'role': role, 'content': content.replace("''", "'")})
 
 
-def message_to_chat(message: Dict[str, str],
-                    chat: List[Dict[str, str]],
+def message_to_chat(message: dict[str, str],
+                    chat: list[dict[str, str]],
                     with_tags: bool = False,
                     with_file: bool = False
                     ) -> None:
@@ -80,6 +79,6 @@ def print_chat_hist(chat, dump=False, source_code=False) -> None:
             print(f"{message['role'].upper()}: {message['content']}")
 
 
-def print_tags_frequency(tags: List[str]) -> None:
+def print_tags_frequency(tags: list[str]) -> None:
     for tag in sorted(set(tags)):
         print(f"- {tag}: {tags.count(tag)}")