paul
/
llm-chat
1
0
Fork 0
Code Issues 5 Pull Requests Packages Projects Releases 3 Wiki Activity

Incorporate bot into chat CLI

This commit is contained in:
Paul Harrison 2024-06-27 16:42:41 +01:00
parent 7968fc0ccf
commit ead1d85ce3
7 changed files with 219 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

96
src/llm_chat/bot.py
View File

@ -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.
class Bot:
"""Custom bot interface."""
def __init__(self, config: BotConfig, bot_dir: Path) -> None:
self.config = config
self.context: list[Message] = []
Returns:
List of system messages to provide as context.
"""
context: list[Message] = []
for context_file in config.context_files:
path = bot_dir / "context" / context_file
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()
self.context.append(Message(role=Role.SYSTEM, content=content))
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:
@ -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)

42
src/llm_chat/chat.py
View File

@ -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]:

29
src/llm_chat/cli/bot.py
View File

@ -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)

12
src/llm_chat/cli/main.py
View File

@ -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)

1
src/llm_chat/models.py
View File

@ -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

64
tests/test_bot.py Normal file
View File

@ -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,
)

5
tests/test_cli.py
View File

@ -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."""