From 774b38bba4c54a4dbd42a513a356341beb4f005a Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Sat, 18 May 2024 21:05:37 +0200 Subject: [PATCH] the 'main' project part and its test functions. handling the command line arguments --- project.py | 112 +++++++++++++++++++++++++++++++++++++++++ test_project.py | 131 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 243 insertions(+) create mode 100644 project.py create mode 100644 test_project.py diff --git a/project.py b/project.py new file mode 100644 index 0000000..475b707 --- /dev/null +++ b/project.py @@ -0,0 +1,112 @@ +import argparse +from scripts.TerminalChat import TerminalBot +from scripts.GUIChat import ChatGUI +import json +import os +import sys + +CONFIG_FILE = "config/config.json" + +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 config is None and args.config is False: + sys.exit("No Config available. please run: 'python project.py --config' to set it up") + + elif args.config is True: + config = configure() + write_config(config) + elif args.m: + config = handle_change_mode(args) + write_config(config) + elif config["mode"] == "terminal": + handle_terminal(args) + elif config["mode"] == "gui": + try: + config["ollamaConfig"]["base_header"] = json.loads(config["ollamaConfig"]["base_header"]) + config["ollamaConfig"]["embeddings_header"] = json.loads(config["ollamaConfig"]["embeddings_header"]) + except json.decoder.JSONDecodeError: + """can be ignored if no header needed""" + pass + # start gui + try: + gui = ChatGUI(**config["ollamaConfig"]) + gui.mainloop() + except TypeError: + sys.exit("The config file seems to be corrupted, please run: 'python project.py --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 = input("Authentication for base model (standard: empty): ") or "" + embeddings_header = input("Authentication 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, "embeddings_header": embeddings_header}} + +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): + with open(CONFIG_FILE, "w") as config_file: + json.dump(config, config_file, indent=4) + +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() + try: + config["ollamaConfig"]["base_header"] = json.loads(config["ollamaConfig"]["base_header"]) + config["ollamaConfig"]["embeddings_header"] = json.loads(config["ollamaConfig"]["embeddings_header"]) + except json.decoder.JSONDecodeError: + """can be ignored if no header needed""" + pass + + if args.p: + try: + bot = TerminalBot(args.p, args.f, **config["ollamaConfig"]) + bot.start() + except TypeError: + sys.exit("The config file seems to be corrupted, please run: 'python project.py --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() \ No newline at end of file diff --git a/test_project.py b/test_project.py new file mode 100644 index 0000000..05ef5bd --- /dev/null +++ b/test_project.py @@ -0,0 +1,131 @@ +import pytest +import json +from pathlib import Path +from project import configure, read_config, write_config, handle_change_mode + +CONFIG_FILE = 'tests/config.json' + + +@pytest.fixture(scope='function') +def setup_config(): + """Fixture to create a dummy config file before each test and remove it after.""" + # Create the config file + initial_config = { + "mode": "terminal", + "ollamaConfig": { + "base_url": "http://localhost:11434", + "embeddings_url": "http://localhost:11434", + "base_model": "mistral", + "embeddings_model": "mxbai-embed-large", + "base_header": "", + "embeddings_header": "" + } + } + with open(CONFIG_FILE, 'w') as f: + json.dump(initial_config, f) + + yield + + # remove the config file after the test + if Path(CONFIG_FILE).exists(): + Path(CONFIG_FILE).unlink() + +def test_configure(monkeypatch): + inputs = iter([ + "I want gui", # Incorrect value + 'terminal', # Correct input after re-prompt + 'https://ai.fabelous.app/v1/ollama/generic', # Base LLM URL + 'http://localhost:11434', # Embeddings URL + 'mistral', # Base model + 'mxbai-embed-large', # Embeddings model + '{"Authorization": "Token xzy"}', # Base header for authentication + '{"Authorization": "Token xzy"}', # Embeddings header for authentication + ]) + + monkeypatch.setattr('builtins.input', lambda _: next(inputs)) + + config = configure() + + # Expected configurations based on the inputs + expected_config = { + "mode": "terminal", + "ollamaConfig": { + "base_url": "https://ai.fabelous.app/v1/ollama/generic", + "embeddings_url": "http://localhost:11434", + "base_model": "mistral", + "embeddings_model": "mxbai-embed-large", + "base_header": '{"Authorization": "Token xzy"}', + "embeddings_header": '{"Authorization": "Token xzy"}' + } + } + + assert config['mode'] == expected_config['mode'], "Mode configuration does not match." + assert config['ollamaConfig'] == expected_config['ollamaConfig'], "OllamaConfig does not match." + + +def test_read_config(setup_config, monkeypatch): + """Test with both existing and non-existing config files.""" + + # non existing test file + monkeypatch.setattr('project.CONFIG_FILE', 'non_existing_config.json') + + result = read_config() + + assert result is None, "The function should return None for a non-existing config file." + + # existing test file + monkeypatch.setattr('project.CONFIG_FILE', CONFIG_FILE) + + config = read_config() + + assert isinstance(config, dict), "The returned configuration is not a dictionary." + assert set(config.keys()) == {'mode', 'ollamaConfig'}, "The returned configuration does not contain expected keys." + + + + + +def test_handle_change_mode(setup_config, monkeypatch): + """Test to correctly update the config for valid modes.""" + monkeypatch.setattr('project.CONFIG_FILE', CONFIG_FILE) + + class Args: + pass + + args = Args() + args.m = 'gui' + updated_config = handle_change_mode(args) + assert updated_config['mode'] == 'gui', "Mode should be updated to 'gui'" + + args.m = 'terminal' + updated_config = handle_change_mode(args) + assert updated_config['mode'] == 'terminal', "Mode should be updated to 'terminal'" + + args.m = 'invalid_mode' + with pytest.raises(SystemExit) as e: + handle_change_mode(args) + assert "Not a valid mode option. Only 'gui' and 'terminal' are valid" in str(e.value) + + + +def test_write_config(monkeypatch): + """Test to ensure changes are written to the config file correctly.""" + monkeypatch.setattr('project.CONFIG_FILE', CONFIG_FILE) + new_config = { + "mode": "terminal", + "ollamaConfig": { + "base_url": "http://localhost:11434", + "embeddings_url": "https://ai.fabelous.app/v1/ollama/generic", + "base_model": "mistral", + "embeddings_model": "mxbai-embed-large", + "base_header": '{"Authorization": "Token xzy"}', + "embeddings_header": '{"Authorization": "Token xzy"}' + } + } + write_config(new_config) + + # Read the file directly to check if the write was successful + with open(CONFIG_FILE, 'r') as f: + config = json.load(f) + assert config == new_config, "The config file was not updated correctly" + Path(CONFIG_FILE).unlink()