Gives the ability to use ChatGPT directly from the command line on the Xonsh Shell
import pytest
from xonsh.xontribs import xontribs_unload, xontribs_load

@pytest.fixture
def loaded_session(xession):
    xontribs_load(["chatgpt"])
    yield xession
    xontribs_unload(["chatgpt"])

def test_it_loads(loaded_session):
    assert "chatgpt" in loaded_session.aliases
    assert 'chatgpt?' in loaded_session.aliases
    assert 'chat-manager' in loaded_session.aliases

def test_it_unloads(loaded_session):
    xontribs_unload(["chatgpt"])
    assert "chatgpt" not in loaded_session.aliases
    assert 'chatgpt?' not in loaded_session.aliases
    assert 'chat-manager' not in loaded_session.aliases
tests/xontrib/test_chatgpt.py
jpal91/xontrib-chatgpt
import shutil
import pytest
from xontrib_chatgpt.events import chat_events

@pytest.fixture
def cm_events(xession):
    events = xession.builtins.events
    for c in chat_events:
        events.doc(*c)
    yield events

@pytest.fixture(scope="module")
def temp_home(tmpdir_factory):
    home = tmpdir_factory.mktemp("home")
    home.mkdir("expected")
    home.mkdir("saved")
    data_dir = home.mkdir("data_dir")
    data_dir.mkdir("chatgpt")
    fixtures = [
        "color_convo.txt",
        "no_color_convo.txt",
        "no_color_convo2.txt",
        "convo.json",
        "convo2.json",
        "long_convo.txt",
    ]
    for f in fixtures:
        shutil.copy(f"tests/fixtures/{f}", f"{home}/expected/{f}")
    shutil.copy(
        f"tests/fixtures/no_color_convo.txt", f"{data_dir}/chatgpt/no_color_convo.txt"
    )
    yield home
tests/xontrib_chatgpt/conftest.py
jpal91/xontrib-chatgpt
import pytest
from argparse import Namespace

from xontrib_chatgpt.args import _gpt_parse


@pytest.fixture(scope="module")
def argparse():
    return _gpt_parse()


@pytest.fixture
def default_namespace_dict():
    return {
        "cmd": "send",
        "text": [],
        "mode": "color",
        "type": "text",
        "name": "",
        "path": "",
        "n": 10,
    }


@pytest.mark.parametrize(
    ("args", "expected"),
    [
        ("something else", {"cmd": "send", "text": ["something", "else"]}),
        ("-p", {"cmd": "print", "text": []}),
        ("-s", {"cmd": "save", "text": []}),
        ("-P path", {"cmd": "send", "text": [], "path": "path"}),
        ("-P path -s", {"cmd": "save", "text": [], "path": "path"}),
        (
            "-P path -s --name name",
            {"cmd": "save", "text": [], "path": "path", "name": "name"},
        ),
        (
            "-P path -s --name name -m json",
            {"cmd": "save", "text": [], "path": "path", "name": "name", "mode": "json"},
        ),
        (
            "-P path -s --name name -m no-color -n 5",
            {
                "cmd": "save",
                "text": [],
                "path": "path",
                "name": "name",
                "mode": "no-color",
                "n": 5,
            },
        ),
        (
            "-P path -s --name name -m color -n 10 something else",
            {
                "cmd": "save",
                "text": ["something", "else"],
                "path": "path",
                "name": "name",
                "mode": "color",
                "n": 10,
            },
        ),
        (
            "-P path -s --name name -m color -t json",
            {
                "cmd": "save",
                "text": [],
                "path": "path",
                "name": "name",
                "mode": "color",
                "type": "json",
            },
        ),
        (
            "-P path -p --name name -m color -t json something else",
            {
                "cmd": "print",
                "text": ["something", "else"],
                "path": "path",
                "name": "name",
                "mode": "color",
                "type": "json",
            },
        ),
        (
            "-P path -p --name name -m color -t text -n 5 something else",
            {
                "cmd": "print",
                "text": ["something", "else"],
                "path": "path",
                "name": "name",
                "mode": "color",
                "type": "text",
                "n": 5,
            },
        ),
    ],
)
def test_parse_args(xession, argparse, args, expected, default_namespace_dict):
    new_dict = {**default_namespace_dict, **expected}
    expected = Namespace(**new_dict)
    assert argparse.parse_args(args.split()) == expected

tests/xontrib_chatgpt/test_args.py
jpal91/xontrib-chatgpt
import io
import json
import shutil
import pytest
from datetime import datetime
from openai.error import RateLimitError
from xontrib_chatgpt.chatgpt import ChatGPT, parse_convo, get_token_list
from xontrib_chatgpt.exceptions import (
    NoApiKeyError,
    UnsupportedModelError,
    NoConversationsError,
    InvalidConversationsTypeError,
)


MARKDOWN_BLOCK = """\
Hello!
```python
print('Hello World!')
```
"""

MARKDOWN_BLOCK_2 = """\
Sure here is a `Hello world` function:
```python
def hello_world():
    print('Hello world!')
```
"""


class DummyAI:
    def __init__(self):
        self.api_key = None

    def create(self, **_):
        return {
            "choices": [{"message": {"content": "test", "role": "assistant"}}],
            "usage": {"prompt_tokens": 1, "completion_tokens": 1},
        }


@pytest.fixture(scope="module")
def temp_home(tmpdir_factory):
    home = tmpdir_factory.mktemp("home")
    home.mkdir("expected")
    home.mkdir("saved")
    data_dir = home.mkdir("data_dir")
    data_dir.mkdir("chatgpt")
    fixtures = [
        "color_convo.txt",
        "no_color_convo.txt",
        "no_color_convo2.txt",
        "convo.json",
        "convo2.json",
        "long_convo.txt",
    ]
    for f in fixtures:
        shutil.copy(f"tests/fixtures/{f}", f"{home}/expected/{f}")
    shutil.copy(
        f"tests/fixtures/no_color_convo.txt", f"{data_dir}/chatgpt/no_color_convo.txt"
    )
    yield home


@pytest.fixture
def monkeypatch_openai(monkeypatch):
    dummy_ai = DummyAI()
    dummy_ai.ChatCompletion = DummyAI()
    monkeypatch.setattr("xontrib_chatgpt.chatgpt.openai", dummy_ai)


@pytest.fixture
def chat(xession):
    return ChatGPT()


@pytest.fixture
def chat_w_alias(xession):
    chat = ChatGPT("gpt")
    yield chat
    del chat


def test_it_loads(load_xontrib, xession):
    load_xontrib("chatgpt")
    assert "chatgpt" in xession.aliases


def test_alias_creation(xession):
    assert "chat" not in xession.aliases
    chat = ChatGPT()
    assert "chat" not in xession.aliases
    chat = ChatGPT("chat")
    assert "chat" in xession.aliases
    del xession.aliases["chat"]


def test_defualt_attribs(xession, chat):
    assert chat.messages == []
    assert chat._tokens == []
    assert chat._max_tokens == 3000
    assert chat.alias == ""
    assert chat.tokens == 53


def test_tokens(xession, chat):
    chat._tokens = [1, 2, 3]
    assert chat.tokens == 59


def test_chat_raises_error_with_no_api_key(xession, chat, monkeypatch_openai):
    xession.env["OPENAI_API_KEY"] = ""
    with pytest.raises(NoApiKeyError):
        chat.chat("test")


def test_chat_raises_error_with_no_chat_model(xession, chat, monkeypatch_openai):
    xession.env["OPENAI_API_KEY"] = "test"
    xession.env["OPENAI_CHAT_MODEL"] = "test"
    with pytest.raises(UnsupportedModelError):
        chat.chat("test")

def test_chat_catches_openai_errors(xession, chat, monkeypatch):
    xession.env["OPENAI_API_KEY"] = "test"
    xession.env["OPENAI_CHAT_MODEL"] = "gpt-3.5-turbo"
    def raise_it(*_, **__):
        raise RateLimitError('test')
    monkeypatch.setattr('xontrib_chatgpt.chatgpt.openai.ChatCompletion.create', raise_it)
    with pytest.raises(SystemExit):
        chat.chat("test")

def test_chat_convo(xession, chat):
    assert chat.chat_convo == chat.base
    chat.messages =[
        {"role": "user", "content": "test"},
        {"role": "assistant", "content": "test"},
    ]
    chat.chat_idx -= 2
    assert chat.chat_convo == chat.base + chat.messages

def test_chat_response(xession, monkeypatch_openai, chat):
    xession.env["OPENAI_API_KEY"] = "test"
    assert chat.chat_idx == 0
    chat.chat("test") == "test"
    assert chat.messages == [
        {"role": "user", "content": "test"},
        {"role": "assistant", "content": "test"},
    ]
    assert chat._tokens == [1, 1]
    assert chat.tokens == 55
    assert chat.chat_idx == -2

@pytest.mark.skip()
def test_trim(xession, chat):
    chat._tokens = [1000, 1000, 900]
    chat.messages = ["test", "test", "test"]
    chat._trim()
    assert len(chat._tokens) == 3
    chat._tokens.append(1000)
    chat.messages.append("test")
    chat._trim()
    assert len(chat._tokens) == 3
    assert len(chat.messages) == 3

def test_trim_convo(xession, chat):
    toks = chat._tokens = [1000, 1000, 900]
    idx = chat.chat_idx = -3
    chat.trim_convo()
    assert chat._tokens == toks
    assert chat.chat_idx == idx
    chat._tokens.append(1000)
    chat.chat_idx -= 1
    chat.trim_convo()
    assert chat._tokens == toks
    assert chat.chat_idx == idx


def test_set_base_msgs(xession, chat):
    assert chat._base_tokens == 53
    chat.base = [{'role': 'system', 'content': 'test'}]
    assert chat._base_tokens == 8


def test__get_json_convo(xession, chat):
    chat.messages.append({"role": "user", "content": "test"})
    res = chat._get_json_convo(n=1)
    assert res == json.dumps([{"role": "user", "content": "test"}], indent=4)


@pytest.mark.parametrize(
    ("n", "n_expected", "color"), [(0, 5, True), (1, 1, True), (-1, 2, False)]
)
def test__get_printed_convo(xession, n, n_expected, color, chat):
    chat.messages.extend(
        [
            {"role": "user", "content": "test1"},
            {"role": "assistant", "content": "test2"},
            {"role": "user", "content": "`test3`"},
        ]
    )
    res = chat._get_printed_convo(n, color)
    assert len(res) == n_expected

    if color:
        assert "\x1b" in res[-1][1]
        assert "`" not in res[-1][1]
    else:
        assert "\x1b" not in res[-1][1]
        assert "`" in res[-1][1]


def test_print_convo_raises_no_convo_error(xession, chat):
    with pytest.raises(NoConversationsError):
        chat.print_convo(0)


def test_print_convo_raises_invalid_convo_error(xession, chat):
    chat.messages.append({"role": "user", "content": "test"})
    with pytest.raises(InvalidConversationsTypeError):
        chat.print_convo(0, mode="invalid")


@pytest.mark.parametrize(
    ("mode", "file"),
    [
        # ("color", "color_convo.txt"),
        ("text", "no_color_convo.txt"),
        ("json", "convo.json"),
    ],
)
def test_saves_convo(xession, chat, temp_home, mode, file, monkeypatch):
    monkeypatch.setenv("USER", "user")
    chat.messages.extend(
        [
            {"role": "user", "content": "Please write me a hello world function"},
            {"role": "assistant", "content": MARKDOWN_BLOCK_2},
        ]
    )
    chat.save_convo(temp_home / "saved" / file, mode=mode)
    assert (temp_home / "saved" / file).exists()
    with open(temp_home / "saved" / file, "r") as f:
        res = f.read().strip()
    with open(temp_home / "expected" / file, "r") as f:
        expected = f.read().strip()
    assert res == expected

def test_saves_with_override(xession, chat, temp_home, monkeypatch):
    monkeypatch.setenv("USER", "user")
    chat.messages.extend(
        [
            {"role": "user", "content": "Please write me a hello world function"},
            {"role": "assistant", "content": MARKDOWN_BLOCK_2},
            {"role": "user", "content": "test"},
        ]
    )
    chat.save_convo(temp_home / 'test.txt', mode='json')
    with open(temp_home / 'test.txt') as f:
        cur = json.load(f)
    chat.messages.pop()
    monkeypatch.setattr('builtins.input', lambda _: 'y')
    chat.save_convo(temp_home / 'test.txt', mode='json')
    with open(temp_home / 'test.txt') as f:
        new = json.load(f)
    assert cur != new
    cur = new
    chat.messages.pop()
    chat.save_convo(temp_home / 'test.txt', mode='json', override=True)
    with open(temp_home / 'test.txt') as f:
        new = json.load(f)
    assert cur != new

@pytest.mark.parametrize(
    ("alias", "json", "name"),
    [
        (True, False, "my_chat"),
        (True, False, ""),
        (False, False, ""),
        (True, True, ""),
        (True, True, "my_chat"),
    ],
)
def test_saves_convo_to_default_location(
    xession, alias, json, name, chat, chat_w_alias, temp_home, monkeypatch
):
    monkeypatch.setenv("USER", "user")
    xession.env["XONSH_DATA_DIR"] = str(temp_home / "data_dir")
    now = datetime.now().strftime("%Y-%m-%d")
    convo = [
        {"role": "user", "content": "Please write me a hello world function"},
        {"role": "assistant", "content": MARKDOWN_BLOCK_2},
    ]
    current_chat = chat_w_alias if alias else chat
    current_chat.messages.extend(convo)
    prefix, ext, mode = (
        f'user_{name or current_chat.alias or "chatgpt"}_{now}',
        ".json" if json else ".txt",
        "json" if json else "text",
    )

    for i in range(3):
        current_chat.save_convo(name=name, mode=mode)
        if i == 0:
            assert (temp_home / "data_dir" / "chatgpt" / f"{prefix}{ext}").exists()
        else:
            assert (temp_home / "data_dir" / "chatgpt" / f"{prefix}_{i}{ext}").exists()


def test_auto_create_alias(xession):
    assert "chat" not in xession.aliases
    chat = ChatGPT("chat")
    assert "chat" in xession.aliases


def test_auto_del_alias(xession):
    chat = ChatGPT("chat")
    assert "chat" in xession.aliases
    del chat
    assert "chat" not in xession.aliases


def test_cli_execution(xession, chat_w_alias, capsys, monkeypatch_openai):
    xession.env["OPENAI_API_KEY"] = "test"
    xession.aliases["gpt"](["hello", "my", "name", "is", "user"])
    out, err = capsys.readouterr()
    out = out.strip().split("\n    ")
    assert "ChatGPT" in out[0]
    assert "test" in out[1]


def test_cli_execution_print(
    xession, chat_w_alias, capsys, monkeypatch, monkeypatch_openai
):
    monkeypatch.setattr(
        "xontrib_chatgpt.chatgpt.ChatGPT.print_convo",
        lambda _, n, mode: print("print_convo", n, mode),
    )
    xession.aliases["gpt"](["-p"])
    out, err = capsys.readouterr()
    assert out.strip() == "print_convo 10 color"
    xession.aliases["gpt"]("-p -n 5 -m no-color".split())
    out, err = capsys.readouterr()
    assert out.strip() == "print_convo 5 no-color"


def test_cli_execution_save(
    xession, chat_w_alias, capsys, monkeypatch, monkeypatch_openai
):
    monkeypatch.setattr(
        "xontrib_chatgpt.chatgpt.ChatGPT.save_convo",
        lambda _, path, name, mode: print("save_convo", path, name, mode),
    )
    xession.aliases["gpt"](["-s"])
    out, err = capsys.readouterr()
    assert out.strip() == "save_convo   text"
    xession.aliases["gpt"]("-s -P path -n 5 -t json --name name".split())
    out, err = capsys.readouterr()
    assert out.strip() == "save_convo path name json"


def test_cli_execution_pipe(xession, chat_w_alias, capsys, monkeypatch_openai):
    xession.env["OPENAI_API_KEY"] = "test"
    stdin = io.StringIO()
    stdin.write("hello")
    xession.aliases["gpt"]([], stdin=stdin)
    out, err = capsys.readouterr()
    out = out.strip().split("\n    ")
    assert "ChatGPT" in out[0]
    assert "test" in out[1]


def test_enter_exit(xession, chat, capsys, monkeypatch_openai):
    xession.env["OPENAI_API_KEY"] = "test"
    exe = xession.execer.exec
    exe("with! chat:\n    hello", glbs={"chat": chat})
    out, err = capsys.readouterr()
    out = out.strip().split("\n    ")
    assert "ChatGPT" in out[0]
    assert "test" in out[1]


def test_loads_from_convo(xession, temp_home):
    chat_file = temp_home / "expected" / "no_color_convo2.txt"
    new_cls = ChatGPT.fromconvo(chat_file)
    assert isinstance(new_cls, ChatGPT)
    assert new_cls.base[0] == {'role': 'system', 'content': 'Test\n'}
    assert "Please write me a hello world function" in new_cls.messages[0]["content"]


def test_loads_from_convo_in_default_dir(xession, temp_home):
    xession.env["XONSH_DATA_DIR"] = str(temp_home / "data_dir")
    new_cls = ChatGPT.fromconvo("no_color_convo.txt")
    assert isinstance(new_cls, ChatGPT)
    assert "Please write me a hello world function" in new_cls.messages[0]["content"]


def test_loads_from_convo_raises_file_not_found(xession, temp_home):
    with pytest.raises(FileNotFoundError):
        ChatGPT.fromconvo("invalid.txt")

def test_loads_and_trims(xession, temp_home, monkeypatch):
    chat_file = temp_home / 'expected' / 'convo2.json'
    def trim_convo(self):
        while self.tokens > 650:
            self.chat_idx += 1
    monkeypatch.setattr('xontrib_chatgpt.chatgpt.ChatGPT.trim_convo', trim_convo)
    new_cls = ChatGPT.fromconvo(chat_file)
    assert new_cls.chat_idx == -1


@pytest.fixture
def inc_test(xession):
    xession.ctx["test"] = 0

    def _inc_test(**_):
        xession.ctx["test"] += 1

    return _inc_test


def test_on_chat_create_handler(xession, cm_events, inc_test):
    cm_events.on_chat_create(inc_test)
    inst = ChatGPT(managed=True)
    assert xession.ctx["test"] == 1


def test_on_chat_change(xession, cm_events, inc_test):
    cm_events.on_chat_create(inc_test)
    inst = ChatGPT(managed=True)
    assert xession.ctx["test"] == 1
    cm_events.on_chat_used(inc_test)
    inst([])
    assert xession.ctx["test"] == 2


def test_on_chat_destroy_handler(xession, cm_events, inc_test):
    cm_events.on_chat_create(inc_test)
    inst = ChatGPT(managed=True)
    assert xession.ctx["test"] == 1
    cm_events.on_chat_destroy(inc_test)
    del inst
    assert xession.ctx["test"] == 2

tests/xontrib_chatgpt/test_chatgpt.py
jpal91/xontrib-chatgpt
import pytest
from datetime import datetime
from textwrap import dedent
from xontrib_chatgpt.chatmanager import ChatManager, convert_to_sys
from xontrib_chatgpt.chatgpt import ChatGPT


@pytest.fixture(scope="module")
def temp_home(tmp_path_factory):
    home = tmp_path_factory.mktemp("home")
    (home / "data_dir").mkdir()
    (home / "data_dir" / "chatgpt").mkdir()
    (home / "data_dir" / "chatgpt" / "dummy").mkdir()
    return home


@pytest.fixture(autouse=True)
def data_dir(xession, temp_home):
    xession.env["XONSH_DATA_DIR"] = str(temp_home / "data_dir")


@pytest.fixture
def test_files(temp_home):
    test_files = ["test1.txt", "test2.txt", "test3.txt"]
    [(temp_home / "data_dir" / "chatgpt" / f).touch() for f in test_files]
    yield test_files
    [(temp_home / "data_dir" / "chatgpt" / f).unlink() for f in test_files]


@pytest.fixture
def cm():
    return ChatManager()


def test_update_inst_dict(xession, cm):
    insts = [
        ("test", ChatGPT("test")),
        ("test1", ChatGPT("test1")),
        ("test2", ChatGPT("test2")),
    ]
    for inst in insts:
        xession.ctx[inst[0]] = inst[1]

    cm._update_inst_dict()

    for inst in insts:
        assert hash(inst[1]) in cm._instances
        assert cm._instances[hash(inst[1])]["name"] == inst[0]
        assert cm._instances[hash(inst[1])]["alias"] == inst[1].alias
        assert cm._instances[hash(inst[1])]["inst"] == inst[1]


def test_on_chat_create_handler(xession, cm_events, cm):
    inst = ChatGPT("test_alias")
    xession.ctx["test_name"] = inst
    cm_events.on_chat_create(lambda *args, **kw: cm.on_chat_create_handler(*args, **kw))
    cm_events.on_chat_create.fire(inst=inst)
    assert cm._current == hash(inst)
    assert hash(inst) in cm._instances
    assert cm._instances[hash(inst)]["alias"] == "test_alias"
    assert cm._instances[hash(inst)]["inst"] == inst


def test_on_chat_destroy_handler(xession, cm_events, cm):
    inst = ChatGPT("test")
    xession.ctx["test"] = inst
    cm._update_inst_dict()
    cm._current = hash(inst)
    cm_events.on_chat_destroy(
        lambda *args, **kw: cm.on_chat_destroy_handler(*args, **kw)
    )
    cm_events.on_chat_destroy.fire(inst=inst)
    assert cm._current is None
    assert hash(inst) not in cm._instances


def test_on_chat_used_handler(xession, cm_events, cm):
    inst = ChatGPT("test")
    assert cm._current is None
    cm_events.on_chat_used(lambda *args, **kw: cm.on_chat_used_handler(*args, **kw))
    cm_events.on_chat_used.fire(inst=inst)
    assert cm._current == hash(inst)


def test_add(xession, cm):
    cm.add("test")
    assert "test" in xession.ctx
    assert isinstance(xession.ctx["test"], ChatGPT)
    assert "test" in xession.aliases


def test_add_with_conflicting_name(xession, cm):
    cm.add("test")
    res = cm.add("test")
    assert res == "Chat with that name already exists!"
    xession.ctx["glob"] = "something"
    res = cm.add("glob")
    assert "glob" in xession.ctx
    assert res == "Variable with that name already exists!"


def test_ls(xession, cm_events, cm, monkeypatch, capsys):
    monkeypatch.setattr(
        "xontrib_chatgpt.chatgpt.ChatGPT.stats", lambda *_, **__: "test_out"
    )
    cm_events.on_chat_create(lambda *args, **kw: cm.on_chat_create_handler(*args, **kw))
    cm.add("test1")
    cm.add("test2")
    assert len(cm._instances.keys()) == 2
    res = cm.ls()
    assert res == "test_out\n\ntest_out"


def test_find_saved(xession, cm, temp_home, test_files):
    res = cm._find_saved()
    assert len(res) == 3
    assert sorted(res) == test_files


@pytest.mark.parametrize(
    ("input", "fname", "name"),
    [
        ("test1", "something_test1.txt", "test1"),
        ("something_test2_01-02-03.json", "something_test2_01-02-03.json", "test2"),
        ("test3", "test3.txt", "test3"),
    ],
)
def test_find_path_from_name(xession, cm, input, fname, name, temp_home):
    data_dir = temp_home / "data_dir"
    chat_dir = data_dir / "chatgpt"
    (chat_dir / fname).touch()
    res = cm._find_path_from_name(input)
    assert res == (str(chat_dir / fname), name)
    (chat_dir / fname).unlink()


def test_choose_from_multiple(xession, cm, monkeypatch):
    monkeypatch.setattr("builtins.input", lambda *_, **__: "2")
    res = cm._choose_from_multiple([("test1", "file1"), ("test2", "file2")])
    assert res == ("test2", "file2")


@pytest.mark.parametrize(
    ("inp", "name"),
    [
        ("data_dir/chatgpt/test1.txt", "test1"),
        ("test2.txt", "test2"),
        ("test3", "test3"),
    ],
)
def test_load(xession, cm, inp, name, temp_home, test_files):
    data_dir = temp_home / "data_dir"
    inp = inp.replace("data_dir", str(data_dir))
    res = cm.load(inp)
    assert name in xession.ctx
    assert f"Loaded chat {name}" in res


def test_load_with_conflicting_name(xession, cm, test_files, temp_home):
    cm.add("test1")
    res = cm.load("test1")
    assert "Loaded chat test10" in res
    assert "test10" in xession.ctx
    assert "test1" in xession.ctx


def test_save(xession, cm, test_files, temp_home, cm_events, monkeypatch):
    monkeypatch.setenv("USER", "user")
    cm_events.on_chat_create(lambda *args, **kw: cm.on_chat_create_handler(*args, **kw))
    with pytest.raises(SystemExit):
        cm.save()
    inst = ChatGPT(alias="new", managed=True)
    inst.messages += [{"role": "user", "content": "test"}]
    res = cm.save()
    assert res is None
    now = datetime.now().strftime("%Y-%m-%d")
    assert (temp_home / "data_dir" / "chatgpt" / f"user_new_{now}.txt").exists()


def test_save_returns_when_key_error(xession, cm):
    with pytest.raises(SystemExit) as s:
        cm.save("nonexistent")

    assert s.value.code == 1


@pytest.mark.parametrize(
    ("action", "args", "expected"),
    [
        ("add", ["add", "test"], (("test",), {})),
        ("ls", ["list"], ((), {"saved": False})),
        ("ls", ["ls", "-s"], ((), {"saved": True})),
        ("save", ["save"], ((), {"chat_name": "", "mode": "text", 'override': False})),
        (
            "save",
            ["save", "-m", "json", '-o', "name"],
            ((), {"chat_name": "name", "mode": "json", 'override': True}),
        ),
        ("load", ["load", "test"], (("test",), {})),
        ("print_chat", ["print"], ((), {"chat_name": "", "n": 10, "mode": "color"})),
        (
            "print_chat",
            ["print", "-n", "20", "-m", "json", "name"],
            ((), {"chat_name": "name", "n": 20, "mode": "json"}),
        ),
        ("help", ["help"], ((), {"tgt": ""})),
        ("help", ["help", "print_chat"], ((), {"tgt": "print_chat"})),
        (
            "edit",
            ["edit", "test"],
            ((), {"chat_name": "test", "sys_msgs": "", "no_code": False}),
        ),
        (
            "edit",
            ["edit", "-s", "[{'role': 'user'}]", "-C"],
            ((), {"chat_name": "", "sys_msgs": "[{'role': 'user'}]", "no_code": True}),
        ),
    ],
)
def test_cli(xession, cm, action, args, expected, monkeypatch):
    monkeypatch.setattr(
        f"xontrib_chatgpt.chatmanager.ChatManager.{action}",
        lambda _, *a, **k: setattr(cm, f"_{action}", (a, k)),
    )
    cm(args)
    assert getattr(cm, f"_{action}") == expected


def test_edit(xession, cm, cm_events):
    cm_events.on_chat_create(lambda *args, **kw: cm.on_chat_create_handler(*args, **kw))
    cm.add("test")
    sys_msg = [{"role": "system", "content": "Hello"}]
    cm.edit(sys_msgs=str(sys_msg))
    inst = cm.get_chat_by_name("test")["inst"]
    assert len(inst.base) == 2
    assert inst.base[1:] == sys_msg
    cm.edit(sys_msgs=str(sys_msg), no_code=True)
    assert len(inst.base) == 1
    assert inst.base == sys_msg

tests/xontrib_chatgpt/test_chatmanager.py
jpal91/xontrib-chatgpt
import pytest
from xontrib_chatgpt.chatmanager import ChatManager

from xontrib_chatgpt.events import (
    add_events,
    rm_events,
)

def test_add_and_rm_events(xession):
    cm = ChatManager()
    add_events(xession, cm)
    events = xession.builtins.events

    assert events.exists('on_chat_create')
    assert events.exists('on_chat_destroy')
    assert events.exists('on_chat_used')

    rm_events(xession)

    assert not events.exists('on_chat_create')
    assert not events.exists('on_chat_destroy')
    assert not events.exists('on_chat_used')


tests/xontrib_chatgpt/test_events.py
jpal91/xontrib-chatgpt
import json
import pytest
from textwrap import dedent

from xontrib_chatgpt.utils import (
    parse_convo,
    get_token_list,
    format_markdown,
    convert_to_sys,
)

@pytest.fixture
def sys_msgs():
    l = """[
    {'role': 'system', 'content': 'Hello'},
    {'role': 'system', 'content': 'Hi there!'},
    ]
    """

    d = '{"content": "Hello"}'

    y = dedent(
        """
    - role: system
      content: Hello
    - role: system
      content: Hi there!
    """
    )

    return l, d, y

MARKDOWN_BLOCK = """\
Hello!
```python
print('Hello World!')
```
"""

def test_format_markdown(xession):
    md = format_markdown(MARKDOWN_BLOCK)
    assert "\x1b" in md
    assert "```" not in md

def test_parses_json(xession, temp_home):
    json_path = temp_home / "expected" / "convo.json"
    with open(json_path) as f:
        exp_json = f.read()
    msg, base = parse_convo(exp_json)
    assert base + msg == json.loads(exp_json)


def test_parses_text(xession, temp_home):
    text_path = temp_home / "expected" / "long_convo.txt"
    with open(text_path) as f:
        exp_text = f.read()

    msgs, base = parse_convo(exp_text)
    assert len(msgs) == 6
    assert len(base) == 1

    for r in msgs:
        assert r["role"] in ["user", "assistant"]
        assert r["content"] != ""
    
    assert base[0] == {'role': 'system', 'content': 'This is a test.\n'}

def test_get_token_list(xession, temp_home):
    json_path = temp_home / "expected" / "convo2.json"
    with open(json_path) as f:
        exp_json = json.load(f)
    res = get_token_list(exp_json)
    assert len(res) == 7
    assert sum(res) == 835

def test_convert_to_sys(xession, sys_msgs):
    l, d, y = sys_msgs
    res = convert_to_sys(l)
    assert res == [
        {"role": "system", "content": "Hello"},
        {"role": "system", "content": "Hi there!"},
    ]
    res = convert_to_sys(d)
    assert res == [{"role": "system", "content": "Hello"}]
    res = convert_to_sys(y)
    assert res == [
        {"role": "system", "content": "Hello"},
        {"role": "system", "content": "Hi there!"},
    ]
tests/xontrib_chatgpt/test_utils.py
jpal91/xontrib-chatgpt
"""chatgpt xontrib"""
from xonsh.built_ins import XonshSession

from xontrib_chatgpt.chatgpt import ChatGPT
from xontrib_chatgpt.chatmanager import ChatManager
from xontrib_chatgpt.events import add_events, rm_events
from xontrib_chatgpt.completers import add_completers, rm_completers


__all__ = ()


def _load_xontrib_(xsh: XonshSession, **_):
    xsh.aliases["chatgpt"] = lambda args, stdin=None: ChatGPT.fromcli(args, stdin)
    xsh.aliases["chatgpt?"] = lambda *_, **__: xsh.help(ChatGPT)

    cm = ChatManager()
    xsh.aliases["chat-manager"] = lambda args, stdin=None: cm(args, stdin)
    xsh.aliases["chat-manager?"] = "chat-manager help"

    add_events(xsh, cm)
    add_completers()

    if "abbrevs" in xsh.ctx:
        xsh.ctx["abbrevs"]["cm"] = "chat-manager"

    return {"ChatGPT": ChatGPT, "chat_manager": cm}


def _unload_xontrib_(xsh: XonshSession, **_):
    del xsh.aliases["chatgpt"]
    del xsh.aliases["chatgpt?"]
    del xsh.aliases["chat-manager"]

    rm_events(xsh)
    rm_completers()

    if "abbrevs" in xsh.ctx:
        del xsh.ctx["abbrevs"]["cm"]
    

xontrib/chatgpt.py
jpal91/xontrib-chatgpt
"""Argument parsers for ChatGPT and ChatManager"""
from argparse import ArgumentParser


def _gpt_parse():
    """Argument parser for chatgpt"""
    cmd_parser = ArgumentParser(
        prog="chatgpt", description="Chat with OpenAI's ChatGPT from the command line"
    )
    cmd_parser.set_defaults(cmd="send")
    p_group = cmd_parser.add_argument_group(title="Print")
    p_group.add_argument(
        "-p",
        "--print",
        dest="cmd",
        const="print",
        action="store_const",
        help="Prints the conversation",
    )
    p_group.add_argument(
        "-n",
        type=int,
        default=10,
        help="Number of messasges to print. Default is 10. Use 0 for all.",
    )
    p_group.add_argument(
        "-m",
        "--mode",
        type=str,
        default="color",
        choices=["color", "no-color", "json"],
        help="Mode to print or save the conversation. Default is color",
    )
    s_group = cmd_parser.add_argument_group(title="Save")
    s_group.add_argument(
        "-s",
        "--save",
        dest="cmd",
        const="save",
        action="store_const",
        help="Saves the conversation",
    )
    s_group.add_argument(
        "-P",
        "--path",
        type=str,
        default="",
        help="File path to save the conversation. Default is $XONSH_DATA_DIR/chatgpt.",
    )
    s_group.add_argument(
        "--name",
        type=str,
        default="",
        help="Name of the conversation file. Default is chatgpt. Ignored when path is specified.",
    )
    s_group.add_argument(
        "-t",
        "--type",
        type=str,
        default="text",
        choices=["text", "json"],
        help="Type of the conversation file. Default is text.",
    )
    cmd_parser.add_argument(
        "text",
        nargs="*",
        default=None,
        help="Text to send to ChatGPT. Will be ignored when other cmd is used.",
    )

    return cmd_parser


def _cm_parse() -> ArgumentParser:
    """Argument parser for chat-manager"""
    parser = ArgumentParser(
        prog="chat-manager",
        description="Chat with OpenAI's ChatGPT from the command line",
    )
    parser.add_argument(
        "-C",
        help="Print information on the current/last used chat and exit",
        const=True,
        default=False,
        action="store_const",
    )

    subparser = parser.add_subparsers(
        dest="cmd", title="Available Commands", metavar="COMMAND"
    )

    p_add = subparser.add_parser(
        "add", help="Add/Create a chat", aliases=["a", "create"]
    )
    p_add.add_argument("name", type=str, help="Name of the chat to create", nargs=1)

    p_list = subparser.add_parser(
        "list", help="List all current or saved chats", aliases=["ls"]
    )
    p_list.add_argument(
        "-s",
        "--saved",
        action="store_true",
        help="List all saved chats from the default directory",
    )

    p_load = subparser.add_parser("load", help="Load a saved chat")
    p_load.add_argument(
        "name", type=str, help="Name or absolute path of the chat to load", nargs=1
    )

    p_save = subparser.add_parser("save", help="Save a chat")
    p_save.add_argument(
        "name",
        type=str,
        default="",
        help="Name of the chat to save. Defaults to last used/current chat",
        nargs="?",
    )
    p_save.add_argument(
        "-m",
        "--mode",
        type=str,
        default="text",
        choices=["text", "json"],
        help="Mode to print or save the conversation. Default is text",
    )
    p_save.add_argument(
        '-o',
        '--overwrite',
        action='store_true',
        help='If set, overwrites the existing file with the same name. Otherwise, the file will be saved with a number appended to the end of the name.',
        dest='override'
    )

    p_print = subparser.add_parser("print", help="Print a chat", aliases=["p"])
    p_print.add_argument(
        "name",
        type=str,
        help="Name of the chat to print. Defaults to last used/current chat",
        nargs="?",
        default="",
    )
    p_print.add_argument(
        "-n", type=int, default=10, help="Number of conversations to print"
    )
    p_print.add_argument(
        "-m",
        "--mode",
        type=str,
        default="color",
        choices=["color", "no-color", "json"],
        help="Mode to print or save the conversation. Default is color",
    )

    p_help = subparser.add_parser("help", help="Print help information", epilog="hello")
    p_help.add_argument(
        "target",
        type=str,
        default="",
        help="Target to print help information on. Defaults to chat-manager",
        nargs="?",
    )

    p_edit = subparser.add_parser('edit', help="Edit a chat's settings", aliases=['e'])
    p_edit.add_argument(
        'name',
        type=str,
        help="Name of the chat to edit. Defaults to last used/current chat",
        nargs="?",
        default="",
    )
    p_edit.add_argument(
        '-s',
        '--system',
        help='Edit base system messages/instructions to be sent to your chat instance. Must be a python list[dict], dict, or yaml equivalent. See documentation for more information.',
        dest='sys_msgs',
        default='',
    ),
    p_edit.add_argument(
        '-C',
        '--no-code',
        help='If set, the default instruction to add markdown code blocks will be removed. Only input system messages will be added.',
        dest='no_code',
        action='store_true'
    )

    return parser


if __name__ == "__main__":
    parser = _cm_parse()
    args = parser.parse_args()
    print(args)

xontrib_chatgpt/args.py
jpal91/xontrib-chatgpt
"""Main ChatGPT class"""
import sys
import os
import json
import weakref
from typing import TextIO
from xonsh.built_ins import XSH
from xonsh.tools import indent
from xonsh.contexts import Block
from xonsh.lazyasd import LazyObject
from xonsh.ansi_colors import ansi_partial_color_format
from openai.error import OpenAIError

from xontrib_chatgpt.args import _gpt_parse
from xontrib_chatgpt.lazyobjs import (
    _openai,
)
from xontrib_chatgpt.utils import (
    get_token_list,
    parse_convo,
    print_res,
    format_markdown,
    get_default_path
)
from xontrib_chatgpt.exceptions import (
    NoApiKeyError,
    UnsupportedModelError,
    NoConversationsError,
    InvalidConversationsTypeError,
)

openai = LazyObject(_openai, globals(), "openai")
parse = LazyObject(_gpt_parse, globals(), "parse")

DOCSTRING = """\
Allows for communication with ChatGPT from the xonsh shell.

This class can be assigned to a variable and used to have a continuous
    conversation with ChatGPT.
Alternatively, an alias 'chatgpt' is automatically registered 
    when the xontrib is loaded, which can be used to have one-off conversations.

Args:
    alias (str, optional): Alias to use for the instance. Defaults to ''.

Usage:
    # With the pre-registered 'chatgpt' alias
    >>> chatgpt [text] # text will be sent to ChatGPT
    >>> echo [text] | chatgpt # text from stdin will be sent to ChatGPT

    # With a ChatGPT instance and context manager
    >>> gpt = ChatGPT()
    >>> with! gpt:
            [text] # text will be sent to ChatGPT
    
    # With a ChatGPT instance and an alias
    >>> gpt = ChatGPT('gpt')
    >>> gpt [text] # text will be sent to ChatGPT
    >>> with! gpt:
            [text] # text will be sent to ChatGPT
    
    # Get Help
    >>> chatgpt? # With default alias
    >>> gpt = ChatGPT()
    # Instance related help / docstrings / src
    >>> gpt?
    >>> gpt.print_convo?
    # CLI/Alias related help
    >>> gpt -h
    >>> chatgpt -h

Methods:
    print_convo - Prints the current conversation to the shell
    save_convo - Saves the current conversation to a file

Envs:
    $OPENAI_API_KEY - OpenAI API Key (REQUIRED)
    $OPENAI_CHAT_MODEL - OpenAI Chat Model
        Default: gpt-3.5-turbo
        Supported: gpt-3.5-turbo, gpt-4

Default Commands/Aliases:
    chatgpt - Alias for ChatGPT.fromcli
    chatgpt? - Alias for xonsh help ChatGPT
    chatgpt -h - Command for CLI help

"""


class ChatGPT(Block):
    """Allows for communication with ChatGPT from the xonsh shell"""

    __xonsh_block__ = str

    def __init__(self, alias: str = "", managed: bool = False) -> None:
        """
        Initializes the ChatGPT instance

        Parameters
        ----------
        alias : str, optional
            Alias to use for the instance. Defaults to ''.
            Will automatically register a xonsh alias under XSH.aliases[alias] if provided.
        """

        self.alias = alias
        self._base: list[dict[str, str]] = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "system",
                "content": "If your responses include code, make sure to wrap it in a markdown code block with the appropriate language.\nExample:\n```python\nprint('Hello World!')\n```",
            },
        ]
        self.messages: list[dict[str, str]] = []
        self._base_tokens: int = 53
        self._tokens: list = []
        self._max_tokens = 3000
        self.chat_idx = 0
        self._managed = managed

        if self.alias:
            # Make sure the __del__ method is called despite the alias pointing to the instance
            ref = weakref.proxy(self)
            XSH.aliases[self.alias] = lambda args, stdin=None: ref(args, stdin)

        if managed:
            XSH.builtins.events.on_chat_create.fire(inst=self)

    def __enter__(self):
        res = self.chat(self.macro_block.strip())
        print_res(res)
        return self

    def __exit__(self, *_):
        del self.macro_block

    def __call__(self, args: list[str], stdin: TextIO = None):
        if self._managed:
            XSH.builtins.events.on_chat_used.fire(inst=self)

        if args:
            pargs = parse.parse_args(args)
        elif stdin:
            pargs = parse.parse_args(stdin.read().strip().split())
        else:
            return

        if pargs.cmd == "send":
            res = self.chat(" ".join(pargs.text))
            print_res(res)
        elif pargs.cmd == "print":
            self.print_convo(pargs.n, pargs.mode)
        elif pargs.cmd == "save":
            self.save_convo(pargs.path, pargs.name, pargs.type)

    def __del__(self):
        if self.alias and self.alias in XSH.aliases:
            print(f"Deleting {self.alias}")
            del XSH.aliases[self.alias]

        if self._managed:
            XSH.builtins.events.on_chat_destroy.fire(inst=self)

    def __str__(self):
        stats = self._stats()
        return "\n".join([f"{k} {v}" for k, v, _, _ in stats])

    def __repr__(self):
        return f"ChatGPT(alias={self.alias or None})"

    @property
    def tokens(self) -> int:
        """Current convo tokens"""
        return self._base_tokens + sum(self._tokens[self.chat_idx:])

    @property
    def base(self) -> list[dict[str, str]]:
        return self._base

    @base.setter
    def base(self, msgs: list[dict[str, str]]) -> None:
        self._base_tokens = sum(get_token_list(msgs))
        self._base = msgs
    
    @property
    def chat_convo(self) -> list[dict[str, str]]:
        return self.base + self.messages[self.chat_idx:]

    def stats(self) -> None:
        """Prints conversation stats to shell"""
        stats = self._stats()
        return "\n".join(
            [
                ansi_partial_color_format(
                    "{} {}{}{} {}".format(tok, col, key, "{BOLD_WHITE}", val)
                )
                for key, val, col, tok in stats
            ]
        )

    def _stats(self) -> str:
        """Helper for stats and __str__"""
        stats = [
            ("Alias:", f"{self.alias or None}", "{BOLD_GREEN}", "🤖"),
            ("Tokens:", self.tokens, "{BOLD_BLUE}", "🪙"),
            ("Trim After:", f"{self._max_tokens} Tokens", "{BOLD_BLUE}", "🔪"),
            ("Messages:", len(self.messages), "{BOLD_BLUE}", "📨"),
        ]
        return stats

    def chat(self, text: str) -> str:
        """
        Main chat function for interfacing with OpenAI API

        This method is not meant to be called directly as calling the
            instance will call this method and print to the cli.
        However, it is avaliable and can be useful for assigning
            the response to a variable.

        Parameters
        ----------
        text : str
            Text to send to ChatGPT

        Returns
        -------
        str: Response from ChatGPT

        """

        model = XSH.env.get("OPENAI_CHAT_MODEL", "gpt-3.5-turbo")
        choices = ["gpt-3.5-turbo", "gpt-4"]
        if model not in choices:
            raise UnsupportedModelError(
                f"Unsupported model: {model} - options are {choices}"
            )

        if not openai.api_key:
            api_key = XSH.env.get("OPENAI_API_KEY", None)
            if not api_key:
                raise NoApiKeyError()
            openai.api_key = api_key

        self.messages.append({"role": "user", "content": text})
        self.chat_idx -= 1

        try:
            response = openai.ChatCompletion.create(
                model=model,
                messages=self.chat_convo,
            )
        except OpenAIError as e:
            self.messages.pop()
            self.chat_idx += 1
            sys.exit(
                ansi_partial_color_format(
                    "{}OpenAI Error{}: {}".format("{BOLD_RED}", "{RESET}", e)
                )
            )

        res_text = response["choices"][0]["message"]
        user_toks, gpt_toks = (
            response["usage"]["prompt_tokens"],
            response["usage"]["completion_tokens"],
        )

        self.messages.append(res_text)
        self._tokens.extend([user_toks, gpt_toks])

        self.chat_idx -= 1
        self.trim_convo()

        return res_text["content"]
    
    def trim_convo(self) -> None:
        while self.chat_idx < -1 and self.tokens > self._max_tokens:
            self.chat_idx += 1

    def _get_json_convo(self, n: int) -> list[dict[str, str]]:
        """Returns the current conversation as a JSON string, up to n last items"""
        n = -n if n != 0 else 0
        messages = self.base + self.messages if n == 0 else self.messages[n:]

        return json.dumps(messages, indent=4)

    def _get_printed_convo(self, n: int, color: bool = True) -> list[tuple[str, str]]:
        """Helper method to get up to n items of conversation, formatted for printing"""
        user = XSH.env.get("USER", "user")
        convo = []
        n = -n if n != 0 else 0
        messages = self.base + self.messages if n == 0 else self.messages[n:]

        for msg in messages:
            if msg["role"] == "user":
                role = (
                    ansi_partial_color_format("{BOLD_GREEN}" + f"{user}:" + "{RESET}")
                    if color
                    else f"{user}:"
                )
            elif msg["role"] == "assistant":
                role = (
                    ansi_partial_color_format("{BOLD_BLUE}" + "ChatGPT:" + "{RESET}")
                    if color
                    else "ChatGPT:"
                )
            elif msg["role"] == "system":
                role = (
                    ansi_partial_color_format("{BOLD_BLUE}" + "System:" + "{RESET}")
                    if color
                    else "System:"
                )
            else:
                continue

            if color:
                content = format_markdown(msg["content"])
            else:
                content = msg["content"]

            convo.append((role, indent(content)))

        return convo

    def print_convo(self, n: int = 10, mode: str = "color") -> str:
        """Prints the current conversation to shell, up to n last items, in the specified mode

        Parameters
        ----------
        n : int, optional
            Number of messages to print. Defaults to 10.
            If n is 0, the entire conversation will be printed.
        mode : str, optional
            Mode to print the conversation in. Defaults to 'color'.
            Options - 'color', 'no-color', 'json'

        Examples
        --------
            >>> chatgpt.print_convo(5, 'color') # prints the last 5 items of the conversation, with color
            >>> chatgpt.print_convo(0, 'json') # prints the entire conversation as a JSON string
            >>> chatgpt.print_convo(mode='no-color') # prints the last 10 items of the conversation,
                    without color or pygments markdown formatting
        """

        rtn_str = "\n"

        if not self.messages:
            raise NoConversationsError()

        if mode in ["color", "no-color"]:
            convo = self._get_printed_convo(n, mode == "color")
            for role, content in convo:
                rtn_str += role + "\n"
                rtn_str += content + "\n"
        elif mode == "json":
            convo = self._get_json_convo(n)
            rtn_str += convo
        else:
            raise InvalidConversationsTypeError(
                f'Invalid mode: "{mode}" -- options are "color", "no-color", and "json"'
            )
        print(rtn_str)

    def save_convo(
        self, path: str = "", name: str = "", mode: str = "text", override: bool = False
    ) -> None:
        """
        Saves conversation to path or default xonsh data directory

        Parameters
        ----------
        path : str, optional
            File path to save the conversation. Defaults to ''.
            If path is not specified, the conversation will be saved to the
            default xonsh data directory.
        name : str, optional
            Name of the conversation file. Defaults to ''.
            Ignored when path is specified.
        mode : str, optional
            Type of the conversation file. Defaults to 'text'.
            Options - 'text', 'json'
        override : bool, optional
            Whether or not to override existing files. Defaults to False.

        Returns
        -------
        None

        Notes
        -----
        Default File Path Structure:
            $XONSH_DATA_DIR/chatgpt/$USER_[name|alias|'chatgpt']_[date]_[index].[text|json]
        """
        if not self.messages:
            raise NoConversationsError()

        if not path:
            path = get_default_path(
                name=name, json_mode=mode == "json", override=override, alias=self.alias
            )
        elif os.path.exists(path) and not override:
            res = input(f"File already exists: {path}\nOverride? [Y/n]: ")

            if res.lower() == "n":
                return

        if mode == "text":
            convo = self._get_printed_convo(0, color=False)
        elif mode == "json":
            convo = self._get_json_convo(0)
        else:
            raise InvalidConversationsTypeError(
                f'Invalid mode: "{mode}" -- options are "text", and "json"'
            )

        with open(path, "w") as f:
            if mode == "json":
                f.write(convo)
            else:
                for role, content in convo:
                    f.write(role + "\n")
                    f.write(content + "\n")

        print("Conversation saved to: " + str(path))
        return

    @staticmethod
    def fromcli(args: list[str], stdin: TextIO = None) -> None:
        """Helper method for one off conversations from the shell.
        Conversation will not save to instance, and will not be saved to history.
        Not meant to be called directly, but from a xonsh alias.
        Alias 'chatgpt' is automatically registered when the xontrib is loaded.

        Parameters
        ----------
        args : list[str]
            Arguments passed from the shell
        stdin : TextIO, optional
            Text from stdin, by default None

        Usage
        -----
            >>> chatgpt [text] # text will be sent to ChatGPT
            >>> echo [text] | chatgpt # text from stdin will be sent to ChatGPT
            >>> cat [text file] | chatgpt # text from file will be sent to ChatGPT
        """
        inst = ChatGPT()
        inst(args, stdin)
        return None

    @staticmethod
    def getdoc() -> str:
        return DOCSTRING

    @classmethod
    def fromconvo(cls, path, alias: str = "", managed: bool = False) -> "ChatGPT":
        """Loads a conversation from a saved file and returns new instance

        Parameters
        ----------
        path : str
            Path to the conversation file
        alias : str, optional
            Alias to use for the instance. Defaults to ''.
        managed : bool, optional
            Whether or not to register the instance with the chat manager.
            Defaults to False.

        Returns
        -------
        ChatGPT
            New instance with the loaded conversation
        """
        if not os.path.exists(path):
            bname = os.path.basename(path)
            default_dir = XSH.env.get(
                "XONSH_DATA_DIR",
                os.path.join(os.path.expanduser("~"), ".local", "share", "xonsh"),
            )
            guess_name = os.path.join(default_dir, "chatgpt", bname)

            if not os.path.exists(guess_name):
                raise FileNotFoundError(f"File not found: {path}")
            else:
                path = guess_name

        with open(path, "r") as f:
            convo = f.read()

        messages, base = parse_convo(convo)
        new_cls = cls(alias=alias, managed=managed)
        new_cls.messages = messages
        if base:
            new_cls.base = base
            new_cls._base_tokens = sum(get_token_list(base))
        new_cls._tokens = get_token_list(messages)
        new_cls.chat_idx = -len(messages)
        new_cls.trim_convo()

        return new_cls

xontrib_chatgpt/chatgpt.py
jpal91/xontrib-chatgpt
"""Module for ChatManager to manage multiple chats"""
import os
import sys
import weakref
from collections import defaultdict
from typing import Optional, Union, TextIO
from argparse import ArgumentParser
from re import Pattern

from xonsh.built_ins import XSH
from xonsh.ansi_colors import ansi_partial_color_format
from xonsh.lazyasd import LazyObject

from xontrib_chatgpt.chatgpt import ChatGPT
from xontrib_chatgpt.utils import convert_to_sys
from xontrib_chatgpt.lazyobjs import _FIND_NAME_REGEX
from xontrib_chatgpt.args import _cm_parse
from xontrib_chatgpt.exceptions import (
    NoConversationsError,
    InvalidConversationsTypeError,
)

FIND_NAME_REGEX: Pattern = LazyObject(_FIND_NAME_REGEX, globals(), "FIND_NAME_REGEX")
PARSER: ArgumentParser = LazyObject(_cm_parse, globals(), "PARSER")

TUTORIAL = """
        {BOLD_WHITE}Usage of {BOLD_BLUE}chat-manager{RESET}
        {BOLD_WHITE}---------------------{RESET}

        First, start off by creating a new chat that we'll call 'gpt':
            {YELLOW}>>> {INTENSE_BLUE}chat-manager {INTENSE_GREEN}add {RESET}gpt
        
        This creates a new conversation with {INTENSE_BLUE}ChatGPT{RESET} which you can interact with.
        Any input and responses will be saved to this instance, so you can continue
        with your conversation as you please.

        From here, you have several ways to interact. The instance is created as 
        a {INTENSE_BLUE}Xonsh{RESET} alias, so you can simply call it as such:
            {YELLOW}>>> {INTENSE_BLUE}gpt{RESET} Hello, how are you?

        Since {INTENSE_BLUE}Xonsh{RESET} aliases can can interact with bash and xsh syntax, you can also
        use the following:
            {YELLOW}>>> {INTENSE_BLUE}echo{RESET} {INTENSE_YELLOW}'Hello, how are you?' {INTENSE_WHITE}| {INTENSE_BLUE}gpt{RESET}
            {YELLOW}>>> {INTENSE_BLUE}gpt{RESET} {INTENSE_WHITE}< {INTENSE_BLUE}echo{RESET} {INTENSE_YELLOW}'Hello, how are you?'
            {YELLOW}>>> {INTENSE_BLUE}cat{RESET} my_input.txt {INTENSE_WHITE}| {INTENSE_BLUE}gpt{RESET}
            {YELLOW}>>> {INTENSE_WHITE}my_input ={RESET} {INTENSE_YELLOW}"Hello, how are you?"
            {YELLOW}>>> {INTENSE_BLUE}echo {INTENSE_PURPLE}@({INTENSE_WHITE}my_input{INTENSE_PURPLE}) {INTENSE_WHITE}| {INTENSE_BLUE}gpt{RESET}
        
        Finally, the instance also acts as a {INTENSE_BLUE}Xonsh{RESET} context block:
            {YELLOW}>>> {INTENSE_GREEN}with{INTENSE_WHITE}! {INTENSE_BLUE}gpt{INTENSE_WHITE}:{RESET}
            {YELLOW}>>>{RESET}     Can you help me fix my python function?
            {YELLOW}>>>{RESET}     def hello_world():
            {YELLOW}>>>{RESET}         return
            {YELLOW}>>>{RESET}         print("Hello World!")
        
        Any content added to the context block will be sent to {INTENSE_BLUE}ChatGPT{RESET}, allowing
        you to send multi-line messages to {INTENSE_BLUE}ChatGPT{RESET}.

        Each conversation has a separate history, so if you want to start a new conversation on
        a different topic, simply create a new instance:
            {YELLOW}>>> {INTENSE_BLUE}chat-manager {INTENSE_GREEN}add {RESET}my_literature_convo

        {BOLD_WHITE}Other Useful Commands{RESET}
        {BOLD_WHITE}---------------------{RESET}

        Each conversation instance has it's own individual commands as well. Use the following
        to learn more about the options available to you:
            {YELLOW}>>> {INTENSE_BLUE}gpt {INTENSE_GREEN}-h{RESET}
        
        Print a conversation:
            {YELLOW}>>> {INTENSE_BLUE}gpt {INTENSE_GREEN}-p{RESET}
            {GREEN}# or{RESET}
            {YELLOW}>>> {INTENSE_BLUE}chat-manager {INTENSE_GREEN}print{RESET} gpt
        
        Save a conversation:
            {YELLOW}>>> {INTENSE_BLUE}my_literature_convo {INTENSE_GREEN}-s{RESET}
            {GREEN}# or{RESET}
            {YELLOW}>>> {INTENSE_BLUE}chat-manager {INTENSE_GREEN}save{RESET} my_literature_convo
        
        You can delete a conversation using {INTENSE_GREEN}Python{RESET} syntax:
            {YELLOW}>>> {INTENSE_PURPLE}del {INTENSE_BLUE}gpt{RESET}
        
        This will delete the instance, unsaved conversation history, and alias. Be careful!

        {BOLD_WHITE}Editing System Instructions{RESET}
        {BOLD_WHITE}---------------------{RESET}

        You can edit the {INTENSE_BLUE}ChatGPT{RESET} system instructions that are sent to the instance
        by using {INTESNE_BLUE}chat-manager edit -s{RESET}. This allows you to customize the
        instructions sent to {INTENSE_BLUE}ChatGPT{RESET} to your liking.

        The system messages must be either a list of dicts, a dict, or a yaml string (if PyYaml is installed).

        For more information, please visit: https://github.com/jpal91/xontrib-chatgpt/edit_sys_messages.md
"""


class ChatManager:
    """Class to manage multiple chats"""

    def __init__(self):
        self._instances: dict[int, dict[str, Union[str, ChatGPT]]] = defaultdict(dict)
        self._current: Optional[int] = None
        self._update_inst_dict()

    def __call__(self, args: list[str], stdin: TextIO = None):
        """Main method to interact with ChatManager via xonsh aliases"""
        if args:
            pargs = PARSER.parse_args(args)
        elif stdin:
            pargs = PARSER.parse_args(stdin.read().strip().split())
        else:
            return PARSER.print_help()

        if pargs.C:
            if self._current is None:
                return "No active chat!"
            else:
                return self._instances[self._current]["inst"].stats()

        if pargs.cmd in ["add", "a", "create"]:
            return self.add(pargs.name[0])
        elif pargs.cmd in ["list", "ls"]:
            return self.ls(saved=pargs.saved)
        elif pargs.cmd == "load":
            return self.load(pargs.name[0])
        elif pargs.cmd == "save":
            return self.save(
                chat_name=pargs.name, mode=pargs.mode, override=pargs.override
            )
        elif pargs.cmd in ["print", "p"]:
            return self.print_chat(chat_name=pargs.name, n=pargs.n, mode=pargs.mode)
        elif pargs.cmd == "help":
            return self.help(tgt=pargs.target)
        elif pargs.cmd in ["edit", "e"]:
            return self.edit(
                chat_name=pargs.name, sys_msgs=pargs.sys_msgs, no_code=pargs.no_code
            )
        else:
            return PARSER.print_help()

    def add(self, chat_name: str) -> str:
        """Create new chat instance

        Parameters
        ----------
        chat_name : str
            Name of the chat instance to create

        Returns
        -------
        str
        """
        if chat_name in self.chat_names():
            return "Chat with that name already exists!"
        elif chat_name in XSH.ctx or chat_name in XSH.aliases:
            return "Variable with that name already exists!"
        inst = ChatGPT(alias=chat_name, managed=True)
        XSH.ctx[chat_name] = inst
        self._instances[hash(inst)]["name"] = chat_name
        return f"Created new chat '{chat_name}'"

    def ls(self, saved: bool = False) -> str:
        """List all active chats or saved chats

        Parameters
        ----------
        saved : bool, optional
            Whether to list saved chats or active chats, by default False

        Returns
        -------
        str
        """
        if saved:
            return (
                ansi_partial_color_format("{BOLD_WHITE}Saved chats:")
                + "\n  "
                + "\n  ".join(self._find_saved())
            )

        if not self._instances:
            return "No active chats."

        return "\n\n".join([inst["inst"].stats() for inst in self._instances.values()])

    def load(self, path_or_name: str) -> str:
        """Load a conversation from a path or a saved chat name

        Parameters
        ----------
        path_or_name : str
            Path to a conversation file or name of a saved chat

        Returns
        -------
        str
        """
        if os.path.exists(path_or_name):
            path, name = path_or_name, FIND_NAME_REGEX.sub(
                r"\1", os.path.basename(path_or_name)
            )
        else:
            path, name = self._find_path_from_name(path_or_name)

        curr_names = self.chat_names()

        if name in curr_names:
            idx = 0
            while name + str(idx) in curr_names:
                idx += 1
            print(
                f"Chat with name {name} already loaded. Updating to {name + str(idx)}"
            )
            name += str(idx)

        inst = ChatGPT.fromconvo(path, alias=name, managed=True)
        XSH.ctx[name] = inst
        self._instances[hash(inst)]["name"] = name

        return f"Loaded chat {name} from {path}"

    def save(
        self, chat_name: str = "", mode: str = "text", override: bool = False
    ) -> Optional[str]:
        """Save a conversation to a file

        Parameters
        ----------
        chat_name : str, optional
            Name of the chat to save. Defaults to last used/current chat, by default ''
        mode : str, optional
            Mode to save the conversation. Defaults to 'text', other option is 'json'

        Returns
        -------
        Optional[str]
        """

        chat = self.get_chat_by_name(chat_name)

        try:
            res = chat["inst"].save_convo(mode=mode, override=override)
        except (NoConversationsError, InvalidConversationsTypeError) as e:
            print(e)
            sys.exit(1)

        return res

    def print_chat(
        self, chat_name: str = "", n: int = 10, mode: str = "color"
    ) -> Optional[str]:
        """Print a conversation

        Parameters
        ----------
        chat_name : str, optional
            Name of the chat to print. Defaults to last used/current chat, by default ''
        n : int, optional
            Number of conversations to print, by default 10
        mode : str, optional
            Mode to print the conversation. Defaults to 'color', other options are 'no-color' and 'json'

        Returns
        -------
        Optional[str]
        """

        chat = self.get_chat_by_name(chat_name)

        try:
            res = chat["inst"].print_convo(n=n, mode=mode)
        except (NoConversationsError, InvalidConversationsTypeError) as e:
            print(e)
            sys.exit(1)

        return res

    def help(self, tgt: str = "") -> Optional[str]:
        """Print help information on the instance, methods, or xontrib

        Parameters
        ----------
        tgt : str, optional
            Target to print help information on. Defaults to chat-manager, by default ''

        Returns
        -------
        Optional[str]
        """
        if not tgt:
            return self.tutorial()

        # TODO: Expand on this more
        if hasattr(self, tgt):
            XSH.help(getattr(self, tgt))
        elif hasattr(ChatGPT, tgt):
            XSH.help(getattr(ChatGPT, tgt))
        else:
            PARSER.print_help()

    def edit(
        self, chat_name: str = "", sys_msgs: str = "", no_code: bool = False
    ) -> Optional[str]:
        """Allows editing attributes of a chat instance

        Parameters
        ----------
        chat_name : str, optional
            Name of the chat to edit. Defaults to last used/current chat, by default ''
        sys_msgs : str, optional
            System messages to edit, by default ''
            Is a string representation of either a python list[dict], dict, or yaml equivalent.

        Returns
        -------
        Optional[str]

        See Also
        --------
        xontrib_chatgpt.chatmanager.convert_to_sys
        """
        chat = self.get_chat_by_name(chat_name)

        # TODO: Remove when more options are added
        if not sys_msgs:
            return "No system messages to edit!"

        sys_msgs = convert_to_sys(sys_msgs)
        if no_code:
            chat["inst"].base = sys_msgs
        else:
            chat["inst"].base = [chat["inst"].base[1]] + sys_msgs

    def chat_names(self) -> list[str]:
        """Returns chat names for current conversations"""
        return [inst["name"] for inst in self._instances.values()]

    def get_chat_by_name(self, chat_name: str) -> Optional[dict]:
        """Returns a chat that matches given name or exits"""
        if not chat_name and not self._current:
            print("No active chat!")
            sys.exit(1)
        elif not chat_name:
            chat = self._instances[self._current]
        else:
            try:
                chat = self._instances[hash(XSH.ctx[chat_name])]
            except KeyError:
                print(f"No chat with name {chat_name} found.")
                sys.exit(1)

        return chat

    def _find_saved(self) -> list[Optional[str]]:
        """Returns a list of saved chat files in the default directory"""
        def_dir = os.path.join(XSH.env["XONSH_DATA_DIR"], "chatgpt")
        if not os.path.exists(def_dir):
            return []
        return [
            f for f in os.listdir(def_dir) if os.path.isfile(os.path.join(def_dir, f))
        ]

    def _find_path_from_name(self, name: str) -> tuple[str, str]:
        """Finds a saved chat file from user input if it's not a path"""
        def_dir = XSH.env["XONSH_DATA_DIR"]
        saved = self._find_saved()

        if name in saved:
            return os.path.join(def_dir, "chatgpt", name), FIND_NAME_REGEX.sub(
                r"\1", name
            )

        # Creates a tuple of (chat_name, file_name)
        names_saved = [
            (n, f)
            for n, f in [(FIND_NAME_REGEX.sub(r"\1", f), f) for f in saved]
            if n == name
        ]

        if not names_saved:
            raise FileNotFoundError(f"No chat with name {name} found.")
        elif len(names_saved) > 1:
            chat_name, file = self._choose_from_multiple(names_saved)
        else:
            chat_name, file = names_saved[0]

        return os.path.join(def_dir, "chatgpt", file), chat_name

    def _choose_from_multiple(self, chats: list[tuple[str, str]]) -> tuple[str, str]:
        """If multiple saved chats are found, allows the user to choose from them"""

        print("Multiple choices found, please choose from:")
        [print(f"  {i + 1}. {os.path.basename(n[1])}") for i, n in enumerate(chats)]

        while True:
            max_choice = len(chats)
            choice = input(f"Number (1 - {max_choice}): ")

            try:
                choice = int(choice)
            except ValueError:
                print("Invalid input, please enter a number.")
                continue

            if not 0 < choice <= max_choice:
                print(
                    f"Invalid input, please enter a number between 1 and {max_choice}."
                )
                continue

            return chats[choice - 1]

    def _update_inst_dict(self):
        """Finds all active instances (active convos) in the global space and updates the internal dict"""
        for key, inst in XSH.ctx.copy().items():
            if isinstance(inst, ChatGPT):
                self._instances[hash(inst)] = {
                    "name": key,
                    "alias": inst.alias,
                    "inst": weakref.proxy(inst),
                }

        if self._current not in self._instances:
            self._current = None

    def on_chat_create_handler(self, inst: ChatGPT) -> None:
        """Handler for on_chat_create. Updates the internal dict with the new chat instance."""
        self._current = hash(inst)

        # 'name' left blank because at this point the instance init isn't complete,
        # so variable name is not yet in the global space.
        # Updated later in the add method
        self._instances[hash(inst)] = {
            "name": "",
            "alias": inst.alias,
            "inst": weakref.proxy(inst),
        }

    def on_chat_destroy_handler(self, inst: ChatGPT) -> None:
        """Handler for on_chat_destroy. Removes the chat instance from the internal dict."""
        inst_hash = hash(inst)

        if inst_hash in self._instances:
            del self._instances[inst_hash]

        if inst_hash == self._current:
            self._current = None

    def on_chat_used_handler(self, inst: ChatGPT) -> None:
        """Handler for on_chat_used. Updates the current chat instance."""
        self._current = hash(inst)

    def tutorial(self) -> str:
        """Returns a usage string for the xontrib."""
        return ansi_partial_color_format(TUTORIAL)


# TODO: Print from a saved file
# TODO:
#   Save/Load name conflicts
#   Possibly add a flag to overwrite, or option to replace the old one
#   or Load chat with similar saved name into different variable name
#   ie load from 'path' into 'name'

xontrib_chatgpt/chatmanager.py
jpal91/xontrib-chatgpt
from xonsh.built_ins import XSH
from xonsh.completers.completer import add_one_completer, remove_completer
from xonsh.completers.tools import RichCompletion, contextual_command_completer_for, contextual_completer
from xonsh.parsers.completion_context import CommandContext, CompletionContext

from xontrib_chatgpt.chatmanager import ChatManager 


@contextual_command_completer_for("chat-manager")
def cm_completer(command: CommandContext) -> set[RichCompletion]:
    """Completions for chat-manager"""
    opts_desc = {
        "add": "Add/Create a new chat",
        "list": "List all current or saved chats",
        "save": "Save a chat to a local file",
        "load": "Load a chat from a local file",
        "print": "Print a chat to the console",
    }
    if command.arg_index < 2:
        return {
                RichCompletion(k, description=v, append_space=True)
                for k, v in opts_desc.items()
                if k.startswith(command.prefix)
            }

@contextual_completer
def print_save_chat_completer(command: CompletionContext) -> set[str]:
    """Completions for chat-manager print"""

    if (
        command.command and
        command.command.arg_index == 2 and
        command.command.args[0].value == 'chat-manager' and
        command.command.args[1].value in ['print', 'save']
    ):
        cm: ChatManager = XSH.ctx['chat_manager']

        return {
                *cm.chat_names()
            }



@contextual_completer
def load_chat_completer(context: CompletionContext) -> set[str]:
    """Completions for chat-manager load"""
    
    if (
        context.command and
        context.command.arg_index == 2 and
        context.command.args[0].value == 'chat-manager' and
        context.command.args[1].value == 'load'
    ):
        cm: ChatManager = XSH.ctx['chat_manager']

        return {
                *cm._find_saved()
            }

def add_completers() -> None:
    add_one_completer("chat-manager", cm_completer, loc="start")
    add_one_completer('cm-load', load_chat_completer, loc='start')
    add_one_completer('cm-print', print_save_chat_completer, loc='start')

def rm_completers() -> None:
    remove_completer("chat-manager")
    remove_completer('cm-load')
    remove_completer('cm-print')
xontrib_chatgpt/completers.py
jpal91/xontrib-chatgpt
"""Events for xontrib_chatgpt."""
from xonsh.built_ins import XonshSession
from xontrib_chatgpt.chatmanager import ChatManager

chat_events = [
    (
        "on_chat_create",
        """
        on_chat_create(inst: ChatGPT) -> None

        Fired when new chat is created. Passes the new instance to be added to
        the chat manager.
        """,
    ),
    (
        "on_chat_destroy",
        """
        on_chat_destroy(inst: ChatGPT) -> None

        Fired when chat is destroyed. Passes the instance to be removed from
        the chat manager.
        """,
    ),
    (
        "on_chat_used",
        """
        on_chat_used(inst: ChatGPT) -> None

        Fires when chat is used, i.e. chatgpt is called. Passes the instance
        of the chat instance to update the current chat for the manager.
        """,
    ),
]


def add_events(xsh: XonshSession, cm: ChatManager):
    events = xsh.builtins.events

    for c in chat_events:
        events.doc(*c)

    events.on_chat_create(lambda *_, **kw: cm.on_chat_create_handler(**kw))
    events.on_chat_destroy(lambda *_, **kw: cm.on_chat_destroy_handler(**kw))
    events.on_chat_used(lambda *_, **kw: cm.on_chat_used_handler(**kw))


def rm_events(xsh: XonshSession):
    events = xsh.builtins.events

    for c in chat_events:
        delattr(events, c[0])

xontrib_chatgpt/events.py
jpal91/xontrib-chatgpt
"""Exceptions for the ChatGPT class"""

#############
# Exceptions
#############
class NoApiKeyError(Exception):
    """Raised when no API key is found"""

    def __init__(self, *_):
        pass

    def __str__(self):
        return "\n\x1b[1;31mNo OpenAI API key found"


class UnsupportedModelError(Exception):
    """Raised when an unsupported model is used"""

    def __init__(self, msg: str, *_):
        self.msg = msg

    def __str__(self):
        return f"\n\x1b[1;31m{self.msg}"


class NoConversationsError(Exception):
    """Raised when attempting to print/save a conversation when there is no conversation history"""

    def __init__(self, *_):
        pass

    def __str__(self):
        return "\n\x1b[1;31mNo conversation history!"


class InvalidConversationsTypeError(Exception):
    """Raised when attempting to print/save a conversation with an invalid type"""

    def __init__(self, msg: str, *_):
        self.msg = msg

    def __str__(self):
        return f"\n\x1b[1;31m{self.msg}"

class MalformedSysMsgError(Exception):
    """Raised when user attempts to set a malformed system message for a chat"""

    def __init__(self, msg: str, *_):
        self.msg = msg

    def __str__(self):
        return f"\x1b[1;31mIncorrect System Message Format:\n{self.msg}\nMust be a python list[dict], dict, or yaml equivalent. See documentation for more information."


#############
xontrib_chatgpt/exceptions.py
jpal91/xontrib-chatgpt
"""Lazy objects for xontrib_chatgpt"""
import re
from collections import namedtuple

#############
# Lazy Objects
#############


def _openai():
    """Imports openai"""
    import openai

    return openai


def _tiktoken():
    """Imports tiktoken"""
    import tiktoken

    return tiktoken.get_encoding("cl100k_base")


def _MULTI_LINE_CODE():
    """Regex to remove multiline code blocks (```code```) from markdown"""
    return re.compile(r"```.*?\n", re.DOTALL)


def _SINGLE_LINE_CODE():
    """Regex to remove single line code blocks (`code`) from markdown"""
    return re.compile(r"`(.*?)`")


def _PYGMENTS():
    """Lazy loading of pygments to avoid slowing down shell startup"""
    from pygments import highlight
    from pygments.lexers import get_lexer_by_name
    from pygments.lexers.python import PythonLexer
    from pygments.formatters import Terminal256Formatter
    from pygments.styles.gh_dark import GhDarkStyle

    container = namedtuple(
        "container",
        [
            "highlight",
            "get_lexer_by_name",
            "PythonLexer",
            "Terminal256Formatter",
            "GhDarkStyle",
        ],
    )
    return container(
        highlight, get_lexer_by_name, PythonLexer, Terminal256Formatter, GhDarkStyle
    )


def _markdown():
    """Formats markdown text using pygments"""
    from pygments import highlight
    from pygments.lexers.markup import MarkdownLexer
    from pygments.formatters import Terminal256Formatter
    from pygments.styles.gh_dark import GhDarkStyle

    return lambda text: highlight(
        text, MarkdownLexer(), Terminal256Formatter(style=GhDarkStyle)
    )


def _FIND_NAME_REGEX():
    """Regex to find the name of a chat from a file name"""
    return re.compile(r"^(?:.+?_)*([a-zA-Z0-9]+)(?:_[0-9\-]+)?\.(?:.*)$", re.DOTALL)

def _YAML():
    """Imports yaml package if it exists"""
    from importlib.util import find_spec

    if find_spec("yaml") is not None:
        import yaml
        return yaml
    else:
        return None
xontrib_chatgpt/lazyobjs.py
jpal91/xontrib-chatgpt
"""Utility Functions for xontrib-chatgpt"""
import os
import json
from typing import Union
from datetime import datetime
from textwrap import dedent

from xonsh.built_ins import XSH
from xonsh.tools import indent
from xonsh.ansi_colors import ansi_partial_color_format
from xonsh.lazyasd import LazyObject

from xontrib_chatgpt.lazyobjs import (
    _tiktoken,
    _MULTI_LINE_CODE,
    _SINGLE_LINE_CODE,
    _markdown,
    _YAML,
)
from xontrib_chatgpt.exceptions import MalformedSysMsgError


tiktoken = LazyObject(_tiktoken, globals(), "tiktoken")
MULTI_LINE_CODE = LazyObject(_MULTI_LINE_CODE, globals(), "MULTI_LINE_CODE")
SINGLE_LINE_CODE = LazyObject(_SINGLE_LINE_CODE, globals(), "SINGLE_LINE_CODE")
markdown = LazyObject(_markdown, globals(), "markdown")
YAML = LazyObject(_YAML, globals(), "YAML")


def parse_convo(convo: str) -> tuple[list[dict[str, str]]]:
    """Parses a conversation from a saved file

    Parameters
    ----------
    convo : str
        Conversation to parse

    Returns
    -------
    tuple[list[dict[str, str]]]
        Parsed conversation and base system messages
    """
    try:
        convo = json.loads(convo)
    except json.JSONDecodeError:
        pass
    else:
        base, messages = [], []
        for msg in convo:
            if msg["role"] == "system":
                base.append(msg)
            else:
                messages.append(msg)
        return messages, base

    convo = convo.split("\n")
    messages, base, user, idx, n = [], [], True, 0, len(convo)

    # Couldn't figure out how to do this with regex, so while loop instead
    while idx < n:
        if not convo[idx]:
            idx += 1
            continue

        if convo[idx].startswith("System"):
            msg = {"role": "system", "content": ""}
        else:
            msg = {"role": "user" if user else "assistant", "content": ""}

        idx += 1

        if idx >= n:
            break

        while idx < n and convo[idx].startswith(" "):
            msg["content"] += convo[idx] + "\n"
            idx += 1

        msg["content"] = dedent(msg["content"])

        if msg["role"] == "system":
            base.append(msg)
        else:
            messages.append(msg)
            user = not user

    return messages, base


def get_token_list(messages: list[dict[str, str]]) -> list[int]:
    """Gets the chat tokens for the loaded conversation

    Parameters
    ----------
    messages : list[dict[str, str]]
        Messages from the conversation

    Returns
    -------
    list[int]
        List of tokens for each message
    """
    tokens_per_message = 3
    tokens = [3]

    for message in messages:
        num_tokens = 0
        num_tokens += tokens_per_message
        for v in message.values():
            num_tokens += len(tiktoken.encode(v))
        tokens.append(num_tokens)

    return tokens


def print_res(res: str) -> None:
    """Called after receiving response from ChatGPT, prints the response to the shell"""
    res = format_markdown(res)
    res = indent(res)
    print(ansi_partial_color_format("\n{BOLD_BLUE}ChatGPT:{RESET}\n"))
    print(res)


def format_markdown(text: str) -> str:
    """Formats the text using the Pygments Markdown Lexer, removes markdown code '`'s"""
    text = markdown(text)
    text = MULTI_LINE_CODE.sub("", text)
    text = SINGLE_LINE_CODE.sub(r"\1", text)
    return text


def get_default_path(
    name: str = "", json_mode: bool = False, override: bool = False, alias: str = ""
) -> str:
    """Helper method to get the default path for saving conversations"""
    user = XSH.env.get("USER", "user")
    data_dir = XSH.env.get(
        "XONSH_DATA_DIR",
        os.path.join(os.path.expanduser("~"), ".local", "share", "xonsh"),
    )
    chat_dir = os.path.join(data_dir, "chatgpt")

    if not os.path.exists(chat_dir):
        os.makedirs(chat_dir)

    date, idx = datetime.now().strftime("%Y-%m-%d"), 1
    path_prefix, ext = (
        os.path.join(chat_dir, f"{user}_{name or alias or 'chatgpt'}_{date}"),
        ".json" if json_mode else ".txt",
    )

    if os.path.exists(f"{path_prefix}{ext}") and not override:
        while os.path.exists(path_prefix + f"_{idx}{ext}"):
            idx += 1
        path = path_prefix + f"_{idx}{ext}"
    else:
        path = path_prefix + f"{ext}"

    return path

def convert_to_sys(msgs: Union[str, dict, list]) -> list[dict]:
    """Returns a list of dicts from a string of python dict, json, or yaml.
    Calls itself recursively until the result is either a list[dict]
    or raises an error.

    Will attempt to parse yaml only if package is installed.

    Parameters
    ----------
    msgs : Union[str, dict, list]
        Messages to convert to system messages

    Returns
    -------
    list[dict]

    Raises
    ------
    MalformedSysMsgError
        Raised when the messages are not properly formatted

    Examples
    --------
    msgs :
        '[{"role": "system", "content": "Hello"}, {"role": "system", "content": "Hi there!"}]'
        '{"role": "system", "content": "Hello"}'
        '- role: system\n  content: Hello\n- role: system\n  content: Hi there!'
    """

    if isinstance(msgs, str):
        msgs = msgs.strip()
        try:
            msgs = eval(msgs)
        except SyntaxError:
            pass

    if isinstance(msgs, dict):
        return convert_to_sys([msgs])
    elif isinstance(msgs, list):
        for m in msgs:
            if "content" not in m:
                raise MalformedSysMsgError(msgs)
            m["role"] = "system"
        return msgs
    elif YAML is not None:
        try:
            return YAML.safe_load(msgs)
        except YAML.YAMLError:
            pass

    raise MalformedSysMsgError(msgs)

xontrib_chatgpt/utils.py
jpal91/xontrib-chatgpt