From c4be19bd513064ab3669bb3772119729bb43e57f Mon Sep 17 00:00:00 2001 From: Paul Harrison Date: Thu, 27 Jun 2024 17:05:24 +0100 Subject: [PATCH] Incorporate custom bots into chat CLI This commit will add the ability to create custom bots in the form of automatically prepending additional bot-specific context to your chat session. A new `chat` CLI argument will be added allowing users to provide the name of an existing bot. For example: ``` llm chat -b "My Bot" ``` An additional `bot` CLI is added for creating and removing bots. Closes #13 --- pyproject.toml | 2 +- src/llm_chat/bot.py | 100 ++++++++++++++++++++++++++++++--------- src/llm_chat/chat.py | 42 ++++++++++++---- src/llm_chat/cli/bot.py | 29 +++++++++++- src/llm_chat/cli/main.py | 12 ++++- src/llm_chat/models.py | 1 + tests/conftest.py | 4 +- tests/test_bot.py | 64 +++++++++++++++++++++++++ tests/test_chat.py | 17 +++++++ tests/test_cli.py | 5 ++ 10 files changed, 239 insertions(+), 37 deletions(-) create mode 100644 tests/test_bot.py diff --git a/pyproject.toml b/pyproject.toml index 84a0213..84db5ce 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "llm-chat" -version = "1.1.5" +version = "2.0.0" description = "A general CLI interface for large language models." authors = ["Paul Harrison "] readme = "README.md" diff --git a/src/llm_chat/bot.py b/src/llm_chat/bot.py index 6d9decd..30fda4c 100644 --- a/src/llm_chat/bot.py +++ b/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 BotDoesNotExist(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,13 +127,36 @@ 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(): - raise BotDoesNotExists(f"Bot {name} does not exist.") + 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) diff --git a/src/llm_chat/chat.py b/src/llm_chat/chat.py index a76d3db..8081ae4 100644 --- a/src/llm_chat/chat.py +++ b/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]: diff --git a/src/llm_chat/cli/bot.py b/src/llm_chat/cli/bot.py index 78cf6ed..65eadb1 100644 --- a/src/llm_chat/cli/bot.py +++ b/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) diff --git a/src/llm_chat/cli/main.py b/src/llm_chat/cli/main.py index 9f7b375..04e0659 100644 --- a/src/llm_chat/cli/main.py +++ b/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) diff --git a/src/llm_chat/models.py b/src/llm_chat/models.py index d3262ca..b820241 100644 --- a/src/llm_chat/models.py +++ b/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 diff --git a/tests/conftest.py b/tests/conftest.py index b315825..861478d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,6 +11,6 @@ def mock_openai_api_key() -> None: @pytest.fixture(autouse=True) -def mock_history_dir(tmp_path: Path) -> None: +def mock_base_dir(tmp_path: Path) -> None: """Set a fake history directory.""" - os.environ["OPENAI_HISTORY_DIR"] = str(tmp_path / ".llm_chat") + os.environ["OPENAI_BASE_DIR"] = str(tmp_path / ".llm_chat") diff --git a/tests/test_bot.py b/tests/test_bot.py new file mode 100644 index 0000000..b30978f --- /dev/null +++ b/tests/test_bot.py @@ -0,0 +1,64 @@ +from pathlib import Path + +import pytest + +from llm_chat.bot import Bot, BotConfig, BotDoesNotExist, BotExists + + +def test_create_load_remove_bot(tmp_path: Path) -> None: + bot_name = "Test Bot" + bot_id = "test-bot" + + with (tmp_path / "context.md").open("w") as f: + f.write("Hello, world!") + + assert not (tmp_path / bot_id).exists() + + Bot.create( + name=bot_name, + bot_dir=tmp_path, + context_files=[tmp_path / "context.md"], + ) + + assert (tmp_path / bot_id).exists() + assert (tmp_path / bot_id / "context").exists() + assert (tmp_path / bot_id / "context" / "context.md").exists() + assert (tmp_path / bot_id / f"{bot_id}.json").exists() + + with (tmp_path / bot_id / f"{bot_id}.json").open() as f: + config = BotConfig.model_validate_json(f.read(), strict=True) + assert config.name == bot_name + assert config.bot_id == bot_id + assert config.context_files == ["context.md"] + + with (tmp_path / bot_id / "context" / "context.md").open() as f: + assert f.read() == "Hello, world!" + + bot = Bot.load(name=bot_name, bot_dir=tmp_path) + assert bot.config == config + assert bot.id == bot_id + assert bot.name == bot_name + + Bot.remove(name=bot_name, bot_dir=tmp_path) + + assert not (tmp_path / bot_id).exists() + + +def test_bot_does_not_exist(tmp_path: Path) -> None: + with pytest.raises(BotDoesNotExist): + Bot.load(name="Test Bot", bot_dir=tmp_path) + + +def test_bot_already_exists(tmp_path: Path) -> None: + bot_name = "Test Bot" + + Bot.create( + name=bot_name, + bot_dir=tmp_path, + ) + + with pytest.raises(BotExists): + Bot.create( + name="Test Bot", + bot_dir=tmp_path, + ) diff --git a/tests/test_chat.py b/tests/test_chat.py index ae8f5e6..c1a0654 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -8,6 +8,7 @@ from openai.types.chat import ChatCompletion, ChatCompletionMessage from openai.types.chat.chat_completion import Choice from openai.types.completion_usage import CompletionUsage +from llm_chat.bot import Bot from llm_chat.chat import Chat, save_conversation from llm_chat.models import Conversation, Message, Role from llm_chat.settings import Model, OpenAISettings @@ -145,3 +146,19 @@ def test_calculate_cost(model: Model, cost: float) -> None: conversation = Chat(settings=settings) _ = conversation.send_message("Hello") assert conversation.cost == cost + + +def test_chat_with_bot(tmp_path: Path) -> None: + settings = OpenAISettings() + bot_name = "Test Bot" + context = "Hello, world!" + + with (tmp_path / "context.md").open("w") as f: + f.write(context) + + Bot.create( + name=bot_name, bot_dir=settings.bot_dir, context_files=[tmp_path / "context.md"] + ) + + chat = Chat(settings=settings, bot=bot_name) + assert chat.conversation.messages[-1].content == context diff --git a/tests/test_cli.py b/tests/test_cli.py index 078f933..0b4b2ec 100644 --- a/tests/test_cli.py +++ b/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."""