Incorporate bot into chat CLI

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 BotDoesNotExists(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.
+
+    Returns:
+        List of system messages to provide as context.
+    """
+    context: list[Message] = []
+    for context_file in config.context_files:
+        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()
+        context.append(Message(role=Role.SYSTEM, content=content))
+    return context
 
 
 class Bot:
-    """Custom bot interface."""
+    """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: 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))
+        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,7 +127,15 @@ 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():
@@ -106,3 +145,18 @@ class Bot:
             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 BotDoesNotExists(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/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."""