File size: 5,588 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import logging
from dataclasses import dataclass
from typing import List, Union
import tiktoken
from langchain.schema import AIMessage, HumanMessage, SystemMessage
# workaround for function moved in:
# https://github.com/langchain-ai/langchain/blob/535db72607c4ae308566ede4af65295967bb33a8/libs/community/langchain_community/callbacks/openai_info.py
try:
from langchain.callbacks.openai_info import (
get_openai_token_cost_for_model, # fmt: skip
)
except ImportError:
from langchain_community.callbacks.openai_info import (
get_openai_token_cost_for_model, # fmt: skip
)
Message = Union[AIMessage, HumanMessage, SystemMessage]
logger = logging.getLogger(__name__)
@dataclass
class TokenUsage:
"""
Represents token usage statistics for a conversation step.
"""
step_name: str
in_step_prompt_tokens: int
in_step_completion_tokens: int
in_step_total_tokens: int
total_prompt_tokens: int
total_completion_tokens: int
total_tokens: int
class Tokenizer:
"""
Tokenizer for counting tokens in text.
"""
def __init__(self, model_name):
self.model_name = model_name
self._tiktoken_tokenizer = (
tiktoken.encoding_for_model(model_name)
if "gpt-4" in model_name or "gpt-3.5" in model_name
else tiktoken.get_encoding("cl100k_base")
)
def num_tokens(self, txt: str) -> int:
"""
Get the number of tokens in a text.
Parameters
----------
txt : str
The text to count the tokens in.
Returns
-------
int
The number of tokens in the text.
"""
return len(self._tiktoken_tokenizer.encode(txt))
def num_tokens_from_messages(self, messages: List[Message]) -> int:
"""
Get the total number of tokens used by a list of messages.
Parameters
----------
messages : List[Message]
The list of messages to count the tokens in.
Returns
-------
int
The total number of tokens used by the messages.
"""
n_tokens = 0
for message in messages:
n_tokens += (
4 # Every message follows <im_start>{role/name}\n{content}<im_end>\n
)
n_tokens += self.num_tokens(message.content)
n_tokens += 2 # Every reply is primed with <im_start>assistant
return n_tokens
class TokenUsageLog:
"""
Represents a log of token usage statistics for a conversation.
"""
def __init__(self, model_name):
self.model_name = model_name
self._cumulative_prompt_tokens = 0
self._cumulative_completion_tokens = 0
self._cumulative_total_tokens = 0
self._log = []
self._tokenizer = Tokenizer(model_name)
def update_log(self, messages: List[Message], answer: str, step_name: str) -> None:
"""
Update the token usage log with the number of tokens used in the current step.
Parameters
----------
messages : List[Message]
The list of messages in the conversation.
answer : str
The answer from the AI.
step_name : str
The name of the step.
"""
prompt_tokens = self._tokenizer.num_tokens_from_messages(messages)
completion_tokens = self._tokenizer.num_tokens(answer)
total_tokens = prompt_tokens + completion_tokens
self._cumulative_prompt_tokens += prompt_tokens
self._cumulative_completion_tokens += completion_tokens
self._cumulative_total_tokens += total_tokens
self._log.append(
TokenUsage(
step_name=step_name,
in_step_prompt_tokens=prompt_tokens,
in_step_completion_tokens=completion_tokens,
in_step_total_tokens=total_tokens,
total_prompt_tokens=self._cumulative_prompt_tokens,
total_completion_tokens=self._cumulative_completion_tokens,
total_tokens=self._cumulative_total_tokens,
)
)
def log(self) -> List[TokenUsage]:
"""
Get the token usage log.
Returns
-------
List[TokenUsage]
A log of token usage details per step in the conversation.
"""
return self._log
def format_log(self) -> str:
"""
Format the token usage log as a CSV string.
Returns
-------
str
The token usage log formatted as a CSV string.
"""
result = "step_name,prompt_tokens_in_step,completion_tokens_in_step,total_tokens_in_step,total_prompt_tokens,total_completion_tokens,total_tokens\n"
for log in self._log:
result += f"{log.step_name},{log.in_step_prompt_tokens},{log.in_step_completion_tokens},{log.in_step_total_tokens},{log.total_prompt_tokens},{log.total_completion_tokens},{log.total_tokens}\n"
return result
def usage_cost(self) -> float:
"""
Return the total cost in USD of the API usage.
Returns
-------
float
Cost in USD.
"""
result = 0
for log in self.log():
result += get_openai_token_cost_for_model(
self.model_name, log.total_prompt_tokens, is_completion=False
)
result += get_openai_token_cost_for_model(
self.model_name, log.total_completion_tokens, is_completion=True
)
return result
|