File size: 1,229 Bytes
d26280a |
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 |
import csv
from io import StringIO
from langchain.schema import HumanMessage, SystemMessage
from gpt_engineer.core.token_usage import TokenUsageLog
def test_format_log():
# arrange
token_usage_log = TokenUsageLog("gpt-4")
request_messages = [
SystemMessage(content="my system message"),
HumanMessage(content="my user prompt"),
]
response = "response from model"
# act
token_usage_log.update_log(request_messages, response, "step 1")
token_usage_log.update_log(request_messages, response, "step 2")
csv_log = token_usage_log.format_log()
# assert
csv_rows = list(csv.reader(StringIO(csv_log)))
assert len(csv_rows) == 3
assert all(len(row) == 7 for row in csv_rows)
def test_usage_cost():
# arrange
token_usage_log = TokenUsageLog("gpt-4")
request_messages = [
SystemMessage(content="my system message"),
HumanMessage(content="my user prompt"),
]
response = "response from model"
# act
token_usage_log.update_log(request_messages, response, "step 1")
token_usage_log.update_log(request_messages, response, "step 2")
usage_cost = token_usage_log.usage_cost()
# assert
assert usage_cost > 0
|