from setuptools import setup, find_packages

with open("requirements.txt") as f:
    requirements = f.read().splitlines()

setup(
    name="chat",
    version="0.1",
    packages=find_packages(where="src"),
    package_dir={"": "src"},
    install_requires=requirements,
    requires=["python (>=3.9)"],
)

setup.py
isak-b/chat
import os
import sys

root_path = os.path.dirname(os.path.abspath(__file__)) + "/"
sys.path.insert(0, root_path)

src/__init__.py
isak-b/chat
from bot import ChatBot


def main():
    ChatBot().cmdloop()


if __name__ == "__main__":
    main()

src/__main__.py
isak-b/chat
import os
import openai
from dotenv import load_dotenv

load_dotenv()
ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/"


def load_openai_api_key(env_filename: str) -> None:
    """Load OpenAI API key from .env file"""
    path = ROOT_PATH + env_filename
    if os.path.exists(path) is False:
        raise FileNotFoundError(f"{path} file not found")
    if os.getenv("OPENAI_API_KEY") is None:
        raise ValueError("OPENAI_API_KEY not found in .env file")

    # Set API key
    openai.api_key = os.getenv("OPENAI_API_KEY")

src/api.py
isak-b/chat
import cmd
from colorama import Fore, Style

from model import get_model
from messages import (
    load_prompt,
    load_history,
    get_question,
    get_answer,
    get_msgs,
    print_history,
    save_history,
    save_msg,
)
from config import load_cfg, save_cfg, print_cfg, process_cfg_value

default_cfg = load_cfg()


class ChatBot(cmd.Cmd):
    intro = "Entering chat session: type '/h' for help, '/c' for commands, '/q' to quit"
    prompt = "> "
    last_line = ""

    def __init__(self, cfg: dict = default_cfg):
        super().__init__()
        self.cfg = cfg
        self.model = get_model(model_name=self.cfg["model"], env_filename=self.cfg["files"]["env"])
        self.chat_prompt = load_prompt(path=self.cfg["dirs"]["prompt"], filename=f"{self.cfg['prompt']}.txt")
        self.chat_history = load_history(path=self.cfg["dirs"]["history"], filename=f"{self.cfg['history']}.json")

        # Add aliases for commands
        self.commands = {
            ChatBot.do_help: ["h", "help"],
            ChatBot.do_quit: ["q", "quit", "exit"],
            ChatBot.do_save: ["s", "save"],
            ChatBot.do_history: ["hs", "history"],
            ChatBot.do_clear_history: ["ch", "clear_history"],
            ChatBot.do_config: ["cfg", "config"],
            ChatBot.do_commands: ["c", "commands"],
        }
        for func, commands in self.commands.items():
            for command in commands:
                setattr(ChatBot, f"do_{command}", func)

    def parseline(self, line: str):
        """Override parseline to allow for '/' commands"""
        if line.startswith("/"):
            return super().parseline(line[1:])
        return None, None, line

    def do_quit(self, _arg):
        """Quit chat session: /q, /quit, /exit"""
        if self.cfg["save_history_on_exit"] is True:
            save_history(self.chat_history, path=self.cfg["dirs"]["history"], filename=f"{self.cfg['history']}.json")
        if self.cfg["save_config_on_exit"] is True:
            save_cfg(self.cfg)
        return True

    def do_save(self, filename: str = None):
        """Save the last message written in console to file: /save <filename: str>"""
        save_msg(self.last_line, path=self.cfg["dirs"]["saved"], filename=filename)

    def do_history(self, i: int = None):
        """Print chat history: /history <i: int>"""
        chat_history = self.chat_history or [{"content": "No chat history"}]
        try:
            chat_history = [chat_history[-1 * int(i)]] if i else chat_history
        except IndexError:
            return print(f"{Fore.RED}>> IndexError: index {i} out of range for {len(chat_history)=}{Style.RESET_ALL}")
        self.last_line = chat_history[-1]["content"]
        print_history(chat_history)

    def do_clear_history(self, _arg):
        """Clear chat history: /clear_history, /ch
        NOTE: This does not automatically delete the history file, but the file may be deleted or overwritten on exit
        """
        self.chat_history = []
        print(f"{Fore.RED}>> History cleared{Style.RESET_ALL}")

    def do_config(self, arg: str = None) -> None:
        """Set config values: /config <key: str>=<value: str>"""
        if not arg:
            return print_cfg(self.cfg)

        args = arg.split("=")
        key, value = args[0], args[1]
        value = process_cfg_value(value)

        # Key validation
        if key not in self.cfg:
            print(f"{Fore.RED}>> {key=} not found in cfg:{Style.RESET_ALL}")
            return print_cfg(self.cfg)
        try:
            if key == "prompt":
                self.chat_prompt = load_prompt(path=self.cfg["dirs"]["prompt"], filename=f"{value}.txt")
            elif key == "history":
                self.chat_history = load_history(path=self.cfg["dirs"]["history"], filename=f"{value}.json")
            elif key == "model":
                self.model = get_model(model_name=value, env_filename=self.cfg["files"]["env"])
        except FileNotFoundError:
            return print(f"{Fore.RED}>> file '{self.cfg['paths'][key]}{value}.*' not found{Style.RESET_ALL}")

        # Set new config value
        self.cfg[key] = value
        print_cfg(self.cfg, highlight_key=key)

    def do_commands(self, _arg) -> None:
        """View available commands: /commands, /c"""
        for func, aliases in self.commands.items():
            print(f"{Fore.CYAN}>> {func.__name__.split('_')[1]}: {[f'/{a}' for a in aliases]}{Style.RESET_ALL}")

    def default(self, line) -> None:
        """Default cmdloop: Get question from user and answer from bot"""
        question = get_question(line)
        msgs = get_msgs(self.chat_prompt, self.chat_history, question, history_size=self.cfg["history_size"])
        answer = get_answer(msgs, model=self.model)
        self.chat_history.extend([question, answer])
        self.last_line = answer["content"]

src/bot.py
isak-b/chat
import os
import yaml
from colorama import Fore, Style

ROOT_PATH = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) + "/"
cfg_path = os.path.join(ROOT_PATH, "config.yaml")


def load_cfg(path: str = cfg_path) -> dict:
    """Load config from file"""
    with open(path, "r") as f:
        cfg = yaml.safe_load(f)
    return cfg


def save_cfg(cfg: dict, path: str = cfg_path) -> None:
    """Save config to file"""
    with open(path, "w") as f:
        yaml.dump(cfg, f)


def print_cfg(cfg: dict, highlight_key: str = None) -> None:
    """Print config to console"""
    cfg_str = ""
    for line in yaml.dump(cfg, sort_keys=False).split("\n"):
        if highlight_key and line.startswith(highlight_key):
            cfg_str += f"{Fore.CYAN}{Style.BRIGHT}>> {line}{Style.RESET_ALL}\n"
        else:
            cfg_str += f"{Fore.LIGHTBLACK_EX}{line}{Style.RESET_ALL}\n"
    print(cfg_str)


def process_cfg_value(value: str) -> str | int | bool:
    """Process config value"""
    if value.lower() in ["true", "false"]:
        value = value.lower() == "true"
    elif value.isdigit():
        value = int(value)
    return value

src/config.py
isak-b/chat
import os
import json
from colorama import Fore, Style

ROOT_PATH = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/"


def load_prompt(path: str, filename: str) -> dict:
    """Load prompt from file"""
    with open(ROOT_PATH + path + filename, "r") as f:
        content = f.read()
    return {"role": "system", "content": content}


def load_history(path: str, filename: str) -> list:
    """Load chat history from file"""
    result = []
    if os.path.exists(ROOT_PATH + path + filename) is True:
        with open(ROOT_PATH + path + filename, "r") as f:
            result = json.load(f)
    return result


def get_question(args: list) -> str:
    """Get question from args"""
    content = " ".join(args) if isinstance(args, list) else str(args)
    return {"role": "user", "content": content}


def get_msgs(prompt: dict, history: list, question: dict, history_size: int = 10) -> list:
    """Get messages for bot to respond to"""
    history = history[-history_size:] if history_size else history
    msgs = [prompt] + history + [question]
    return msgs


def get_answer(msgs: list, model, print_answer: bool = True) -> dict:
    """Get answer from bot"""
    answer = model(messages=msgs).choices[0]["message"]
    if print_answer is True:
        print(f"{Fore.CYAN}{Style.BRIGHT}>> {answer['content']}{Style.RESET_ALL}")
    return answer


def print_history(history: list) -> None:
    """Print chat history to console"""
    for i, msg in enumerate(history):
        print(f"{Fore.LIGHTBLACK_EX}>> history[{len(history) - i}]: {msg['content']}{Style.RESET_ALL}")


def save_msg(msg: str, path: str, filename: str = None) -> None:
    """Save message to file"""
    if not filename:
        files = os.listdir(ROOT_PATH + path)
        max_num = 0
        for file in files:
            if file.startswith("output_"):
                try:
                    num = int(file.split("_")[1].split(".")[0])
                except ValueError:
                    continue
                max_num = max(num, max_num)
        new_num = max_num + 1
        filename = f"output_{str(new_num).zfill(3)}.txt"

    output_path = f"{ROOT_PATH}{path}{filename}"
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, "w") as f:
        f.write(msg)
    print(f"{Fore.GREEN}>> Output saved: {output_path}{Style.RESET_ALL}")


def save_history(history: list, path: str, filename: str) -> None:
    """Save chat history to file"""
    output_path = f"{ROOT_PATH}{path}{filename}"
    if history:
        with open(output_path, "w") as f:
            json.dump(history, f)
            print(f"{Fore.GREEN}>> History saved: {output_path}{Style.RESET_ALL}")
    else:
        if os.path.exists(output_path):
            os.remove(output_path)
            print(f"{Fore.RED}>> No history to save; file removed: {output_path}{Style.RESET_ALL}")

src/messages.py
isak-b/chat
import functools
import openai

from api import load_openai_api_key


class MockModel:
    content = "This is a mock reply.\nYou can change to the default model with /config model=default"
    choices: list = [{"message": {"role": "assistant", "content": content}}]

    def __call__(self, *_args, **_kwargs):
        return self


gpt3_5 = functools.partial(
    openai.ChatCompletion.create,
    model="gpt-3.5-turbo",
)
gpt4 = functools.partial(
    openai.ChatCompletion.create,
    model="gpt-4",
)
models = {
    "default": gpt3_5,
    "mock": MockModel(),
    "gpt3.5": gpt3_5,
    "gpt4": gpt4,
}


def get_model(model_name: str, env_filename: str = None):
    model = models[model_name]
    if model_name != "mock":
        load_openai_api_key(env_filename=env_filename)
    return model

src/model.py
isak-b/chat
import os
import sys

import pytest

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/"
sys.path.insert(0, root_path)

from src.config import load_cfg  # noqa: E402

test_cfg = load_cfg()


@pytest.fixture(scope="session")
def cfg():
    return test_cfg

tests/conftest.py
isak-b/chat
import openai

from src.api import load_openai_api_key


def test_load_openai_api_key(cfg):
    load_openai_api_key(env_filename=cfg["files"]["env"])
    assert openai.api_key is not None
    assert openai.api_key[:3] == "sk-"

tests/test_api.py
isak-b/chat
import io
import tempfile
from contextlib import redirect_stdout

from src.config import save_cfg, load_cfg, print_cfg

buffer = io.StringIO()


def test_config():
    with tempfile.NamedTemporaryFile() as temp_file:
        temp_file_path = temp_file.name

        # Test save_cfg
        cfg = {"key_a": "val_a"}
        save_cfg(cfg, temp_file_path)

        # Test load_cfg
        loaded_cfg = load_cfg(temp_file_path)
        assert loaded_cfg == cfg

        # Test print_cfg
        with redirect_stdout(buffer):
            print_cfg(cfg)
        printed_value = buffer.getvalue()
        assert "key_a: val_a" in printed_value

tests/test_config.py
isak-b/chat
from src.model import get_model, models


def test_get_gpt3_5_model(cfg):
    model = get_model(model_name="gpt3.5", env_filename=cfg["files"]["env"])
    assert model == models["gpt3.5"]


def test_get_gpt4_model(cfg):
    model = get_model(model_name="gpt4", env_filename=cfg["files"]["env"])
    assert model == models["gpt4"]


def test_get_mock_model():
    model = get_model(model_name="mock")
    assert model == models["mock"]

tests/test_model.py
isak-b/chat