109 lines
3.1 KiB
Python
109 lines
3.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import shutil
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from pydantic import BaseModel
|
||
|
|
|
||
|
|
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 BotConfig(BaseModel):
|
||
|
|
"""Bot configuration class."""
|
||
|
|
|
||
|
|
bot_id: str
|
||
|
|
name: str
|
||
|
|
context_files: list[str]
|
||
|
|
|
||
|
|
|
||
|
|
class BotExists(Exception):
|
||
|
|
"""Bot already exists error."""
|
||
|
|
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class BotDoesNotExists(Exception):
|
||
|
|
"""Bot already exists error."""
|
||
|
|
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class Bot:
|
||
|
|
"""Custom bot interface."""
|
||
|
|
|
||
|
|
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))
|
||
|
|
|
||
|
|
@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: list[Path] = []) -> 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.
|
||
|
|
"""
|
||
|
|
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 existing bot."""
|
||
|
|
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.")
|
||
|
|
|
||
|
|
with (bot_path / f"{bot_id}.json").open("r") as f:
|
||
|
|
config = BotConfig(**json.load(f))
|
||
|
|
|
||
|
|
return cls(config, bot_dir)
|