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

Compare commits

base: paul:main
Branches Tags
paul:main
paul:v2.0.0
paul:v1.1.0
paul:v1.0.0
..
compare: paul:v1.1.0
Branches Tags
paul:main
paul:v2.0.0
paul:v1.1.0
paul:v1.0.0

No commits in common. "main" and "v1.1.0" have entirely different histories.
main ... v1.1.0

16 changed files with 80 additions and 539 deletions
Show Stats Expand all files Collapse all files

33
.gitignore vendored
View File

@ -8,39 +8,6 @@
### asdf ### ### asdf ###
/.tool-versions /.tool-versions
### macOS ###
# General
.DS_Store
.AppleDouble
.LSOverride
# Icon must end with two \r
Icon
# Thumbnails
._*
# Files that might appear in the root of a volume
.DocumentRevisions-V100
.fseventsd
.Spotlight-V100
.TemporaryItems
.Trashes
.VolumeIcon.icns
.com.apple.timemachine.donotpresent
# Directories potentially created on remote AFP share
.AppleDB
.AppleDesktop
Network Trash Folder
Temporary Items
.apdisk
### macOS Patch ###
# iCloud generated files
*.icloud
### Python ### ### Python ###
# Byte-compiled / optimized / DLL files # Byte-compiled / optimized / DLL files
__pycache__/ __pycache__/

2
poetry.lock generated
View File

@ -764,4 +764,4 @@ files = [
[metadata] [metadata]
lock-version = "2.0" lock-version = "2.0"
python-versions = ">=3.11,<3.12" python-versions = ">=3.11,<3.12"
content-hash = "8d76898eeb53fd3848f3be2f6aa1662517f9dbd80146db8dfd6f2932021ace48" content-hash = "e8c5d78c5c95eaadb03e603c5b4ceada8aa27aaa049e6c0d72129c1f2dc53ed9"

2
pyproject.toml
View File

@ -1,6 +1,6 @@
[tool.poetry] [tool.poetry]
name = "llm-chat" name = "llm-chat"
version = "2.0.0" version = "1.1.0"
description = "A general CLI interface for large language models." description = "A general CLI interface for large language models."
authors = ["Paul Harrison <paul@harrison.sh>"] authors = ["Paul Harrison <paul@harrison.sh>"]
readme = "README.md" readme = "README.md"

162
src/llm_chat/bot.py
View File

@ -1,162 +0,0 @@
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)

66
src/llm_chat/chat.py
View File

@ -1,4 +1,3 @@
import math
from datetime import datetime from datetime import datetime
from enum import StrEnum, auto from enum import StrEnum, auto
from pathlib import Path from pathlib import Path
@ -9,7 +8,6 @@ from openai import OpenAI
from openai.types.chat import ChatCompletion from openai.types.chat import ChatCompletion
from openai.types.completion_usage import CompletionUsage from openai.types.completion_usage import CompletionUsage
from llm_chat.bot import Bot
from llm_chat.models import Conversation, Message, Role from llm_chat.models import Conversation, Message, Role
from llm_chat.settings import Model, OpenAISettings from llm_chat.settings import Model, OpenAISettings
@ -54,10 +52,6 @@ class ChatProtocol(Protocol):
conversation: Conversation conversation: Conversation
@property
def bot(self) -> str:
"""Get the name of the bot the conversation is with."""
@property @property
def cost(self) -> float: def cost(self) -> float:
"""Get the cost of the conversation.""" """Get the cost of the conversation."""
@ -80,29 +74,21 @@ class ChatProtocol(Protocol):
class Chat: class Chat:
"""Interface class for OpenAI's ChatGPT chat API. """Interface class for OpenAI's ChatGPT chat API.
Args: Arguments:
settings: Settings for the chat. Defaults to reading from environment settings (optional): Settings for the chat. Defaults to reading from
variables. environment variables.
context: Context for the chat. Defaults to an empty list. context (optional): 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]] = { _pricing: dict[Model, dict[Token, float]] = {
Model.GPT3: { Model.GPT3: {
Token.COMPLETION: 0.0005, Token.COMPLETION: 0.002,
Token.PROMPT: 0.0015, Token.PROMPT: 0.0015,
}, },
Model.GPT4: { Model.GPT4: {
Token.COMPLETION: 0.06, Token.COMPLETION: 0.06,
Token.PROMPT: 0.03, Token.PROMPT: 0.03,
}, },
Model.GPT4_TURBO: {
Token.COMPLETION: 0.03,
Token.PROMPT: 0.01,
},
} }
def __init__( def __init__(
@ -110,23 +96,16 @@ class Chat:
settings: OpenAISettings | None = None, settings: OpenAISettings | None = None,
context: list[Message] = [], context: list[Message] = [],
name: str = "", name: str = "",
bot: str = "",
initial_system_messages: bool = True, initial_system_messages: bool = True,
) -> None: ) -> None:
self._settings = settings 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( self.conversation = Conversation(
messages=context, messages=INITIAL_SYSTEM_MESSAGES + context
if initial_system_messages
else context,
model=self.settings.model, model=self.settings.model,
temperature=self.settings.temperature, temperature=self.settings.temperature,
name=name, name=name,
bot=bot,
) )
self._start_time = datetime.now(tz=ZoneInfo("UTC")) self._start_time = datetime.now(tz=ZoneInfo("UTC"))
self._client = OpenAI( self._client = OpenAI(
@ -135,7 +114,7 @@ class Chat:
@classmethod @classmethod
def load( def load(
cls, path: Path, api_key: str | None = None, base_dir: Path | None = None cls, path: Path, api_key: str | None = None, history_dir: Path | None = None
) -> ChatProtocol: ) -> ChatProtocol:
"""Load a chat from a file.""" """Load a chat from a file."""
with path.open() as f: with path.open() as f:
@ -146,8 +125,8 @@ class Chat:
} }
if api_key is not None: if api_key is not None:
args["api_key"] = api_key args["api_key"] = api_key
if base_dir is not None: if history_dir is not None:
args["base_dir"] = base_dir args["history_dir"] = history_dir
settings = OpenAISettings(**args) settings = OpenAISettings(**args)
return cls( return cls(
@ -163,11 +142,6 @@ class Chat:
self._settings = OpenAISettings() self._settings = OpenAISettings()
return self._settings return self._settings
@property
def bot(self) -> str:
"""Get the name of the bot the conversation is with."""
return self.conversation.bot
@property @property
def cost(self) -> float: def cost(self) -> float:
"""Get the cost of the conversation.""" """Get the cost of the conversation."""
@ -199,17 +173,12 @@ class Chat:
"""Calculate the cost of a request.""" """Calculate the cost of a request."""
self.conversation.completion_tokens += usage.completion_tokens self.conversation.completion_tokens += usage.completion_tokens
self.conversation.prompt_tokens += usage.prompt_tokens self.conversation.prompt_tokens += usage.prompt_tokens
self.conversation.cost = ( self.conversation.cost = round(
math.floor(
1000000
* (
(self.conversation.completion_tokens / 1000) (self.conversation.completion_tokens / 1000)
* self._pricing[self.settings.model][Token.COMPLETION] * self._pricing[self.settings.model][Token.COMPLETION]
+ (self.conversation.prompt_tokens / 1000) + (self.conversation.prompt_tokens / 1000)
* self._pricing[self.settings.model][Token.PROMPT] * self._pricing[self.settings.model][Token.PROMPT],
) 6,
)
/ 1000000
) )
def save(self) -> None: def save(self) -> None:
@ -237,13 +206,10 @@ class Chat:
def get_chat( def get_chat(
settings: OpenAISettings | None = None, settings: OpenAISettings | None = None, context: list[Message] = [], name: str = ""
context: list[Message] = [],
name: str = "",
bot: str = "",
) -> ChatProtocol: ) -> ChatProtocol:
"""Get a chat object.""" """Get a chat object."""
return Chat(settings=settings, context=context, name=name, bot=bot) return Chat(settings=settings, context=context, name=name)
def get_chat_class() -> Type[Chat]: def get_chat_class() -> Type[Chat]:

67
src/llm_chat/cli/main.py → src/llm_chat/cli.py
View File

@ -7,12 +7,16 @@ from rich.console import Console
from rich.markdown import Markdown from rich.markdown import Markdown
from llm_chat.chat import ChatProtocol, get_chat, get_chat_class from llm_chat.chat import ChatProtocol, get_chat, get_chat_class
from llm_chat.cli import bot
from llm_chat.models import Message, Role from llm_chat.models import Message, Role
from llm_chat.settings import Model, OpenAISettings from llm_chat.settings import (
DEFAULT_HISTORY_DIR,
DEFAULT_MODEL,
DEFAULT_TEMPERATURE,
Model,
OpenAISettings,
)
app = typer.Typer() app = typer.Typer()
app.add_typer(bot.app, name="bot", help="Manage custom bots.")
def prompt_continuation(width: int, *args: Any) -> str: def prompt_continuation(width: int, *args: Any) -> str:
@ -75,8 +79,6 @@ def run_conversation(current_chat: ChatProtocol) -> None:
) )
if current_chat.name: if current_chat.name:
console.print(f"[bold green]Name:[/bold green] {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: while not finished:
prompt = read_user_input(session) prompt = read_user_input(session)
@ -104,15 +106,15 @@ def chat(
), ),
] = None, ] = None,
model: Annotated[ model: Annotated[
Optional[Model], Model,
typer.Option(..., "--model", "-m", help="Model to use.", show_choices=True), typer.Option(..., "--model", "-m", help="Model to use.", show_choices=True),
] = None, ] = DEFAULT_MODEL,
temperature: Annotated[ temperature: Annotated[
Optional[float], float,
typer.Option( typer.Option(
..., "--temperature", "-t", help="Model temperature (i.e. creativeness)." ..., "--temperature", "-t", help="Model temperature (i.e. creativeness)."
), ),
] = None, ] = DEFAULT_TEMPERATURE,
context: Annotated[ context: Annotated[
list[Path], list[Path],
typer.Option( typer.Option(
@ -121,7 +123,7 @@ def chat(
"-c", "-c",
help=( help=(
"Path to a file containing context text. " "Path to a file containing context text. "
"Can provide multiple times for multiple files." "Can provide multiple time for multiple files."
), ),
exists=True, exists=True,
file_okay=True, file_okay=True,
@ -129,18 +131,18 @@ def chat(
readable=True, readable=True,
), ),
] = [], ] = [],
base_dir: Annotated[ history_dir: Annotated[
Optional[Path], Path,
typer.Option( typer.Option(
..., ...,
"--base-dir", "--history-dir",
"-d", "-d",
help=( help="Path to the directory where conversation history will be saved.",
"Path to the base directory in which conversation " exists=True,
"configuration and history will be saved." dir_okay=True,
file_okay=False,
), ),
), ] = DEFAULT_HISTORY_DIR,
] = None,
name: Annotated[ name: Annotated[
str, str,
typer.Option( typer.Option(
@ -150,31 +152,24 @@ def chat(
help="Name of the chat.", help="Name of the chat.",
), ),
] = "", ] = "",
bot: Annotated[
str,
typer.Option(
..., "--bot", "-b", help="Name of bot with whom you want to chat."
),
] = "",
) -> None: ) -> None:
"""Start a chat session.""" """Start a chat session."""
# TODO: Add option to provide context string as an argument. # TODO: Add option to provide context string as an argument.
args: dict[str, Any] = {}
if api_key is not None: if api_key is not None:
args |= {"api_key": api_key} settings = OpenAISettings(
if model is not None: api_key=api_key,
args |= {"model": model} model=model,
if temperature is not None: temperature=temperature,
args |= {"temperature": temperature} history_dir=history_dir,
if base_dir is not None: )
args |= {"base_dir": base_dir} else:
settings = OpenAISettings(**args) settings = OpenAISettings(
model=model, temperature=temperature, history_dir=history_dir
)
context_messages = [load_context(path) for path in context] context_messages = [load_context(path) for path in context]
current_chat = get_chat( current_chat = get_chat(settings=settings, context=context_messages, name=name)
settings=settings, context=context_messages, name=name, bot=bot
)
run_conversation(current_chat) run_conversation(current_chat)

3
src/llm_chat/cli/__init__.py
View File

@ -1,3 +0,0 @@
from llm_chat.cli.main import app
__all__ = ["app"]

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

@ -1,79 +0,0 @@
from pathlib import Path
from typing import Annotated, Any, Optional
import typer
from llm_chat.bot import Bot
from llm_chat.settings import OpenAISettings
app = typer.Typer()
@app.command("create")
def create(
name: Annotated[
str,
typer.Argument(help="Name of bot to create."),
],
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,
context_files: Annotated[
list[Path],
typer.Option(
...,
"--context",
"-c",
help=(
"Path to a file containing context text. "
"Can provide multiple times for multiple files."
),
exists=True,
file_okay=True,
dir_okay=False,
readable=True,
),
] = [],
) -> None:
"""Create a new bot."""
args: dict[str, Any] = {}
if base_dir is not None:
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)

1
src/llm_chat/models.py
View File

@ -32,7 +32,6 @@ class Conversation(BaseModel):
model: Model model: Model
temperature: float = DEFAULT_TEMPERATURE temperature: float = DEFAULT_TEMPERATURE
name: str = "" name: str = ""
bot: str = ""
completion_tokens: int = 0 completion_tokens: int = 0
prompt_tokens: int = 0 prompt_tokens: int = 0
cost: float = 0.0 cost: float = 0.0

37
src/llm_chat/settings.py
View File

@ -10,14 +10,11 @@ class Model(StrEnum):
GPT3 = "gpt-3.5-turbo" GPT3 = "gpt-3.5-turbo"
GPT4 = "gpt-4" GPT4 = "gpt-4"
GPT4_TURBO = "gpt-4-turbo-preview"
DEFAULT_MODEL = Model.GPT3 DEFAULT_MODEL = Model.GPT3
DEFAULT_TEMPERATURE = 0.7 DEFAULT_TEMPERATURE = 0.7
DEFAULT_BASE_DIR = Path.home() / ".llm-chat" DEFAULT_HISTORY_DIR = Path.home() / ".llm_chat"
DEFAULT_BOT_PATH = "bots"
DEFAULT_HISTORY_PATH = "history"
class OpenAISettings(BaseSettings): class OpenAISettings(BaseSettings):
@ -26,9 +23,7 @@ class OpenAISettings(BaseSettings):
api_key: str = "" api_key: str = ""
model: Model = DEFAULT_MODEL model: Model = DEFAULT_MODEL
temperature: float = DEFAULT_TEMPERATURE temperature: float = DEFAULT_TEMPERATURE
base_dir: Path = DEFAULT_BASE_DIR history_dir: Path = DEFAULT_HISTORY_DIR
bot_path: str = DEFAULT_BOT_PATH
history_path: str = DEFAULT_HISTORY_PATH
model_config: SettingsConfigDict = SettingsConfigDict( # type: ignore[misc] model_config: SettingsConfigDict = SettingsConfigDict( # type: ignore[misc]
env_file=".env", env_file=".env",
@ -38,25 +33,9 @@ class OpenAISettings(BaseSettings):
use_enum_values=True, use_enum_values=True,
) )
@field_validator("base_dir") @field_validator("history_dir")
def base_dir_must_exist(cls, base_dir: Path) -> Path: def history_dir_must_exist(cls, history_dir: Path) -> Path:
"""Ensure that the base directory exists.""" """Ensure that the history directory exists."""
if not base_dir.exists(): if not history_dir.exists():
base_dir.mkdir(parents=True) history_dir.mkdir(parents=True)
return base_dir return history_dir
@property
def bot_dir(self) -> Path:
"""Return bot directory Path object, creating if required."""
path = self.base_dir / self.bot_path
if not path.exists():
path.mkdir()
return path
@property
def history_dir(self) -> Path:
"""Return history directory Path object, creating if required."""
path = self.base_dir / self.history_path
if not path.exists():
path.mkdir()
return path

9
src/llm_chat/utils.py
View File

@ -1,9 +0,0 @@
import re
def kebab_case(string: str) -> str:
"""Convert a string to kebab case."""
string = string.replace("-", " ")
string = re.sub("([A-Z][a-z]+)", r" \1", string)
string = re.sub("([A-Z]+)", r" \1", string)
return "-".join(string.split()).lower()

4
tests/conftest.py
View File

@ -11,6 +11,6 @@ def mock_openai_api_key() -> None:
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
def mock_base_dir(tmp_path: Path) -> None: def mock_history_dir(tmp_path: Path) -> None:
"""Set a fake history directory.""" """Set a fake history directory."""
os.environ["OPENAI_BASE_DIR"] = str(tmp_path / ".llm_chat") os.environ["OPENAI_HISTORY_DIR"] = str(tmp_path / ".llm_chat")

64
tests/test_bot.py
View File

@ -1,64 +0,0 @@
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,
)

29
tests/test_chat.py
View File

@ -8,7 +8,6 @@ from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage from openai.types.completion_usage import CompletionUsage
from llm_chat.bot import Bot
from llm_chat.chat import Chat, save_conversation from llm_chat.chat import Chat, save_conversation
from llm_chat.models import Conversation, Message, Role from llm_chat.models import Conversation, Message, Role
from llm_chat.settings import Model, OpenAISettings from llm_chat.settings import Model, OpenAISettings
@ -71,21 +70,19 @@ def test_load(tmp_path: Path) -> None:
) )
# Save the conversation to a file # Save the conversation to a file
history_dir = tmp_path / "history" file_path = tmp_path / "conversation.json"
history_dir.mkdir()
file_path = history_dir / "conversation.json"
with file_path.open("w") as f: with file_path.open("w") as f:
f.write(conversation.model_dump_json()) f.write(conversation.model_dump_json())
# Load the conversation from the file # Load the conversation from the file
loaded_chat = Chat.load(file_path, api_key="foo", base_dir=tmp_path) loaded_chat = Chat.load(file_path, api_key="foo", history_dir=tmp_path)
# Check that the loaded conversation matches the original conversation # Check that the loaded conversation matches the original conversation
assert loaded_chat.settings.model == conversation.model assert loaded_chat.settings.model == conversation.model
assert loaded_chat.settings.temperature == conversation.temperature assert loaded_chat.settings.temperature == conversation.temperature
assert loaded_chat.conversation.messages == conversation.messages assert loaded_chat.conversation.messages == conversation.messages
assert loaded_chat.settings.api_key == "foo" assert loaded_chat.settings.api_key == "foo"
assert loaded_chat.settings.base_dir == tmp_path assert loaded_chat.settings.history_dir == tmp_path
# We don't want to load the tokens or cost from the previous session # We don't want to load the tokens or cost from the previous session
assert loaded_chat.conversation.completion_tokens == 0 assert loaded_chat.conversation.completion_tokens == 0
@ -119,9 +116,7 @@ def test_send_message() -> None:
assert response == "Hello!" assert response == "Hello!"
@pytest.mark.parametrize( @pytest.mark.parametrize("model,cost", [(Model.GPT3, 0.000043), (Model.GPT4, 0.00105)])
"model,cost", [(Model.GPT3, round(0.000027, 6)), (Model.GPT4, 0.00105)]
)
def test_calculate_cost(model: Model, cost: float) -> None: def test_calculate_cost(model: Model, cost: float) -> None:
with patch("llm_chat.chat.Chat._make_request") as mock_make_request: with patch("llm_chat.chat.Chat._make_request") as mock_make_request:
mock_make_request.return_value = ChatCompletion( mock_make_request.return_value = ChatCompletion(
@ -146,19 +141,3 @@ def test_calculate_cost(model: Model, cost: float) -> None:
conversation = Chat(settings=settings) conversation = Chat(settings=settings)
_ = conversation.send_message("Hello") _ = conversation.send_message("Hello")
assert conversation.cost == cost 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

33
tests/test_cli.py
View File

@ -9,9 +9,9 @@ from pytest import MonkeyPatch
from rich.console import Console from rich.console import Console
from typer.testing import CliRunner from typer.testing import CliRunner
import llm_chat.cli import llm_chat
from llm_chat.chat import ChatProtocol from llm_chat.chat import ChatProtocol
from llm_chat.cli.main import app from llm_chat.cli import app
from llm_chat.models import Conversation, Message, Role from llm_chat.models import Conversation, Message, Role
from llm_chat.settings import Model, OpenAISettings from llm_chat.settings import Model, OpenAISettings
@ -37,11 +37,6 @@ class ChatFake:
def _set_args(self, **kwargs: Any) -> None: def _set_args(self, **kwargs: Any) -> None:
self.args = kwargs self.args = kwargs
@property
def bot(self) -> str:
"""Get the name of the bot the conversation is with."""
return self.args.get("bot", "")
@property @property
def cost(self) -> float: def cost(self) -> float:
"""Get the cost of the conversation.""" """Get the cost of the conversation."""
@ -82,9 +77,9 @@ def test_chat(monkeypatch: MonkeyPatch) -> None:
mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) mock_read_user_input = MagicMock(side_effect=["Hello", "/q"])
monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat)
monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console)
monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input)
result = runner.invoke(app, ["chat"]) result = runner.invoke(app, ["chat"])
assert result.exit_code == 0 assert result.exit_code == 0
@ -112,9 +107,9 @@ def test_chat_with_context(
mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) mock_read_user_input = MagicMock(side_effect=["Hello", "/q"])
monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat)
monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console)
monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input)
result = runner.invoke(app, ["chat", argument, str(tmp_file)]) result = runner.invoke(app, ["chat", argument, str(tmp_file)])
assert result.exit_code == 0 assert result.exit_code == 0
@ -144,9 +139,9 @@ def test_chat_with_name(
mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) mock_read_user_input = MagicMock(side_effect=["Hello", "/q"])
monkeypatch.setattr(llm_chat.cli.main, "get_chat", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_chat", mock_get_chat)
monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console)
monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input)
result = runner.invoke(app, ["chat", argument, name]) result = runner.invoke(app, ["chat", argument, name])
assert result.exit_code == 0 assert result.exit_code == 0
@ -184,9 +179,9 @@ def test_load(monkeypatch: MonkeyPatch, tmp_path: Path) -> None:
mock_read_user_input = MagicMock(side_effect=["Hello", "/q"]) mock_read_user_input = MagicMock(side_effect=["Hello", "/q"])
monkeypatch.setattr(llm_chat.cli.main, "get_chat_class", mock_get_chat) monkeypatch.setattr(llm_chat.cli, "get_chat_class", mock_get_chat)
monkeypatch.setattr(llm_chat.cli.main, "get_console", mock_get_console) monkeypatch.setattr(llm_chat.cli, "get_console", mock_get_console)
monkeypatch.setattr(llm_chat.cli.main, "read_user_input", mock_read_user_input) monkeypatch.setattr(llm_chat.cli, "read_user_input", mock_read_user_input)
# Load the conversation from the file # Load the conversation from the file
result = runner.invoke(app, ["load", str(file_path)]) result = runner.invoke(app, ["load", str(file_path)])

22
tests/test_utils.py
View File

@ -1,22 +0,0 @@
import pytest
from llm_chat.utils import kebab_case
@pytest.mark.parametrize(
"string,expected",
[
("fooBar", "foo-bar"),
("FooBar", "foo-bar"),
("Foo Bar", "foo-bar"),
("1Foo2Bar3", "1-foo2-bar3"),
],
ids=[
"fooBar",
"FooBar",
"Foo Bar",
"1Foo2Bar3",
],
)
def test_kebab_case(string: str, expected: str) -> None:
assert kebab_case(string) == expected