72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
from datetime import datetime
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
from zoneinfo import ZoneInfo
|
|
|
|
import pytest
|
|
|
|
from llm_chat.chat import Chat, save_conversation
|
|
from llm_chat.models import Conversation, Message, Role
|
|
from llm_chat.settings import Model, OpenAISettings
|
|
|
|
|
|
def test_save_conversation(tmp_path: Path) -> None:
|
|
conversation = Conversation(
|
|
messages=[
|
|
Message(role=Role.SYSTEM, content="Hello!"),
|
|
Message(role=Role.USER, content="Hi!"),
|
|
Message(role=Role.ASSISTANT, content="How are you?"),
|
|
],
|
|
model=Model.GPT3,
|
|
temperature=0.5,
|
|
completion_tokens=10,
|
|
prompt_tokens=15,
|
|
cost=0.000043,
|
|
)
|
|
|
|
path = tmp_path / ".history"
|
|
expected_file_path = path / "20210101120000.json"
|
|
dt = datetime(2021, 1, 1, 12, 0, 0, tzinfo=ZoneInfo("UTC"))
|
|
|
|
assert not path.exists()
|
|
|
|
save_conversation(
|
|
conversation=conversation,
|
|
history_dir=path,
|
|
dt=dt,
|
|
)
|
|
|
|
assert path.exists()
|
|
assert path.is_dir()
|
|
assert expected_file_path in path.iterdir()
|
|
|
|
with expected_file_path.open() as f:
|
|
conversation_from_file = Conversation.model_validate_json(f.read())
|
|
|
|
assert conversation == conversation_from_file
|
|
|
|
|
|
def test_send_message() -> None:
|
|
with patch("llm_chat.chat.Chat._make_request") as mock_make_request:
|
|
mock_make_request.return_value = {
|
|
"choices": [{"message": {"content": "Hello!"}}],
|
|
"usage": {"completion_tokens": 1, "prompt_tokens": 1},
|
|
}
|
|
conversation = Chat()
|
|
response = conversation.send_message("Hello")
|
|
assert isinstance(response, str)
|
|
assert response == "Hello!"
|
|
|
|
|
|
@pytest.mark.parametrize("model,cost", [(Model.GPT3, 0.000043), (Model.GPT4, 0.00105)])
|
|
def test_calculate_cost(model: Model, cost: float) -> None:
|
|
with patch("llm_chat.chat.Chat._make_request") as mock_make_request:
|
|
mock_make_request.return_value = {
|
|
"choices": [{"message": {"content": "Hello!"}}],
|
|
"usage": {"completion_tokens": 10, "prompt_tokens": 15},
|
|
}
|
|
settings = OpenAISettings(model=model)
|
|
conversation = Chat(settings=settings)
|
|
_ = conversation.send_message("Hello")
|
|
assert conversation.cost == cost
|