diff --git a/poetry.lock b/poetry.lock index fdc8b3b..9de1f39 100644 --- a/poetry.lock +++ b/poetry.lock @@ -764,4 +764,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = ">=3.11,<3.12" -content-hash = "e8c5d78c5c95eaadb03e603c5b4ceada8aa27aaa049e6c0d72129c1f2dc53ed9" +content-hash = "8d76898eeb53fd3848f3be2f6aa1662517f9dbd80146db8dfd6f2932021ace48" diff --git a/src/llm_chat/bot.py b/src/llm_chat/bot.py new file mode 100644 index 0000000..6d9decd --- /dev/null +++ b/src/llm_chat/bot.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +from pydantic import BaseModel + +from llm_chat.models import Message, Role +from llm_chat.utils import kebab_case + + +def _bot_id_from_name(name: str) -> str: + """Create bot ID from name.""" + return kebab_case(name) + + +class BotConfig(BaseModel): + """Bot configuration class.""" + + bot_id: str + name: str + context_files: list[str] + + +class BotExists(Exception): + """Bot already exists error.""" + + pass + + +class BotDoesNotExists(Exception): + """Bot already exists error.""" + + pass + + +class Bot: + """Custom bot interface.""" + + def __init__(self, config: BotConfig, bot_dir: Path) -> None: + self.config = config + self.context: list[Message] = [] + for context_file in config.context_files: + path = bot_dir / "context" / context_file + if not path.exists(): + raise ValueError(f"{path} does not exist.") + if not path.is_file(): + raise ValueError(f"{path} is not a file") + with path.open("r") as f: + content = f.read() + self.context.append(Message(role=Role.SYSTEM, content=content)) + + @property + def id(self) -> str: + """Return the bot ID.""" + return self.config.bot_id + + @property + def name(self) -> str: + """Return the bot name.""" + return self.config.name + + @classmethod + def create(cls, name: str, bot_dir: Path, context_files: list[Path] = []) -> None: + """Create a custom bot. + + This command creates the directory structure for the custom bot and copies + the context files. The bot directory is stored within the base bot directory + (e.g. `~/.llm_chat/bots/`), which is stored as the snake case version of + the name argument. the directory contains a settings file `.json` and a + directory of context files. + + Args: + name: Name of the custom bot. + bot_dir: Path to where custom bot contexts are stored. + context_files: Paths to context files. + """ + bot_id = _bot_id_from_name(name) + path = bot_dir / bot_id + if path.exists(): + raise BotExists(f"The bot {name} already exists.") + (path / "context").mkdir(parents=True) + + config = BotConfig( + bot_id=bot_id, + name=name, + context_files=[context.name for context in context_files], + ) + + with (path / f"{bot_id}.json").open("w") as f: + f.write(config.model_dump_json() + "\n") + + for context in context_files: + shutil.copy(context, path / "context" / context.name) + + @classmethod + def load(cls, name: str, bot_dir: Path) -> Bot: + """Load existing bot.""" + bot_id = _bot_id_from_name(name) + bot_path = bot_dir / bot_id + if not bot_path.exists(): + raise BotDoesNotExists(f"Bot {name} does not exist.") + + with (bot_path / f"{bot_id}.json").open("r") as f: + config = BotConfig(**json.load(f)) + + return cls(config, bot_dir) diff --git a/src/llm_chat/cli/__init__.py b/src/llm_chat/cli/__init__.py new file mode 100644 index 0000000..ff0731d --- /dev/null +++ b/src/llm_chat/cli/__init__.py @@ -0,0 +1,3 @@ +from llm_chat.cli.main import app + +__all__ = ["app"] diff --git a/src/llm_chat/cli/bot.py b/src/llm_chat/cli/bot.py new file mode 100644 index 0000000..78cf6ed --- /dev/null +++ b/src/llm_chat/cli/bot.py @@ -0,0 +1,52 @@ +from pathlib import Path +from typing import Annotated, Any, Optional + +import typer + +from llm_chat.bot import Bot +from llm_chat.settings import OpenAISettings + +app = typer.Typer() + + +@app.command("create") +def create( + name: Annotated[ + str, + typer.Argument(help="Name of bot."), + ], + base_dir: Annotated[ + Optional[Path], + typer.Option( + ..., + "--base-dir", + "-d", + help=( + "Path to the base directory in which conversation " + "configuration and history will be saved." + ), + ), + ] = None, + context_files: Annotated[ + list[Path], + typer.Option( + ..., + "--context", + "-c", + help=( + "Path to a file containing context text. " + "Can provide multiple times for multiple files." + ), + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + ), + ] = [], +) -> None: + """Create a new bot.""" + args: dict[str, Any] = {} + if base_dir is not None: + args |= {"base_dir": base_dir} + settings = OpenAISettings(**args) + Bot.create(name, settings.bot_dir, context_files=context_files) diff --git a/src/llm_chat/cli.py b/src/llm_chat/cli/main.py similarity index 96% rename from src/llm_chat/cli.py rename to src/llm_chat/cli/main.py index 8bf4824..9f7b375 100644 --- a/src/llm_chat/cli.py +++ b/src/llm_chat/cli/main.py @@ -7,10 +7,12 @@ from rich.console import Console from rich.markdown import Markdown from llm_chat.chat import ChatProtocol, get_chat, get_chat_class +from llm_chat.cli import bot from llm_chat.models import Message, Role from llm_chat.settings import Model, OpenAISettings app = typer.Typer() +app.add_typer(bot.app, name="bot", help="Manage custom bots.") def prompt_continuation(width: int, *args: Any) -> str: @@ -117,7 +119,7 @@ def chat( "-c", help=( "Path to a file containing context text. " - "Can provide multiple time for multiple files." + "Can provide multiple times for multiple files." ), exists=True, file_okay=True, @@ -157,7 +159,7 @@ def chat( if temperature is not None: args |= {"temperature": temperature} if base_dir is not None: - args |= {"history_dir": base_dir} + args |= {"base_dir": base_dir} settings = OpenAISettings(**args) context_messages = [load_context(path) for path in context] diff --git a/src/llm_chat/settings.py b/src/llm_chat/settings.py index 0d2ce5b..f7027af 100644 --- a/src/llm_chat/settings.py +++ b/src/llm_chat/settings.py @@ -15,7 +15,7 @@ class Model(StrEnum): DEFAULT_MODEL = Model.GPT3 DEFAULT_TEMPERATURE = 0.7 -DEFAULT_BASE_DIR = Path.home() / ".llm_chat" +DEFAULT_BASE_DIR = Path.home() / ".llm-chat" DEFAULT_BOT_PATH = "bots" DEFAULT_HISTORY_PATH = "history" diff --git a/src/llm_chat/utils.py b/src/llm_chat/utils.py new file mode 100644 index 0000000..f7a7141 --- /dev/null +++ b/src/llm_chat/utils.py @@ -0,0 +1,9 @@ +import re + + +def kebab_case(string: str) -> str: + """Convert a string to kebab case.""" + string = string.replace("-", " ") + string = re.sub("([A-Z][a-z]+)", r" \1", string) + string = re.sub("([A-Z]+)", r" \1", string) + return "-".join(string.split()).lower() diff --git a/tests/test_cli.py b/tests/test_cli.py index c039be5..078f933 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -9,9 +9,9 @@ from pytest import MonkeyPatch from rich.console import Console from typer.testing import CliRunner -import llm_chat +import llm_chat.cli from llm_chat.chat import ChatProtocol -from llm_chat.cli import app +from llm_chat.cli.main import app from llm_chat.models import Conversation, Message, Role from llm_chat.settings import Model, OpenAISettings @@ -77,9 +77,9 @@ def test_chat(monkeypatch: MonkeyPatch) -> None: mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) - monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat) - monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) - monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) + monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) + monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) + monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) result = runner.invoke(app, ["chat"]) assert result.exit_code == 0 @@ -107,9 +107,9 @@ def test_chat_with_context( mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) - monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat) - monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) - monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) + monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) + monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) + monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) result = runner.invoke(app, ["chat", argument, str(tmp_file)]) assert result.exit_code == 0 @@ -139,9 +139,9 @@ def test_chat_with_name( mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) - monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat) - monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) - monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) + monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) + monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) + monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) result = runner.invoke(app, ["chat", argument, name]) assert result.exit_code == 0 @@ -179,9 +179,9 @@ def test_load(monkeypatch: MonkeyPatch, tmp_path: Path) -> None: mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) - monkeypatch.setattr(llm_chat.cli, "get_chat_class", mock_get_chat) - monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console) - monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input) + monkeypatch.setattr(llm_chat.cli.main, "get_chat_class", mock_get_chat) + monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) + monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) # Load the conversation from the file result = runner.invoke(app, ["load", str(file_path)]) diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..1bafc36 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,22 @@ +import pytest + +from llm_chat.utils import kebab_case + + +@pytest.mark.parametrize( + "string,expected", + [ + ("fooBar", "foo-bar"), + ("FooBar", "foo-bar"), + ("Foo Bar", "foo-bar"), + ("1Foo2Bar3", "1-foo2-bar3"), + ], + ids=[ + "fooBar", + "FooBar", + "Foo Bar", + "1Foo2Bar3", + ], +) +def test_kebab_case(string: str, expected: str) -> None: + assert kebab_case(string) == expected