File size: 1,541 Bytes
e7c3249 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
from langchain.chat_models.base import BaseChatModel
from langchain_community.chat_models.fake import FakeListChatModel
from gpt_engineer.core.ai import AI
def mock_create_chat_model(self) -> BaseChatModel:
return FakeListChatModel(responses=["response1", "response2", "response3"])
def test_start(monkeypatch):
monkeypatch.setattr(AI, "_create_chat_model", mock_create_chat_model)
ai = AI("gpt-4")
# act
response_messages = ai.start("system prompt", "user prompt", step_name="step name")
# assert
assert response_messages[-1].content == "response1"
def test_next(monkeypatch):
# arrange
monkeypatch.setattr(AI, "_create_chat_model", mock_create_chat_model)
ai = AI("gpt-4")
response_messages = ai.start("system prompt", "user prompt", step_name="step name")
# act
response_messages = ai.next(
response_messages, "next user prompt", step_name="step name"
)
# assert
assert response_messages[-1].content == "response2"
def test_token_logging(monkeypatch):
# arrange
monkeypatch.setattr(AI, "_create_chat_model", mock_create_chat_model)
ai = AI("gpt-4")
# act
response_messages = ai.start("system prompt", "user prompt", step_name="step name")
usageCostAfterStart = ai.token_usage_log.usage_cost()
ai.next(response_messages, "next user prompt", step_name="step name")
usageCostAfterNext = ai.token_usage_log.usage_cost()
# assert
assert usageCostAfterStart > 0
assert usageCostAfterNext > usageCostAfterStart
|