the 'main' project part and its test functions. handling the command line arguments
This commit is contained in:
parent
a7585d3a50
commit
774b38bba4
|
|
@ -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()
|
||||||
|
|
@ -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()
|
||||||
Loading…
Reference in New Issue