163 lines
4.3 KiB
Python
163 lines
4.3 KiB
Python
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)
|