Fabelous-Ai-Chat/project.py

import argparse
from scripts.TerminalChat import TerminalBot
from scripts.GUIChat import ChatGUI
import json
import os
import sys

CONFIG_FILE = "config/config.json"
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()
    # 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")
    # add config
    elif args.config is True:
       config = configure()
       write_config(config)
    # change ui mode
    elif args.m:
        config = handle_change_mode(args)
        write_config(config)
    # start terminal ui
    elif config["mode"] == "terminal":
        handle_terminal(args)
    # start customtkinter ui
    elif config["mode"] == "gui":
        # start gui
        try:
            gui = ChatGUI(**config["ollamaConfig"])
            gui.mainloop()
        except TypeError:
            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"
    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 ""

    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)


def write_config(config: dict):
    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()
    # if a prompt exists
    if args.p:
        try:
            bot = TerminalBot(args.p, args.f, **config["ollamaConfig"])
            bot.start()
        except TypeError:
            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()