fix: roles parser (#5)

README.md

@@ -41,6 +41,10 @@ Let AskGPT check grammar errors (from clipboard) for you:
 
 ## Changelog
 
+[v0.6.1](https://github.com/phguo/AskGPT/releases/tag/v0.6.1) - Apr. 4, 2023
+  - Fix #5 that are related to the roles parser.
+  - Add a configuration for printing user-inputted content.
+
 [v0.6](https://github.com/phguo/AskGPT/releases/tag/v0.6) - Apr. 2, 2023
 
 <img src="https://raw.githubusercontent.com/phguo/AskGPT/main/video/v0.6_User_Configuration.png"  width="70%" height="70%">

gpt.py

@@ -9,8 +9,10 @@
 chat_cache = "chat.json"
 
 
-def pre_process(api_key, content, model, temprature, context_time, chat_number, clear_keyword, clipboard_keyword, roles, signature):
+def pre_process(api_key, content, model, temprature, context_time, chat_number, clear_keyword, clipboard_keyword, roles, signature, print_input):
     # context_time
+    if float(context_time) < 0:
+        context_time = float("inf")
     if os.path.isfile(chat_cache):
         mod_time = os.path.getmtime(chat_cache)
         current_time = time.time()
@@ -33,26 +35,34 @@ def pre_process(api_key, content, model, temprature, context_time, chat_number,
 
     # roles
     for role in roles.split('\n'):
-        k, v = role.split(':')
+        index = role.find(":")
+        k = role[:index]
+        v = role[index + 1:].strip()
         if k in content:
             content = content.replace(k, '')
             content = v + content
 
     # signature
     signature = int(signature)
 
+    # print_input
+    print_input = int(print_input)
+
     messages = []
     if os.path.exists(chat_cache):
         with open(chat_cache, "r") as chat_json:
             messages = json.load(chat_json)
     messages += [{'role': 'user', 'content': content}]
 
-    return api_key, model, temprature, messages, signature, chat_number
+    return api_key, model, temprature, messages, signature, print_input, chat_number
 
 
-def ask_gpt(api_key, model, temprature, messages, signature):
+def ask_gpt(api_key, model, temprature, messages, signature, print_input):
     openai.api_key = api_key
 
+    if print_input:
+        keyboard.write('>>> ' + messages[-1]["content"] + '\n')
+
     response = openai.ChatCompletion.create(
         model=model,
         messages=messages,
@@ -71,7 +81,7 @@ def ask_gpt(api_key, model, temprature, messages, signature):
             if signature == 1:
                 from datetime import datetime
                 current = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                keyboard.write("\n[Generated by ChatGPT, {}]".format(current))
+                keyboard.write("\n[Generated by ChatGPT, {}]\n\n".format(current))
 
     messages += [{'role': 'assistant', 'content': assistant_content}]
     return messages
@@ -86,7 +96,7 @@ def post_process(messages, chat_number):
 
 if __name__ == '__main__':
     params = sys.argv[1:]
-    api_key, model, temprature, messages, signature, chat_number = pre_process(*params)
+    api_key, model, temprature, messages, signature, print_input, chat_number = pre_process(*params)
     if messages[-1].get("content", None):
-        messages = ask_gpt(api_key, model, temprature, messages, signature)
+        messages = ask_gpt(api_key, model, temprature, messages, signature, print_input)
         post_process(messages, chat_number)