Fabelous-Ai-Chat/project.py

113 lines
4.5 KiB
Python
Raw Permalink Normal View History

import argparse
from scripts.TerminalChat import TerminalBot
from scripts.GUIChat import ChatGUI
import json
import os
import sys
CONFIG_FILE = "config/config.json"
2024-05-20 09:07:57 +00:00
ERROR_CORRPUTED_CONFIG = "The config file seems to be corrupted, please run: 'python project.py --config'"
def main():
parser = argparse.ArgumentParser(description="Write with Ollama, using standard chat or RAG system, for \n\
for first usage run 'python project.py --config' all options except 'mode' are optional")
parser.add_argument('--config', action='store_true', help='Enable configuration mode')
parser.add_argument('-f', type=str, help='Path to the input file (only in terminal mode)')
parser.add_argument('-p', type=str, help='User prompt (only in terminal mode)')
parser.add_argument('-m', type=str, choices=["gui", "terminal"], help='change mode')
args = parser.parse_args()
config = read_config()
2024-05-20 11:24:27 +00:00
# If no config file exists and user does not go into "config mode"
if config is None and args.config is False:
sys.exit("No Config available. please run: 'python project.py --config' to set it up")
2024-05-20 11:24:27 +00:00
# add config
elif args.config is True:
config = configure()
write_config(config)
2024-05-20 11:24:27 +00:00
# change ui mode
elif args.m:
config = handle_change_mode(args)
write_config(config)
2024-05-20 11:24:27 +00:00
# start terminal ui
elif config["mode"] == "terminal":
handle_terminal(args)
2024-05-20 11:24:27 +00:00
# start customtkinter ui
elif config["mode"] == "gui":
# start gui
try:
gui = ChatGUI(**config["ollamaConfig"])
gui.mainloop()
except TypeError:
2024-05-20 09:07:57 +00:00
sys.exit(ERROR_CORRPUTED_CONFIG)
else:
sys.exit(ERROR_CORRPUTED_CONFIG)
def configure():
print("Configuration mode enabled.")
mode = input("Enter terminal or gui mode (terminal/gui): ")
while mode.lower() not in ["terminal", "gui"]:
print("Invalid input. Please enter 'terminal' or 'gui':")
mode = input("Enter terminal or gui mode (terminal/gui): ")
base_llm_url = input("Enter base LLM URL (standard: http://localhost:11434): ") or "http://localhost:11434"
embeddings_url = input("Enter embeddings URL (standard: http://localhost:11434): ") or "http://localhost:11434"
base_model = input("Enter base model (standard: 'mistral'): ") or "mistral"
embeddings_model = input("Enter embeddings model (standard: 'mxbai-embed-large'): ") or "mxbai-embed-large"
2024-05-19 12:03:20 +00:00
base_header_key = input("Authentication Key for base model (standard: empty): ") or ""
base_header_value = input("Authentication Value for base model (standard: empty): ") or ""
embeddings_header_key = input("Authentication Key for embeddings model (standard: empty): ") or ""
embeddings_header_value = input("Authentication Value for embeddings model (standard: empty): ") or ""
2024-05-19 12:03:20 +00:00
return {"mode": mode, "ollamaConfig":{"base_url": base_llm_url, "embeddings_url": embeddings_url, "base_model": base_model,
"embeddings_model": embeddings_model, "base_header":{base_header_key: base_header_value}
,"embeddings_header" :{embeddings_header_key: embeddings_header_value}}}
def read_config():
if not os.path.exists(CONFIG_FILE):
return None
with open(CONFIG_FILE, "r") as f:
return json.load(f)
2024-05-20 11:24:27 +00:00
def write_config(config: dict):
2024-05-20 11:24:27 +00:00
try:
with open(CONFIG_FILE, "w") as config_file:
json.dump(config, config_file, indent=4)
except FileNotFoundError:
sys.exit(ERROR_CORRPUTED_CONFIG)
def handle_change_mode(args):
config = read_config()
if args.m:
if args.m == "gui":
config["mode"] = "gui"
elif args.m == "terminal":
config["mode"] = "terminal"
else:
sys.exit("Not a valid mode option. Only 'gui' and 'terminal' are valid")
return config
def handle_terminal(args):
config = read_config()
2024-05-20 11:24:27 +00:00
# if a prompt exists
if args.p:
try:
bot = TerminalBot(args.p, args.f, **config["ollamaConfig"])
bot.start()
except TypeError:
2024-05-20 09:07:57 +00:00
sys.exit(ERROR_CORRPUTED_CONFIG)
elif args.f:
sys.exit("failure: prompt needed")
else:
sys.exit("usage in terminal mode: project.py -p 'prompt' and optional: -f 'filename'")
if __name__ == "__main__":
main()