Incorporate custom bots into chat CLI

This commit will add the ability to create custom bots in the form of
automatically prepending additional bot-specific context to your chat
session. A new `chat` CLI argument will be added allowing users to
provide the name of an existing bot. For example:

```
llm chat -b "My Bot"
```

An additional `bot` CLI is added for creating and removing bots.

Closes #13

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "llm-chat"
-version = "1.1.5"
+version = "2.0.0"
 description = "A general CLI interface for large language models."
 authors = ["Paul Harrison <paul@harrison.sh>"]
 readme = "README.md"

src/llm_chat/bot.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import json
 import shutil
 from pathlib import Path
+from typing import Iterable
 
 from pydantic import BaseModel
 
@@ -10,9 +11,16 @@ 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 BotExists(Exception):
+    """Bot already exists error."""
+
+    pass
+
+
+class BotDoesNotExist(Exception):
+    """Bot does not exist error."""
+
+    pass
 
 
 class BotConfig(BaseModel):
@@ -23,33 +31,48 @@ class BotConfig(BaseModel):
     context_files: list[str]
 
 
-class BotExists(Exception):
-    """Bot already exists error."""
+def _bot_id_from_name(name: str) -> str:
+    """Create bot ID from name.
 
-    pass
+    Args:
+        name: Bot name in full prose (e.g. My Amazing Bot).
+    """
+    return kebab_case(name)
 
 
-class BotDoesNotExists(Exception):
-    """Bot already exists error."""
+def _load_context(config: BotConfig, bot_dir: Path) -> list[Message]:
+    """Load text from context files.
 
-    pass
+    Args:
+        config: Bot configuration.
 
-
-class Bot:
-    """Custom bot interface."""
-
-    def __init__(self, config: BotConfig, bot_dir: Path) -> None:
-        self.config = config
-        self.context: list[Message] = []
+    Returns:
+        List of system messages to provide as context.
+    """
+    context: list[Message] = []
     for context_file in config.context_files:
-            path = bot_dir / "context" / context_file
+        path = bot_dir / config.bot_id / "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))
+        context.append(Message(role=Role.SYSTEM, content=content))
+    return context
+
+
+class Bot:
+    """Custom bot interface.
+
+    Args:
+        config: Bot configuration instance.
+        bot_dir: Path to directory of bot configurations.
+    """
+
+    def __init__(self, config: BotConfig, bot_dir: Path) -> None:
+        self.config = config
+        self.context = _load_context(config, bot_dir)
 
     @property
     def id(self) -> str:
@@ -62,7 +85,12 @@ class Bot:
         return self.config.name
 
     @classmethod
-    def create(cls, name: str, bot_dir: Path, context_files: list[Path] = []) -> None:
+    def create(
+        cls,
+        name: str,
+        bot_dir: Path,
+        context_files: Iterable[Path] = tuple(),
+    ) -> None:
         """Create a custom bot.
 
         This command creates the directory structure for the custom bot and copies
@@ -75,6 +103,9 @@ class Bot:
             name: Name of the custom bot.
             bot_dir: Path to where custom bot contexts are stored.
             context_files: Paths to context files.
+
+        Returns:
+            Instantiated Bot instance.
         """
         bot_id = _bot_id_from_name(name)
         path = bot_dir / bot_id
@@ -96,13 +127,36 @@ class Bot:
 
     @classmethod
     def load(cls, name: str, bot_dir: Path) -> Bot:
-        """Load existing bot."""
+        """Load an existing bot.
+
+        Args:
+            name: Name of the custom bot.
+            bot_dir: Path to where custom bot contexts are stored.
+
+        Returns:
+            Instantiated Bot instance.
+        """
         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.")
+            raise BotDoesNotExist(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)
+
+    @classmethod
+    def remove(cls, name: str, bot_dir: Path) -> None:
+        """Remove an existing bot.
+
+        Args:
+            name: Name of the custom bot.
+            bot_dir: Path to where custom bot contexts are stored.
+        """
+        bot_id = _bot_id_from_name(name)
+        bot_path = bot_dir / bot_id
+        if not bot_path.exists():
+            raise BotDoesNotExist(f"Bot {name} does not exist.")
+
+        shutil.rmtree(bot_path)

src/llm_chat/chat.py

@@ -9,6 +9,7 @@ from openai import OpenAI
 from openai.types.chat import ChatCompletion
 from openai.types.completion_usage import CompletionUsage
 
+from llm_chat.bot import Bot
 from llm_chat.models import Conversation, Message, Role
 from llm_chat.settings import Model, OpenAISettings
 
@@ -53,6 +54,10 @@ class ChatProtocol(Protocol):
 
     conversation: Conversation
 
+    @property
+    def bot(self) -> str:
+        """Get the name of the bot the conversation is with."""
+
     @property
     def cost(self) -> float:
         """Get the cost of the conversation."""
@@ -75,10 +80,14 @@ class ChatProtocol(Protocol):
 class Chat:
     """Interface class for OpenAI's ChatGPT chat API.
 
-    Arguments:
-        settings (optional): Settings for the chat. Defaults to reading from
-            environment variables.
-        context (optional): Context for the chat. Defaults to an empty list.
+    Args:
+        settings: Settings for the chat. Defaults to reading from environment
+            variables.
+        context: Context for the chat. Defaults to an empty list.
+        name: Name of the chat.
+        bot: Name of bot to chat with.
+        initial_system_messages: Whether to include the standard initial system
+            messages.
     """
 
     _pricing: dict[Model, dict[Token, float]] = {
@@ -101,16 +110,23 @@ class Chat:
         settings: OpenAISettings | None = None,
         context: list[Message] = [],
         name: str = "",
+        bot: str = "",
         initial_system_messages: bool = True,
     ) -> None:
         self._settings = settings
+
+        if bot:
+            context = Bot.load(bot, self.settings.bot_dir).context + context
+
+        if initial_system_messages:
+            context = INITIAL_SYSTEM_MESSAGES + context
+
         self.conversation = Conversation(
-            messages=INITIAL_SYSTEM_MESSAGES + context
-            if initial_system_messages
-            else context,
+            messages=context,
             model=self.settings.model,
             temperature=self.settings.temperature,
             name=name,
+            bot=bot,
         )
         self._start_time = datetime.now(tz=ZoneInfo("UTC"))
         self._client = OpenAI(
@@ -147,6 +163,11 @@ class Chat:
             self._settings = OpenAISettings()
         return self._settings
 
+    @property
+    def bot(self) -> str:
+        """Get the name of the bot the conversation is with."""
+        return self.conversation.bot
+
     @property
     def cost(self) -> float:
         """Get the cost of the conversation."""
@@ -216,10 +237,13 @@ class Chat:
 
 
 def get_chat(
-    settings: OpenAISettings | None = None, context: list[Message] = [], name: str = ""
+    settings: OpenAISettings | None = None,
+    context: list[Message] = [],
+    name: str = "",
+    bot: str = "",
 ) -> ChatProtocol:
     """Get a chat object."""
-    return Chat(settings=settings, context=context, name=name)
+    return Chat(settings=settings, context=context, name=name, bot=bot)
 
 
 def get_chat_class() -> Type[Chat]:

src/llm_chat/cli/bot.py

@@ -13,7 +13,7 @@ app = typer.Typer()
 def create(
     name: Annotated[
         str,
-        typer.Argument(help="Name of bot."),
+        typer.Argument(help="Name of bot to create."),
     ],
     base_dir: Annotated[
         Optional[Path],
@@ -50,3 +50,30 @@ def create(
         args |= {"base_dir": base_dir}
     settings = OpenAISettings(**args)
     Bot.create(name, settings.bot_dir, context_files=context_files)
+
+
+@app.command("remove")
+def remove(
+    name: Annotated[
+        str,
+        typer.Argument(help="Name of bot to remove."),
+    ],
+    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,
+) -> None:
+    """Remove an existing bot."""
+    args: dict[str, Any] = {}
+    if base_dir is not None:
+        args |= {"base_dir": base_dir}
+    settings = OpenAISettings(**args)
+    Bot.remove(name, settings.bot_dir)

src/llm_chat/cli/main.py

@@ -75,6 +75,8 @@ def run_conversation(current_chat: ChatProtocol) -> None:
     )
     if current_chat.name:
         console.print(f"[bold green]Name:[/bold green] {current_chat.name}")
+    if current_chat.bot:
+        console.print(f"[bold green]Bot:[/bold green] {current_chat.bot}")
 
     while not finished:
         prompt = read_user_input(session)
@@ -148,6 +150,12 @@ def chat(
             help="Name of the chat.",
         ),
     ] = "",
+    bot: Annotated[
+        str,
+        typer.Option(
+            ..., "--bot", "-b", help="Name of bot with whom you want to chat."
+        ),
+    ] = "",
 ) -> None:
     """Start a chat session."""
     # TODO: Add option to provide context string as an argument.
@@ -164,7 +172,9 @@ def chat(
 
     context_messages = [load_context(path) for path in context]
 
-    current_chat = get_chat(settings=settings, context=context_messages, name=name)
+    current_chat = get_chat(
+        settings=settings, context=context_messages, name=name, bot=bot
+    )
 
     run_conversation(current_chat)
 

src/llm_chat/models.py

@@ -32,6 +32,7 @@ class Conversation(BaseModel):
     model: Model
     temperature: float = DEFAULT_TEMPERATURE
     name: str = ""
+    bot: str = ""
     completion_tokens: int = 0
     prompt_tokens: int = 0
     cost: float = 0.0

tests/conftest.py

@@ -11,6 +11,6 @@ def mock_openai_api_key() -> None:
 
 
 @pytest.fixture(autouse=True)
-def mock_history_dir(tmp_path: Path) -> None:
+def mock_base_dir(tmp_path: Path) -> None:
     """Set a fake history directory."""
-    os.environ["OPENAI_HISTORY_DIR"] = str(tmp_path / ".llm_chat")
+    os.environ["OPENAI_BASE_DIR"] = str(tmp_path / ".llm_chat")

tests/test_bot.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+
+import pytest
+
+from llm_chat.bot import Bot, BotConfig, BotDoesNotExist, BotExists
+
+
+def test_create_load_remove_bot(tmp_path: Path) -> None:
+    bot_name = "Test Bot"
+    bot_id = "test-bot"
+
+    with (tmp_path / "context.md").open("w") as f:
+        f.write("Hello, world!")
+
+    assert not (tmp_path / bot_id).exists()
+
+    Bot.create(
+        name=bot_name,
+        bot_dir=tmp_path,
+        context_files=[tmp_path / "context.md"],
+    )
+
+    assert (tmp_path / bot_id).exists()
+    assert (tmp_path / bot_id / "context").exists()
+    assert (tmp_path / bot_id / "context" / "context.md").exists()
+    assert (tmp_path / bot_id / f"{bot_id}.json").exists()
+
+    with (tmp_path / bot_id / f"{bot_id}.json").open() as f:
+        config = BotConfig.model_validate_json(f.read(), strict=True)
+        assert config.name == bot_name
+        assert config.bot_id == bot_id
+        assert config.context_files == ["context.md"]
+
+    with (tmp_path / bot_id / "context" / "context.md").open() as f:
+        assert f.read() == "Hello, world!"
+
+    bot = Bot.load(name=bot_name, bot_dir=tmp_path)
+    assert bot.config == config
+    assert bot.id == bot_id
+    assert bot.name == bot_name
+
+    Bot.remove(name=bot_name, bot_dir=tmp_path)
+
+    assert not (tmp_path / bot_id).exists()
+
+
+def test_bot_does_not_exist(tmp_path: Path) -> None:
+    with pytest.raises(BotDoesNotExist):
+        Bot.load(name="Test Bot", bot_dir=tmp_path)
+
+
+def test_bot_already_exists(tmp_path: Path) -> None:
+    bot_name = "Test Bot"
+
+    Bot.create(
+        name=bot_name,
+        bot_dir=tmp_path,
+    )
+
+    with pytest.raises(BotExists):
+        Bot.create(
+            name="Test Bot",
+            bot_dir=tmp_path,
+        )

tests/test_chat.py

@@ -8,6 +8,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionMessage
 from openai.types.chat.chat_completion import Choice
 from openai.types.completion_usage import CompletionUsage
 
+from llm_chat.bot import Bot
 from llm_chat.chat import Chat, save_conversation
 from llm_chat.models import Conversation, Message, Role
 from llm_chat.settings import Model, OpenAISettings
@@ -145,3 +146,19 @@ def test_calculate_cost(model: Model, cost: float) -> None:
         conversation = Chat(settings=settings)
         _ = conversation.send_message("Hello")
         assert conversation.cost == cost
+
+
+def test_chat_with_bot(tmp_path: Path) -> None:
+    settings = OpenAISettings()
+    bot_name = "Test Bot"
+    context = "Hello, world!"
+
+    with (tmp_path / "context.md").open("w") as f:
+        f.write(context)
+
+    Bot.create(
+        name=bot_name, bot_dir=settings.bot_dir, context_files=[tmp_path / "context.md"]
+    )
+
+    chat = Chat(settings=settings, bot=bot_name)
+    assert chat.conversation.messages[-1].content == context

tests/test_cli.py

@@ -37,6 +37,11 @@ class ChatFake:
     def _set_args(self, **kwargs: Any) -> None:
         self.args = kwargs
 
+    @property
+    def bot(self) -> str:
+        """Get the name of the bot the conversation is with."""
+        return self.args.get("bot", "")
+
     @property
     def cost(self) -> float:
         """Get the cost of the conversation."""