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/`), 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. 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)