-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
59 lines (52 loc) · 2.16 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import sys
import os
import shlex
sys.stdout.reconfigure(encoding='utf-8')
from src.colors import colors
from src.chats import Chat
import src.config as config
from src.args import parse_arguments
import warnings
from urllib3.exceptions import InsecureRequestWarning
warnings.simplefilter('ignore', InsecureRequestWarning)
outro = f'\n {colors["cyan"]}Goodbye...{colors["reset"]}\n'
def chat_loop(chat):
input_prefix = f'{colors["grey"]}\u252C {colors["cyan"]}💬{chat.name}\n{colors["grey"]}\u2514\u2500 {colors["bright_cyan"]}{os.getenv("USERNAME")}: {colors["reset"]}'
while True:
try:
user_input = input(input_prefix)
try:
args_list = shlex.split(user_input)
except:
args_list = None
if user_input.startswith('-') and args_list is not None:
args = parse_arguments(args_list)
handle_args(args)
else:
output_prefix = f'{colors["grey"]}\u252C {colors["cyan"]}\u2699 {config.get_config("model")}\n{colors["grey"]}\u2514\u2500 {colors["bright_cyan"]}ChatGPT: {colors["bright_blue"]}'
print(f"{output_prefix}{chat.add_message('user', user_input)}")
except KeyboardInterrupt:
print(outro)
break
def handle_args(args):
if args.model:
config.set_config('model', args.model)
elif args.select_chat:
chat = Chat.load_chat(args.select_chat)
chat_loop(chat)
elif args.delete_chat:
Chat.delete_chat(args.delete_chat)
elif args.list_chats:
Chat.list_chats()
elif args.question:
chat = Chat(name='temp')
output_prefix = f'{colors["grey"]}\u252C {colors["cyan"]}\u2699 {config.get_config("model")}\n{colors["grey"]}\u2514\u2500 {colors["bright_cyan"]}ChatGPT: {colors["bright_blue"]}'
print(f"{output_prefix}{chat.add_message('user', args.question)}")
Chat.delete_chat('temp')
args = parse_arguments()
handle_args(args)
if not any(vars(args).values()):
Chat.list_chats()
name = input('Enter the name of the chat you want to load, or enter a name for a new chat: ')
chat = Chat.load_chat(name)
chat_loop(chat)