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