from dataclasses import dataclass
from ai.analysis.model.assistant_chat_model import AssistantChatModel
from ai.analysis.model.token_type import TokenType
from ai.analysis.money.money import Money
[docs]
@dataclass
class ModelPriceInfo:
"""Cost information for input and output tokens of a model.
Attributes:
input_token_cost (Money): Cost per input token.
output_token_cost (Money): Cost per output token.
"""
input_token_cost: Money
output_token_cost: Money
[docs]
def get_cost_for(self, token_type: TokenType) -> Money:
"""Return the cost for the specified token type.
Args:
token_type (TokenType): The type of token (INPUT or OUTPUT).
Returns:
Money: The cost associated with the token type.
Raises:
ValueError: If an unsupported token type is provided.
"""
if token_type == TokenType.INPUT:
return self.input_token_cost
elif token_type == TokenType.OUTPUT:
return self.output_token_cost
raise ValueError("Wrong token type")
[docs]
class TokenPriceCalculator:
"""Calculator for determining token costs of supported assistant models."""
model_price_infos: dict[AssistantChatModel, ModelPriceInfo] = {
AssistantChatModel.GPT_4O: ModelPriceInfo(Money(2.5e-6), Money(10e-6)),
AssistantChatModel.GPT_4O_MINI: ModelPriceInfo(Money(0.15e-6), Money(0.6e-6)),
AssistantChatModel.GPT_4_1: ModelPriceInfo(Money(2e-6), Money(8e-6)),
AssistantChatModel.GPT_4_1_MINI: ModelPriceInfo(Money(0.40e-6), Money(1.60e-6)),
AssistantChatModel.GPT_4_1_NANO: ModelPriceInfo(Money(0.1e-6), Money(0.4e-6)),
AssistantChatModel.LLAMA_70B: ModelPriceInfo(Money(0.12e-6), Money(0.3e-6)),
AssistantChatModel.QWEN_2_5_7B: ModelPriceInfo(Money(0.025e-6), Money(0.05e-6)),
AssistantChatModel.LLAMA_3B: ModelPriceInfo(Money(0.015e-6), Money(0.025e-6)),
AssistantChatModel.QWEN_2_5_7B_TOGETHER: ModelPriceInfo(Money(0.3e-6), Money(0.3e-6)),
}
[docs]
def get_cost(self, model: AssistantChatModel, token_type: TokenType) -> Money:
"""Calculate the cost for a given model and token type.
Args:
model (AssistantChatModel): The chat model to query.
token_type (TokenType): The type of token (INPUT or OUTPUT).
Returns:
Money: The cost for one token of the specified type for the given model.
Raises:
ValueError: If the model is not found in the price information.
"""
if model in self.model_price_infos:
return self.model_price_infos[model].get_cost_for(token_type)
raise ValueError(f"Model not found {model.value}")