llm-chat/src/llm_chat/bot.py

from __future__ import annotations

import json
import shutil
from pathlib import Path
from typing import Iterable

from pydantic import BaseModel

from llm_chat.models import Message, Role
from llm_chat.utils import kebab_case


class BotExists(Exception):
    """Bot already exists error."""

    pass


class BotDoesNotExist(Exception):
    """Bot does not exist error."""

    pass


class BotConfig(BaseModel):
    """Bot configuration class."""

    bot_id: str
    name: str
    context_files: list[str]


def _bot_id_from_name(name: str) -> str:
    """Create bot ID from name.

    Args:
        name: Bot name in full prose (e.g. My Amazing Bot).
    """
    return kebab_case(name)


def _load_context(config: BotConfig, bot_dir: Path) -> list[Message]:
    """Load text from context files.

    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.

    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:
        """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: Iterable[Path] = tuple(),
    ) -> 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/<name>`), which is stored as the snake case version of
        the name argument. the directory contains a settings file `<name>.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.

        Returns:
            Instantiated Bot instance.
        """
        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 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 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)